jax_reshape_function

probly.representation.sample.jax_functions.jax_reshape_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 reshape wrapper for sample arrays.

The sample axis survives a reshape if it maps onto exactly one axis of the result; this mirrors the numpy backend. Jax arrays are always C-contiguous, so order="A" is equivalent to order="C".