Unused import?

This split should happen before rng is used to generate random numbers, e.g.

rng, next_rng = jax.random.split(rng)

otherwise you may have subtle correlational bugs in the pseudo-random numbers.

As I have changed to use vmap, so now I do the splitting in the vmap argument: jax.random.split(key, n_resamples)

vmap is fine here, but do you need multiple vmaps if theta_hat_b has more than two dimensions?

should I use a control flow here to apply multiple vmap if I find that theta_hat_b have more than 2 dimensions?

You can use Python for loops to apply multiple vmap transforms; something like this:

import jax
import jax.numpy as jnp

def f(x):
  return x.sum()

x = jnp.ones((2, 3, 4, 5))
for i in range(x.ndim - 1):
  f = jax.vmap(f)
print(f(x).shape)
# (2, 3, 4)

Is there a reason you chose scan instead of vmap? Because it's a serial operation, scan is going to be far slower than vmap on accelerators.

nice suggestion!
for _bootstrap_resample_and_compute_statistic, vmap does run faster than scan

Same comment here regarding scan - we should express this in terms of vmap instead to take advantage of parallelism on accelerators.

I found an issue here @jakevdp

for _jackknife_resample_and_compute_statistic, I weirdly find that vmap runs SLOWER than scan.
On CPU Colab with 10,000 samples, vmap speed is ~210 ms while scan speed is ~110 ms.

Here is the Colab notebook that I did my benchmarking:
https://colab.research.google.com/drive/1abKv-zyI-CZ4BEtXc3oh5Ebx_rcFq0dp?usp=sharing

I am not sure if its because my implementation is inefficient, or I have made some errrors.

Interesting! On CPU scan and friends are generally not too bad, but I suspect if you benchmarked on GPU/TPU you'd find scan to be much slower than vmap.

This would be more consistent with JAX's typical approach if (1) key were the first argument, and (2) batching were handled via vmap at the call-site, rather than via an argument passed to the function. So the API would be something like this:

def _bootstrap_resample_and_compute_statistic(key, sample, statistic):
  ...

and rather than calling it like

_bootstrap_resample_and_compute_statistic(sample, statistic, n_resamples, key)

you could instead call it like

keys = random.split(key, n_resamples)
vmap(_bootstrap_resample_and_compute_statistic, (0, None, None))(keys, sample, statistic)

The benefits would be (1) more explicit handling of key splitting by the user of the function, and (2) vmap at the outer levels may be somewhat more efficient (I'm not entirely sure on that, but I think it is the case) and (3) it's more maintainable, because it makes use of JAX's composable transforms in the way they're intended to be used, rather than hiding them behind less flexible batch-aware APIs.

Same comment here. Can we define _jackknife_resample_and_compute_statistic so it natively handles a single batch, and then use vmap as appropriate at the call-site?

Rather than vmapping, can we use the axis argument to jnp.percentile?

I don't understand the purpose of empty indexing here.

Typically APIs don't require object identity for boolean values.

You can use extra_params here to document the key argument.

That said, I'm starting to wonder if this should really be considered a wrapper of scipy.stats.bootstrap, because its API is now substantially different. In numpy's case, we don't provide any wrappers for numpy.random functionality, instead using a different key-based API in jax.random. I'm starting to think that the same treatment would make sense here, because as written jax.scipy.bootstrap must be called with a different signature than scipy.bootstrap.

It also would solve the issue of how to handle irrelevant params like vectorized, and we could write the API in a way that is more typical of JAX library functions (i.e. keep batching orthogonal to the implementation, rather than calling vmap within.

What do you think?

@jakevdp Is there a decision/consensus on whether to adhere to the original API?