jax_take_along_axis_function

probly.representation.sample.jax_functions.jax_take_along_axis_function(func: Callable, args: tuple[Any, ...], kwargs: dict[str, Any], create_sample: JaxSampleCreator, array: JaxLike, sample_axis: int, weights: jax.Array | None) → Any

Implementation of the jax take_along_axis wrapper for sample arrays.

The sample axis only survives if the gather runs along a different axis and leaves the number of samples untouched.