Jax equivalent of tf.gather_nd

[AlexiaJM]

How to implement tf.gather_nd in Jax?

This is similar to #3658 and the answer by @shoyer seemed simple enough.

[hawkinsp]

Without batch dimensions, x[tuple(jnp.moveaxis(indices, -1, 0))] should work.

If you need batch dimensions, you can use vmap. (It's slightly more awkward with many batch dimensions.)

[AlexiaJM]

Thank you, I will try this!

[AlexiaJM]

For those curious, here is my batched version:

def gather_nd_unbatched(params, indices):
  return params[tuple(jnp.moveaxis(indices, -1, 0))]

def gather_nd(params, indices, batch=False):
  if not batch:
    return gather_nd_unbatched(params, indices)
  else:
    return vmap(gather_nd_unbatched, (0, 0), 0)(params, indices)

Note that batch_dims=0 would mean batch=False and batch_dims=1 would mean batch=True, so it doesn't implement batch_dims > 1.