Containers

Sequential and structured container modules.

class taktiny.nn.Sequential(modules)

Bases: Module

Apply modules in order, feeding each output into the next module.

Extra positional and keyword arguments are broadcast to every layer; they are not filtered by signature. Tuple/dict outputs remain a single input PyTree and are not unpacked. An empty sequence is the identity operation. Slices retain references to the original child modules.

Parameters:

modules (Sequence[Module]) – The sequence of modules to chain.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> model = nn.Sequential([nn.Dropout(0.1), nn.Dropout(0.2)])
>>> with nn.set_context_rng(nn.Rngs(0)):
...     y = model(jnp.ones((8, 4)))
>>> y.shape
(8, 4)

class taktiny.nn.Stack(modules, *, axis_name=None, spmd_axis_name=None)

Bases: Module

A module stack that uses vmap to apply stacked modules.

All layers must share type, static configuration, leaf shapes and dtypes. Creates stacked arrays without changing the original modules. in_axes maps positional inputs, with None broadcasting an input to every layer. A tuple provides one axis/PyTree specification per positional argument. Keyword arguments are always broadcast unchanged. out_axes controls only result axes; stored module state always uses leading axis zero.

Array and owned-Rngs updates are retained. Static configuration and state shapes/dtypes must not change during a call. Under jit, pass a stateful container in and return it to carry updates between compiled calls. An outer context RNG cannot be captured by vmap; use independent owned RNGs or map independent keys to a callback that establishes its own context. Parameter logical names and partition specs acquire a replicated layer axis; each mapped call sees the original per-layer metadata.

Individual layer restrictions still apply: BatchNorm’s current running statistics update is eager-only; use eval mode or track_running_stats=False.

Parameters:

• modules (Iterable[Module]) – The modules to stack and vectorize.

• axis_name (Any | None) – The name of the mapped axis. Defaults to None.

• spmd_axis_name (Any | tuple[Any, ...] | None) – The name of the SPMD mapped axis. Defaults to None.

Example

>>> from taktiny import nn
>>> import jax, jax.numpy as jnp
>>> layers = nn.Stack([nn.Dropout(0.5, rngs=nn.Rngs(i)) for i in range(2)])
>>> @jax.jit
... def step(model, x):
...     y = model(x, in_axes=None)
...     return y, model
>>> y, layers = step(layers, jnp.ones((8, 4)))
>>> y.shape
(2, 8, 4)

class taktiny.nn.SeqStack(modules, *, reverse=False, unroll=1, split_transpose=False)

Bases: Module

A sequential module stack that uses scan to apply stacked modules.

Consecutive layers with identical types, static configuration, leaf shapes and dtypes are grouped into separate scans. New stacked arrays are created; the original modules are not updated. The callback receives an individual layer and carry, and must return (next_carry, output). Carry structure, shapes and dtypes must be constant within each scan. Group outputs must have matching structures, trailing shapes and dtypes, or all be None. Outputs retain original layer order even when execution is reversed.

Array and owned-Rngs updates are stored back into the stack. Calls must not change module structure, static configuration, or state shapes/dtypes. For stateful JIT execution, pass this container into the compiled function and return it alongside the result. Context RNGs must be threaded through carry and installed inside the callback, not captured from outside scan. Logical names and partition specs gain a replicated leading layer axis; the callback sees the original per-layer metadata.

Individual layer restrictions still apply: BatchNorm’s current running statistics update is eager-only; use eval mode or track_running_stats=False.

Parameters:

• modules (Iterable[Module]) – The sequence of modules to stack and scan over.

• reverse (bool) – Whether to scan in reverse. Defaults to False.

• unroll (int | bool) – Loop unrolling factor. Defaults to 1.

• split_transpose (bool) – Whether to split transpose. Defaults to False.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> layers = nn.SeqStack([nn.Linear(4, 4, rngs=nn.Rngs(i)) for i in range(2)])
>>> def apply(layer, carry):
...     return layer(carry), None
>>> result, _ = layers(apply, jnp.ones((3, 4)))
>>> result.shape
(3, 4)

class taktiny.nn.List(modules)

Bases: Module

A list-like module container.

Stores the supplied module objects without copying them. Supports iteration, integer indexing and slicing; slices share child modules with the original. Parameters participate in PyTree traversal and recursive train/eval calls. This is a storage container, not a callable pipeline.

Parameters:

modules (Sequence[Module]) – The sequence of modules to store.

Example

>>> from taktiny import nn
>>> layers = nn.List([nn.Dropout(0.1), nn.Dropout(0.2)])
>>> len(layers[:1])
1

class taktiny.nn.Dict(modules)

Bases: Module

A dictionary-like module container indexed by stable string keys.

Keys must be nonempty strings without dots, preserving unambiguous parameter paths. Iteration follows insertion order. Children are shared, not copied; the container supports recursive train/eval and parameter traversal but does not define a forward call.

Parameters:

modules (Mapping[str, Module]) – A mapping of string keys to modules.

Example

>>> from taktiny import nn
>>> layers = nn.Dict({'dropout': nn.Dropout(0.1)})
>>> list(layers)
['dropout']