You are viewing the development documentation (main). Switch to the stable release docs.

jax_copy_function¶

probly.representation.sample.jax_functions.jax_copy_function(func: Callable, args: tuple[Any, ...], kwargs: dict[str, Any], create_sample: JaxSampleCreator, array: JaxLike, sample_axis: int, weights: jax.Array | None) → Any[source]¶: Implementation of the jax copy wrapper for sample arrays.