jax_internals_override¶

probly.representation.sample.jax_functions.jax_internals_override(jax_sample_param_name: str) → Callable[[_JaxBoundFunctionWithInternals], _JaxFunction][source]¶
probly.representation.sample.jax_functions.jax_internals_override(*, jax_sample_param_pos: int) → Callable[[_JaxBoundFunctionWithInternals], _JaxFunction]: Decorator to convert a function taking a sample array argument.