Take

Returns a new tensor with the elements of input at the given indices (flattened).

Abstract Signature:

Take(input: Tensor, index: Tensor)

PyTorch

API: torch.take
Strategy: Direct Mapping
Official Docs ↗

JAX (Core)

API: jax.numpy.take
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: numpy.take
Strategy: Direct Mapping
Official Docs ↗

Keras

API: keras.ops.take
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.gather
Strategy: Macro 'tf.gather(tf.reshape({input}, [-1]), {index})'

Apple MLX

API: mlx.core.take
Strategy: Macro 'mlx.core.take(mlx.core.flatten({input}), {index})'

Flax NNX

API: jax.numpy.take
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.numpy.take
Strategy: Direct Mapping
Official Docs ↗