jax_stack_function¶ probly.representation.sample.jax_functions.jax_stack_function(func: Callable, args: tuple[Any, ...], kwargs: dict[str, Any]) → Any[source]¶ Implementation of the jax stack wrapper for sample arrays.