jax_function_dispatch

probly.representation.jax_functions.jax_function_dispatch(*names: str, unpack: tuple[str, ...] = ()) → _SignaturePreservingDecorator

Decorate a native implementation with JAX function override dispatch.

Parameters:

• *names – Names of individual parameters participating in dispatch.

• unpack – Names of sequence parameters whose elements participate in dispatch. These parameters must contain reusable iterables.

Returns:

A decorator preserving the implementation’s signature, annotations, and documentation. Overrides receive the decorated public function as their dispatch key and the original positional and keyword arguments.

Raises:

ValueError – If selectors are empty, repeated, unknown, or variadic.

Examples

Use @jax_function_dispatch("a", "where") for a reduction, or @jax_function_dispatch(unpack=("arrays",)) for a sequence operation.