Flax NNX¶

tinax.nn captures independent graph snapshots, restores them with explicit copy policy, clones mutable variables while preserving graph aliases and concrete module types, displays objects through Flax's Treescope integration, and closes the validation gaps in Flax's attention-mask helpers.

from tinax.nn import clone_graph, display, make_causal_mask, restore_graph, snapshot_graph

snapshot = snapshot_graph(model)
independent_model = restore_graph(snapshot, copy=True)
clone = clone_graph(model)
display(clone)
mask = make_causal_mask(jnp.zeros((batch, length)), dtype=jnp.bool_)

display delegates to flax.nnx.display: inside IPython it renders each argument with Treescope, and outside it prints them. make_attention_mask, make_causal_mask, and combine_masks mirror the Flax helpers but reject rank-mismatched inputs, negative or boolean extra_batch_dims, and non-real dtypes, which Flax accepts silently. Model layers and flax.nnx.dot_product_attention / flax.nnx.MultiHeadAttention stay upstream: Flax already raises specific errors for their shape, depth, head-count, and dropout contracts.

snapshot_graph ¶

snapshot_graph(module: T) -> GraphState[T]

Capture independent variable state while preserving the module's internal graph aliases.

Parameters:

Name	Type	Description	Default
`module`	`T`	Flax NNX module to snapshot.	required

Returns:

Type	Description
`GraphState[T]`	A `flax.nnx.GraphState` holding independent variable state with the
`GraphState[T]`	module's internal aliases preserved.

Raises:

Type	Description
`TypeError`	If `module` is not a `flax.nnx.Module`.

restore_graph ¶

restore_graph(
    snapshot: GraphState[T], *, copy: bool
) -> T

Instantiate a graph snapshot with explicit NNX graph-copy policy.

Parameters:

Name	Type	Description	Default
`snapshot`	`GraphState[T]`	Graph state previously produced by `snapshot_graph`.	required
`copy`	`bool`	If `True`, copy variables while merging; if `False`, reuse them.	required

Returns:

Type	Description
`T`	A reconstructed Flax NNX module of the snapshot's concrete type.

Raises:

Type	Description
`TypeError`	If `snapshot` is not a `flax.nnx.GraphState`, `copy` is not a `bool`, or the snapshot does not describe a `flax.nnx.Module`.

clone_graph ¶

clone_graph(module: T) -> T

Clone mutable variables while preserving aliases and the concrete NNX module type.

Parameters:

Name	Type	Description	Default
`module`	`T`	Flax NNX module to clone.	required

Returns:

Type	Description
`T`	An independent clone of `module` with the same concrete type and preserved
`T`	internal aliases.

Raises:

Type	Description
`TypeError`	If `module` is not a `flax.nnx.Module`, or Flax NNX returns a clone of an unexpected concrete type.

display ¶

display(*args: object) -> None

Display objects with Treescope inside IPython and print them otherwise.

Delegates to flax.nnx.display: inside IPython each argument is rendered with the Treescope pretty-printer, and outside it each argument is printed with print. Calling it with no arguments renders nothing.

Parameters:

Name	Type	Description	Default
`*args`	`object`	Objects to display, typically Flax NNX modules, graph state, or arrays.	`()`

make_attention_mask ¶

make_attention_mask(
    query_input: Array | ndarray,
    key_input: Array | ndarray,
    *,
    pairwise_fn: Callable[..., Any] = multiply,
    extra_batch_dims: int = 0,
    dtype: Any = float32,
) -> Array

Build a broadcast attention mask from two flat inputs of equal rank.

Mirrors flax.nnx.make_attention_mask: inputs shaped [batch..., length] produce [batch..., 1, len_q, len_kv], plus one leading singleton axis per extra_batch_dims. Flax broadcasts rank-mismatched inputs silently; this requires the ranks to match. Host numpy.ndarray inputs are converted with jnp.asarray.

Parameters:

Name	Type	Description	Default
`query_input`	`Array \| ndarray`	Query-side flat input of shape `[batch..., len_q]`.	required
`key_input`	`Array \| ndarray`	Key-side flat input of shape `[batch..., len_kv]`.	required
`pairwise_fn`	`Callable[..., Any]`	Broadcasting elementwise function, `jnp.multiply` by default.	`multiply`
`extra_batch_dims`	`int`	Number of leading singleton axes to add.	`0`
`dtype`	`Any`	Mask dtype; a boolean, integer, or floating dtype.	`float32`

Returns:

Type	Description
`Array`	The broadcast mask as a JAX array.

Raises:

Type	Description
`TypeError`	If an input is not a `jax.Array` or `numpy.ndarray`, `pairwise_fn` is not callable, `extra_batch_dims` is not an integer (booleans rejected), or `dtype` is not a dtype.
`ValueError`	If the inputs do not share one rank of at least one, `extra_batch_dims` is negative, or `dtype` is not a real dtype.

make_causal_mask ¶

make_causal_mask(
    x: Array | ndarray,
    *,
    extra_batch_dims: int = 0,
    dtype: Any = float32,
) -> Array

Build a causal self-attention mask for inputs shaped [batch..., length].

Mirrors flax.nnx.make_causal_mask: the result has shape [batch..., 1, length, length], plus one leading singleton axis per extra_batch_dims. Host numpy.ndarray inputs are converted with jnp.asarray.

Parameters:

Name	Type	Description	Default
`x`	`Array \| ndarray`	Input array whose last axis is the sequence length.	required
`extra_batch_dims`	`int`	Number of leading singleton axes to add.	`0`
`dtype`	`Any`	Mask dtype; a boolean, integer, or floating dtype.	`float32`

Returns:

Type	Description
`Array`	The causal mask as a JAX array.

Raises:

Type	Description
`TypeError`	If `x` is not a `jax.Array` or `numpy.ndarray`, `extra_batch_dims` is not an integer (booleans rejected), or `dtype` is not a dtype.
`ValueError`	If `x` has no dimensions, `extra_batch_dims` is negative, or `dtype` is not a real dtype.

combine_masks ¶

combine_masks(
    *masks: Array | ndarray | None, dtype: Any = float32
) -> Array | None

Combine attention masks by logical AND, skipping None entries.

Mirrors flax.nnx.combine_masks and returns None when every entry is None. Host numpy.ndarray masks are converted with jnp.asarray.

Parameters:

Name	Type	Description	Default
`*masks`	`Array \| ndarray \| None`	Masks to combine; any entry may be `None`.	`()`
`dtype`	`Any`	Mask dtype; a boolean, integer, or floating dtype.	`float32`

Returns:

Type	Description
`Array \| None`	The combined mask as a JAX array, or `None` when no mask was given.

Raises:

Type	Description
`TypeError`	If a non-`None` entry is not a `jax.Array` or `numpy.ndarray`, or `dtype` is not a dtype.
`ValueError`	If `dtype` is not a real dtype, or the masks do not share one rank.