sample.jax_functions

JAX function implementations for sample arrays.

Classes

JaxSampleCreator	Protocol for creating sample arrays.
JaxSampleInternals	Internal information about a sample array.

Functions

jax_concatenate_function	Implementation of the jax concatenate wrapper for sample arrays.
jax_conj_function	Implementation of the sample-axis-preserving jax conjugate wrapper for sample arrays.
jax_copy_function	Implementation of the jax copy wrapper for sample arrays.
jax_expand_dims_function	Implementation of the jax expand_dims wrapper for sample arrays.
jax_function_override	Decorator to convert a bound jax function to a jax function.
jax_internals_override	Decorator to convert a function taking a sample array argument.
jax_matrix_transpose_function	Implementation of the jax matrix transpose wrapper for sample arrays.
jax_moveaxis_function	Implementation of the jax moveaxis wrapper for sample arrays.
jax_reduction_function	Implementation of axis-reducing jax functions with a keepdims parameter.
jax_reshape_function	Implementation of the jax reshape wrapper for sample arrays.
jax_sample_internals	Get internals for a sample array.
jax_squeeze_function	Implementation of the jax squeeze wrapper for sample arrays.
jax_stack_function	Implementation of the jax stack wrapper for sample arrays.
jax_swapaxes_function	Implementation of the jax swapaxes wrapper for sample arrays.
jax_take_along_axis_function	Implementation of the jax take_along_axis wrapper for sample arrays.
jax_transpose_function	Implementation of the jax transpose wrapper for sample arrays.