jax_reduction_function

probly.representation.sample.jax_functions.jax_reduction_function(func: Callable, args: tuple[Any, ...], kwargs: dict[str, Any]) → Any

Implementation of axis-reducing jax functions with a keepdims parameter.

Unlike the torch mirror there is no out handling: jax.numpy reductions reject any non-None out argument, so a sample array can never be an output buffer.

Mirroring the numpy backend, the sample weights are never injected into jax_average: they only apply to a reduction over the sample axis, which the caller expresses through sample_mean() instead.