Utils

jax_gather_nd

jax_gather_nd(params, indices)

Slice a params array at a set of indices.

This is a reimplementation of TensorFlow's gather_nd function: link

Parameters:

• params (Float[Array, ' N *rest']) –

  an arbitrary array with leading axes of length $N$ upon which we shall slice.

• indices (Int[Array, ' M 1']) –

  an integer array of length $M$ with values in the range $[0, N)$ whose value at index $i$ will be used to slice params at index $i$.

Returns:

• Float[Array, ' M *rest'] –

  An arbitrary array with leading axes of length $M$.