Gather

Given data tensor of rank r >= 1, and indices tensor of rank q, gather entries of the axis dimension of data (by default outer-most one as axis=0) indexed by indices, and concatenates them in an output tensor of rank q + (r - 1). It is an indexing operation that indexes into the input data…

Abstract Signature:

Gather(data: Tensor, indices, axis: int)

PyTorchJAX (Core)TensorFlow

PyTorch

API: torch.gather

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: jax.numpy.take_along_axis

Strategy: Plugin (gather_adapter)

TensorFlow

API: tf.gather

Strategy: Direct Mapping

Official Docs ↗