brainevent.general_batching_rule

brainevent.general_batching_rule#

brainevent.general_batching_rule(prim, args, axes, **kwargs)[source]#

General-purpose batching rule for custom JAX primitives.

Implements batching by separating batched and non-batched arguments, moving all batch dimensions to axis 0, and then applying the primitive to each element in the batch via jax.lax.scan.

This function is registered as the default batching rule for every XLACustomKernel during initialization.

Parameters:

  • prim (Primitive) – The JAX primitive operation to be batched.

  • args (Sequence) – Input arguments to the primitive.

  • axes (Sequence[int | None]) – Batch dimension index for each argument. None indicates that the corresponding argument is not batched.

  • **kwargs – Additional keyword arguments forwarded to the primitive.

Return type:

Tuple[Any, Any]

Returns:

  • outs (pytree) – The batched outputs from applying the primitive.

  • out_dim (pytree) – A pytree with the same structure as outs. When at least one operand is batched, every leaf is 0 (the batch dimension is the leading axis of each output). When no operand is batched (every entry of axes is None), every leaf is None to signal that the outputs carry no batch dimension (the value JAX uses internally as the not_mapped sentinel).

Notes

All batch dimensions are moved to axis 0 before scanning. The scan carry is unused (always 0); only the stacked scan outputs are returned. As a special case, if no operand is batched the function short-circuits: it binds the primitive once on the original arguments and reports every output as unbatched, avoiding an empty jax.lax.scan (which would raise ValueError: scan got no values to scan over).

See also

XLACustomKernel.register_general_batching

Registers this function as the batching rule for a primitive.

XLACustomKernel.def_batching_rule

Override with a custom batching rule.

Examples

>>> import functools
>>> from jax.interpreters import batching
>>> batching.primitive_batchers[my_prim] = functools.partial(
...     general_batching_rule, my_prim
... )