JAX equivalent to TensorFlow tf.gather_nd()

[conceptofmind]

Hi,

Is there an equivalent to tf.gather_nd() in Jax?

For example, I have a defined a function to gather indices:

def _gather(x, indices, gather_axis):
    all_indices = jnp.full(indices.shape)
    gather_locations = jnp.reshape(indices, [indices.shape.num_elements()])
    gather_indices = []
    for axis in range(len(indices.shape)):
        if axis == gather_axis:
            gather_indices.append(gather_locations)
        else:
            gather_indices.append(all_indices[:, axis])

    gather_indices = jnp.stack(gather_indices, axis=-1)
    gathered = tf.gather_nd(x, gather_indices)
    reshaped = jnp.reshape(gathered, indices.shape)
    return reshaped

I greatly appreciate any help.

Thank you,

Enrico

[YouJiacheng]

def jax_gather_nd(params, indices):
    tuple_indices = tuple(indices[..., i] for i in range(indices.shape[-1]))
    return params[tuple_indices]

In your case:

def _gather(x, indices, gather_axis):
    all_indices = jnp.full(indices.shape)
    gather_locations = jnp.reshape(indices, [indices.shape.num_elements()])
    gather_indices = []
    for axis in range(len(indices.shape)):
        if axis == gather_axis:
            gather_indices.append(gather_locations)
        else:
            gather_indices.append(all_indices[:, axis])
    gathered = x[tuple(gather_indices)]
    reshaped = jnp.reshape(gathered, indices.shape)
    return reshaped

[sharadmv]

You could refer also to TFP's implementation of gather_nd (TFP reimplements TF functions in terms of JAX).