jax.lax.index_take

jax.lax.index_take(src, idxs, axes)

Parameters

• src (Any) –

• idxs (Any) –

• axes (Sequence[int]) –

Return type

Any