Functional Transforms

JAX functional transformations (vmap, scan) adapted for Taktiny PyTree models.

taktiny.utils.transforms.vmap(fun=None, in_axes=0, out_axes=0, axis_name=None, axis_size=None, spmd_axis_name=None, sum_match=False)

Vectorize arbitrary PyTrees with JAX’s vmap semantics.

Supports vmap(function), @vmap and @vmap(…). Positional inputs follow in_axes; keyword-argument array leaves map along axis zero, as in JAX. out_axes=None retains JAX’s unmapped-output meaning, not scan’s discard meaning. No module metadata is changed and RNGs are not split implicitly. Return updated state explicitly; use Stack for module-state conveniences.

Return type:

Any

Parameters:

• fun (F | None)

• in_axes (int | None | Sequence[Any])

• out_axes (Any)

• axis_name (Any | None)

• axis_size (int | None)

• spmd_axis_name (Any | tuple[Any, ...] | None)

• sum_match (bool)

Examples

>>> import jax.numpy as jnp
>>> @vmap(in_axes=(0, None))
... def scale(x, factor):
...     return x * factor
>>> scale(jnp.arange(3), 2).tolist()
[0, 2, 4]

taktiny.utils.transforms.scan(fun=None, *, in_axes=0, out_axes=0, length=None, reverse=False, unroll=1, _split_transpose=False)

Build a generic scan callable with configurable input and output axes.

The transformed function accepts (init, xs, *args, **kwargs). Extra arguments are broadcast across iterations and passed to the scan body after carry and x.

in_axes is an integer, None, or a PyTree prefix of xs. Integer axes select the iteration dimension (negative axes are supported); None broadcasts a subtree unchanged. All mapped dimensions must have equal length. Supply length when there are no mapped leaves, including when xs is None.

out_axes is an integer or PyTree prefix of body outputs, selecting where each stacked iteration axis appears. None discards a subtree before stacking; it does not select an invariant or final output. Put a single final result in carry. Final carry is never rearranged by out_axes.

Carry structure, shapes and dtypes remain fixed, as in lax.scan. reverse, unroll and _split_transpose are forwarded to JAX. Reversing execution does not reverse the returned output order. Module metadata is untouched, and mutable state is not implicitly preserved. Thread RNGs/state through carry or supply independent keys in xs. Use SeqStack for module conveniences.

Return type:

Any

Parameters:

• fun (F | None)

• in_axes (Any)

• out_axes (Any)

• length (int | None)

• reverse (bool)

• unroll (int | bool)

• _split_transpose (bool)

Examples

>>> import jax.numpy as jnp
>>> @scan(in_axes=1, out_axes=-1)
... def accumulate(carry, x, *, factor=1):
...     carry = carry + x * factor
...     return carry, carry
>>> final, history = accumulate(jnp.zeros(2), jnp.ones((2, 3)), factor=2)
>>> final.tolist(), history.shape
([6.0, 6.0], (2, 3))

>>> @scan(in_axes=None, out_axes=None, length=3)
... def repeat(carry, increment):
...     return carry + increment, carry
>>> final, history = repeat(0, 2)
>>> int(final), history
(6, None)