Batching rule for 'gather' not implemented

Here is a minimal example to reproduce the error.

import jax
import jax.numpy as np

a = np.ones((3, 3))
def f(a):
    inds = np.arange(a.shape[0])
    return a[inds]
jax.vmap(f)(a)

NotImplementedError: Batching rule for 'index_take' not implemented

So if you want to use vmap, the only form of indexing you can do is scalar indexing. Am I interpreting this correctly? For me that’s a huge limitation.

I would try to implement the batch rule but I think it will be complicated and I have a workaround: np.stack([a[ind] for ind in inds]) instead of a[inds]. With an ideal JIT compiler those two would be equivalent, but my guess is that the performance of this workaround is not optimal.

[mattjj]

Since #307 removed index_take as a primitive, running the OP’s code now results in

NotImplementedError: Batching rule for 'gather' not implemented

I changed the title of the issue to reflect the new topic: we want a batching rule for lax.gather!

[mattjj]

Whew! We finally got this one done in #351. The challenge was that XLA’s Gather HLO is a bit complex. In fact, it supports more fancy indexing than you can even express in NumPy. But we finally completed the batching rule for it, and as a result we can vmap through all of NumPy’s advanced integer array indexing now.

(The story isn’t totally over until we get #353 in too, but that should get in momentarily.)