jax_concatenate_function

probly.representation.sample.jax_functions.jax_concatenate_function(func: Callable, args: tuple[Any, ...], kwargs: dict[str, Any]) Any[source]

Implementation of the jax concatenate wrapper for sample arrays.