TakeAlongAxis

Take elements from an array.

Abstract Signature:

TakeAlongAxis(arr: Tensor, indices: Tensor, axis: int | None = -1, mode: str | None, fill_value)

JAX (Core)

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

NumPy

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