jax.lax.index_in_dim

jax.lax.index_in_dim(operand, index, axis=0, keepdims=True)

Convenience wrapper around slice to perform int indexing.

Parameters

• operand (Any) –

• index (int) –

• axis (int) –

• keepdims (bool) –

Return type

Any