Vectorized sampling along one axis with jax.random.choice

[Qottmann]

I have an ndarray of some shape (N, d0, d1) and want to sample along the first axis. I.e. I would want something along the way of

key = jax.random.PRNGKey(42)
a = jnp.arange(15*5*3).reshape(15, 5, 3)
jax.random.choice(key, a, shape=(5, 5, 3)).shape == (5, 5, 3)

Is there an efficient way to do this without iterating through d0 and d1?

Bonus:
Do the same with a (N, d0, d1) shaped p

[davisyoshida]

You just want shape = (5,) since what choice does is give you a random subset of the tensor along the requested axis.

I'm not sure what it would mean to use an (N, d0, d1) shaped p though, since all the sampling is happening along the first axis. Wouldn't you just want an (N,) shaped p?

[Qottmann]

Thanks! The first part works :)

I'm not sure what it would mean to use an (N, d0, d1) shaped p though, since all the sampling is happening along the first axis. Wouldn't you just want an (N,) shaped p?

It would just be a different distribution for each value of the 2nd and 3rd axis.
You can think of it as a 1d p of size (N,) for each value of d0, d1.
I tried vmapping but it is difficult with keywars. I also reduced the complexity of the problem to only one secondary axis. Here is my attempt at some hacky solution, which runs but yields the wrong result

key = jax.random.PRNGKey(42)
a = jnp.arange(15*5).reshape(15, 5)
p = jnp.arange(15*5).reshape(15, 5)

def new_choice(a, p):
    return jax.random.choice(key, a, shape=(5,), p=p)
jax.vmap(new_choice, [None, 1])(a, p) # shape is (5, 5, 5)

[davisyoshida]

Hmm maybe I'm not quite sure what you're going for. My understanding is that jax.random.choice on an MxN array would be taking random rows of the matrix, so having different ps for each column wouldn't really fit.

[jakevdp]

For your first question, I believe @davisyoshida is right that you just need to pass shape=(5,). An alternative way to do this kind of sampling is by using vmap; it might look like this:

split_key = jax.random.split(key, a[0].size).reshape(*a[0].shape, *key.shape)
choice = partial(jax.random.choice, shape=(5,))
choice = jax.vmap(choice, in_axes=(0, 1), out_axes=1)
choice = jax.vmap(choice, in_axes=(1, 2), out_axes=2)
out = choice(split_key, a)
out.shape  # (5, 5, 3)

The vmap approach is somewhat more complicated, but it allows you to express exactly the mapping semantics that you wish, rather than relying on the implicit broadcasting/vectorization semantics of the function itself. This becomes relevant for the second part of your question, which can be directly addressed using the vmap approach:

import numpy as np
from functools import partial

p = np.random.rand(*a.shape)
p /= p.sum(0, keepdims=True)

@partial(jax.vmap, in_axes=(1, 2, 2), out_axes=2)
@partial(jax.vmap, in_axes=(0, 1, 1), out_axes=1)
def choice(key, a, p):
  return jax.random.choice(key, a, shape=(5,), p=p)

out = choice(split_key, a, p)
out.shape  # (5, 5, 3)

[Qottmann]

Amazing, thank you so much @jakevdp 🚀