Implementing rejection sampling with vmap

[Binbose]

I am trying to implement rejection sampling in jax for an arbitrary function ('apply_fun') that is parameterized by 'params', and then vmap over different sets of parameters.
For this I am allocating an array of length 'num_samples' and then iterate in a while loop until I have num_samples, see code below (I tried a python while loop and a lax.while_loop):

def sample_fun(rng, params, num_samples, ymax):
         def rejection_sample(args):
            rng, all_x, params, i, ymax = args
            rng, split_rng = jax.random.split(rng)
            x = jax.random.uniform(split_rng, minval=0, maxval=1, shape=(1,))
            rng, split_rng = jax.random.split(rng)
            y = jax.random.uniform(split_rng, minval=0, maxval=ymax, shape=(1,))

            passed = (y < apply_fun(params, x)).astype(bool)
            all_x = all_x.at[i].add((passed * x)[0])
            i = i + passed
            # all_x = jax.lax.cond(passed, lambda x: all_x.at(0).set(x), lambda x: all_x, x)
            # i = jax.lax.cond(passed, lambda i: i+1, lambda i: i, i)

            return rng, all_x, params, i, ymax


         i = 0
         all_x = np.zeros(num_samples)
         jax.lax.while_loop(lambda i: i[3] < num_samples, rejection_sample, (rng, all_x, params, i, ymax))

         # while i < num_samples:
         #    rng, split_rng = jax.random.split(rng)
         #    x = jax.random.uniform(split_rng, minval=0, maxval=1, shape=(1,))
         #    rng, split_rng = jax.random.split(rng)
         #    y = jax.random.uniform(split_rng, minval=0, maxval=ymax, shape=(1,))
         #
         #    passed = (y < apply_fun(params, x)).astype(bool)
         #    all_x = all_x.at[i].add((passed * x)[0])
         #    i = i + passed

         return all_x

sample_fun_vec = jit(partial(vmap(sample_fun, in_axes=(None, 0))))

However, this fails either with a concretization error in the python while loop, or with this error:

TypeError: body_fun output and input must have identical types, got
('ShapedArray(uint32[2])', 'ShapedArray(float32[10])', 'ShapedArray(float32[2])', 'DIFFERENT ShapedArray(int32[1]) vs. ShapedArray(int32[])', 'ShapedArray(float32[])').

in the lax.while_loop.
Any ideas how to implement this correctly?

[YouJiacheng]

Hmmm, I think your code is fine...

def sample_fun(rng, params, num_samples, ymax):
         def rejection_sample(args):
            rng, all_x, i= args
            rng, split_rng = jax.random.split(rng)
            x = jax.random.uniform(split_rng, minval=0, maxval=1, shape=(1,))
            rng, split_rng = jax.random.split(rng)
            y = jax.random.uniform(split_rng, minval=0, maxval=ymax, shape=(1,))

            passed = (y < apply_fun(params, x)).astype(bool)
            all_x = all_x.at[i].add((passed * x)[0])
            i = i + passed[0] # fix here
            return rng, all_x, i
         all_x = np.zeros(num_samples)
         _, all_x, _ = jax.lax.while_loop(lambda i: i[2] < num_samples, rejection_sample, (rng, all_x, 0))
         return all_x

[Binbose]

Unfortunately, it doesnt work, I added a small dummy function for a minimal self contained example to reproduce the error (based on your code for readability)

def apply_fun(params, x):
   return params[0] * x**2 + params[1] * x


def sample_fun(rng, params, num_samples, ymax):
   def rejection_sample(args):
      rng, all_x, i = args
      rng, split_rng = jax.random.split(rng)
      x = jax.random.uniform(split_rng, minval=0, maxval=1, shape=(1,))
      rng, split_rng = jax.random.split(rng)
      y = jax.random.uniform(split_rng, minval=0, maxval=ymax, shape=(1,))

      passed = (y < apply_fun(params, x)).astype(bool)
      all_x = all_x.at[i].add((passed * x)[0])
      i = i + passed
      return rng, all_x, i

   all_x = np.zeros(num_samples)
   _, all_x, _ = jax.lax.while_loop(lambda i: i[2] < num_samples, rejection_sample, (rng, all_x, 0))
   return all_x

sample_fun_vec = partial(vmap(sample_fun, in_axes=(None, 0, None, None)))
apply_fun_vec = vmap(apply_fun, in_axes=(None, 0))

rng = jax.random.PRNGKey(0)
params = jax.random.uniform(rng, minval=0, maxval=1, shape=(3, 2))
xx = np.linspace(0, 1, 100)

fig, ax = plt.subplots()
ax.plot(xx, apply_fun_vec(params, xx))
s = sample_fun_vec(rng, params, 10, 2)
ax.hist(np.array(s), density=True, bins=100)

ax.grid(True)
ax.legend(loc='best')
plt.show()

Which produces the aforementioned error:

TypeError: body_fun output and input must have identical types, got
('ShapedArray(uint32[2])', 'ShapedArray(float32[10])', 'DIFFERENT ShapedArray(int32[1]) vs. ShapedArray(int32[])').

[YouJiacheng]

i = i + passed => i = i + passed[0] will fix it! I didn't notice that!

[Binbose]

Hey, you are right, now it doesn't throw the error anymore!
However, another problem is now, that vmap passes the same rng key to all instances of the function, which means it always creates the same sample

[YouJiacheng]

You can just use jax.random.split to create multiple keys and also vmap on rng.

[Binbose]

Awesome, that works, thank you!