jax_take_along_axis

probly.representation.jax_functions.jax_take_along_axis(arr: JaxArrayLike, /, indices: JaxArrayLike, axis: int | None = -1, mode: str | None = None, fill_value: Any = None) → jax.Array[source]

Take values along an axis, mirroring jax.numpy.take_along_axis.

Parameters:

  • arr – The array to take values from.

  • indices – The indices to take.

  • axis – The axis to take along. None flattens the input first.

  • mode – How out-of-bounds indices are handled.

  • fill_value – The value used for out-of-bounds indices in "fill" mode.

Returns:

The gathered array.