Batched jax.lax.gather

[andreasveit]

I have a tensor of batched embeddings of shape
embs.shape = [batch_size, num_embeddings_per_example, embedding_dim]
I further have a tensor of batched indices into this embedding tensor
embs_indices.shape = [batch_size, num_indices_per_example]
num_indices_per_example <= num_embeddings_per_example

I would like to gather the embeddings from embs into a new tensor of shape
output.shape = [batch_size, num_indices_per_example, embedding_dim]

Similar to the pytorch operation of embs.gather(dim=-2, index=embs_indices)

I have looked into jax.lax.gather, but have not found a way to specify batch dimensions.

For a single example, I was able to get this to work:

embs = jnp.reshape(jnp.arange(100), [10, 10])
embs_indices = jnp.array([[2], [2]])
jax.lax.gather(
    input, indices,
    jax.lax.GatherDimensionNumbers(
        offset_dims=(0, 1),
        collapsed_slice_dims=(),
        start_index_map=(0,)),
    slice_sizes=(1, embs.shape[1])).transpose([-1, -3, -2])

Providing the following result.

Array([[[20, 21, 22, 23, 24, 25, 26, 27, 28, 29]],
       [[20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]], dtype=int32)

Any recommendation for how to add a leading batch dimension to this?

[jakevdp]

In general, I'd avoid calling lax.gather directly. Instead, I would use numpy-style indexing, which is far more flexible and easy to use.

In your case, it would look like this:

result = embs[embs_indices]

Regarding how to "add a leading batch dimension to this" – I'm not entirely clear on what you want to do. But if you're asking about adding a batch dimension to both embs and embs_indices, you could do it with vmap:

embs_batched = jnp.stack([embs + i for i in range(4)])
embs_indices_batched = jnp.stack([embs_indices + i for i in range(4)])
batched_result = jax.vmap(lambda a, b: a[b])(embs_batched, embs_indices_batched)

The vmap transformation here will essentially lower this to a single batched gather operation, without you having to worry about getting all the parameters correct.