Flax NNX¶

tinax.nnx captures independent graph snapshots, restores them with explicit copy policy, and clones mutable variables while preserving graph aliases and concrete module types.

from tinax.nnx import clone_graph, restore_graph, snapshot_graph

snapshot = snapshot_graph(model)
independent_model = restore_graph(snapshot, copy=True)
clone = clone_graph(model)

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.