API Reference

Auto-generated from docstrings via mkdocstrings. Sections are ordered from most foundational to most specialized; within each section, classes precede free functions and type aliases.

1. Top-level package

nnx

nnx — lightweight PyTorch training / eval / visualization toolkit.

The package is organized under nnx.nn (model, params, datasets, enums, nets, callbacks) and two top-level helpers (nnx.utils.Utils, nnx.vis_utils.VisUtils). The curated re-exports below give a flat surface for the most common imports without forbidding the deep paths existing notebook code relies on.

__version__ = _version('thekaveh-nnx') module-attribute

LRFinderResult dataclass

Result of an :func:lr_finder sweep.

Attributes:

Name	Type	Description
lrs	list[float]	list of learning rates actually exercised. Length matches losses. May be shorter than num_iter if the sweep early-exited due to loss divergence.
losses	list[float]	list of loss values, one per LR.
suggested_lr	float	the recommended max_lr for a subsequent real training run — the LR at the steepest-descent point of the smoothed loss curve.
figure	Figure	Plotly Figure plotting loss vs log(LR) with the suggested LR marked.

Source code in nnx/lr_finder.py

@dataclass(frozen=True, kw_only=True, slots=True)
class LRFinderResult:
    """Result of an :func:`lr_finder` sweep.

    Attributes:
        lrs: list of learning rates actually exercised. Length matches
            ``losses``. May be shorter than ``num_iter`` if the sweep
            early-exited due to loss divergence.
        losses: list of loss values, one per LR.
        suggested_lr: the recommended ``max_lr`` for a subsequent
            real training run — the LR at the steepest-descent
            point of the smoothed loss curve.
        figure: Plotly ``Figure`` plotting loss vs log(LR) with the
            suggested LR marked.
    """

    lrs: list[float]
    losses: list[float]
    suggested_lr: float
    figure: go.Figure

set_seed(seed: int, strict: bool = False) -> None

Pin every RNG that affects training and toggle cuDNN deterministic.

Parameters:

Name	Type	Description	Default
seed	int	integer seed shared across Python random, NumPy, and PyTorch (CPU + CUDA). Also written to os.environ["PYTHONHASHSEED"] so DataLoader workers started via the spawn method (default on Windows + macOS/Py3.8+) inherit a deterministic hash seed. Note: the current Python interpreter's hash state was fixed at startup — this assignment only affects spawned subprocesses. For full hash determinism in the current process, set PYTHONHASHSEED=<N> in the shell BEFORE invoking Python.	required
strict	bool	when True also calls torch.use_deterministic_algorithms(True) and sets CUBLAS_WORKSPACE_CONFIG. Slower and may raise on ops that lack a deterministic CUDA implementation; opt in only when full bit-for-bit reproducibility matters.	False

Source code in nnx/seeding.py

def set_seed(seed: int, strict: bool = False) -> None:
    """Pin every RNG that affects training and toggle cuDNN deterministic.

    Args:
        seed: integer seed shared across Python `random`, NumPy, and PyTorch
            (CPU + CUDA). Also written to `os.environ["PYTHONHASHSEED"]`
            so DataLoader workers started via the `spawn` method (default
            on Windows + macOS/Py3.8+) inherit a deterministic hash seed.
            Note: the current Python interpreter's hash state was fixed at
            startup — this assignment only affects spawned subprocesses.
            For full hash determinism in the current process, set
            `PYTHONHASHSEED=<N>` in the shell BEFORE invoking Python.
        strict: when True also calls torch.use_deterministic_algorithms(True)
            and sets CUBLAS_WORKSPACE_CONFIG. Slower and may raise on ops
            that lack a deterministic CUDA implementation; opt in only when
            full bit-for-bit reproducibility matters.
    """
    # PYTHONHASHSEED governs hash randomization in spawned subprocesses.
    # The current interpreter's hash state was fixed at startup and is
    # NOT affected by this assignment — but DataLoader workers using the
    # `spawn` start method (default on Windows + macOS/Py3.8+) and any
    # `subprocess.Popen` started downstream inherit this env var. Without
    # it, those children re-randomize their dict/set hash order, so any
    # code path that iterates a dict/set populated in the worker can
    # differ between runs even when every other RNG is seeded. Setting
    # it here is the cheap defensive move; the explicit caveat is in the
    # function docstring.
    os.environ["PYTHONHASHSEED"] = str(seed)

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    # cuDNN: turn off benchmarking (which picks the fastest kernel based on
    # input shapes and can introduce non-determinism) and turn on the
    # deterministic flag.
    if torch.backends.cudnn.is_available():
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    if strict:
        # Required for deterministic CUDA matmul / convolutions; see
        # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
        torch.use_deterministic_algorithms(True)

dataloader_worker_init_fn(worker_id: int) -> None

DataLoader worker_init_fn that pins each worker's numpy/python seed deterministically from the worker_id + the parent torch seed.

Pass as: DataLoader(..., worker_init_fn=dataloader_worker_init_fn).

Source code in nnx/seeding.py

def dataloader_worker_init_fn(worker_id: int) -> None:
    """DataLoader `worker_init_fn` that pins each worker's numpy/python seed
    deterministically from the worker_id + the parent torch seed.

    Pass as: `DataLoader(..., worker_init_fn=dataloader_worker_init_fn)`.
    """
    # torch.initial_seed() returns the base seed propagated to this worker.
    base_seed = torch.initial_seed() % (2**32)
    worker_seed = (base_seed + worker_id) % (2**32)
    np.random.seed(worker_seed)
    random.seed(worker_seed)

env_snapshot(force_refresh: bool = False) -> dict

Capture a snapshot of the runtime environment for reproducibility.

Returned dict is JSON-serializable. Includes Python / torch / numpy versions, GPU info if any, OS, and the git commit hash if running inside a git repo. Safe to call from anywhere — failures degrade to None per field rather than raising.

Result is memoized within the process (versions/hardware don't change between calls). Caveat: the git_commit / git_dirty fields are frozen at first call too, so a long session that commits mid-run records the session-start git state in later runs' metadata.yaml. Pass force_refresh=True to re-compute — useful in tests that mutate the environment, or to re-stamp git state.

Source code in nnx/seeding.py

def env_snapshot(force_refresh: bool = False) -> dict:
    """Capture a snapshot of the runtime environment for reproducibility.

    Returned dict is JSON-serializable. Includes Python / torch / numpy
    versions, GPU info if any, OS, and the git commit hash if running
    inside a git repo. Safe to call from anywhere — failures degrade to
    `None` per field rather than raising.

    Result is memoized within the process (versions/hardware don't
    change between calls). Caveat: the ``git_commit`` / ``git_dirty``
    fields are frozen at first call too, so a long session that commits
    mid-run records the session-start git state in later runs'
    metadata.yaml. Pass ``force_refresh=True`` to re-compute — useful
    in tests that mutate the environment, or to re-stamp git state.
    """
    global _ENV_SNAPSHOT_CACHE
    if _ENV_SNAPSHOT_CACHE is not None and not force_refresh:
        return dict(_ENV_SNAPSHOT_CACHE)
    import platform
    import subprocess

    def _git_commit() -> Optional[str]:
        try:
            return (
                subprocess.check_output(
                    ["git", "rev-parse", "HEAD"],
                    stderr=subprocess.DEVNULL,
                    timeout=2,
                )
                .decode()
                .strip()
            )
        # Broad catch is deliberate: env_snapshot() is opportunistic —
        # the caller isn't in a git repo (CI tarball install, fresh
        # `pip install` from PyPI, `tempfile.TemporaryDirectory` runs),
        # `git` isn't on PATH, the 2-second timeout fired, or the
        # subprocess crashed for any other reason. metadata.yaml just
        # omits the field; we never want this lookup to surface as an
        # exception to the user.
        except Exception:
            return None

    def _git_dirty() -> Optional[bool]:
        try:
            out = (
                subprocess.check_output(
                    ["git", "status", "--porcelain"],
                    stderr=subprocess.DEVNULL,
                    timeout=2,
                )
                .decode()
                .strip()
            )
            return bool(out)
        # Same opportunistic-fallback rationale as `_git_commit` above —
        # the dirty flag is metadata.yaml decoration, not a precondition.
        except Exception:
            return None

    def _nnx_version() -> Optional[str]:
        # PyPI distribution is `thekaveh-nnx` (PR #49) — the bare `nnx` name
        # is squatted by an abandoned JAX library, so a stale `version("nnx")`
        # here silently returned None on every clean install of the renamed
        # package, defeating metadata.yaml's reproducibility job. Mirror the
        # same lookup `nnx.__version__` uses (`src/nnx/__init__.py`).
        try:
            from importlib.metadata import version

            return version("thekaveh-nnx")
        except Exception:
            return None

    snap = {
        "nnx": _nnx_version(),
        "python": platform.python_version(),
        # torch.__version__ is a TorchVersion subclass that yaml.dump
        # can't serialize as a plain scalar — coerce to str so the
        # resulting metadata.yaml stays yaml.safe_load-compatible.
        "torch": str(torch.__version__),
        "numpy": str(np.__version__),
        "platform": platform.platform(),
        "cuda_available": bool(torch.cuda.is_available()),
        "cuda_device_count": int(torch.cuda.device_count()) if torch.cuda.is_available() else 0,
        "git_commit": _git_commit(),
        "git_dirty": _git_dirty(),
    }
    _ENV_SNAPSHOT_CACHE = dict(snap)
    return snap

2. Orchestrators

2.1. NNModel — supervised orchestrator

nnx.nn.nn_model.NNModel

Bases: _HubMixinBase

Top-level training/eval/predict wrapper around an nn.Module.

Inherits from :class:huggingface_hub.PyTorchModelHubMixin (when the thekaveh-nnx[hub] extra is installed) to gain save_pretrained / push_to_hub / from_pretrained. Without the extra installed, those three methods raise a clear ImportError pointing at the extra; no other NNModel functionality is affected.

Source code in nnx/nn/nn_model.py

class NNModel(_HubMixinBase):
    """Top-level training/eval/predict wrapper around an ``nn.Module``.

    Inherits from :class:`huggingface_hub.PyTorchModelHubMixin` (when the
    ``thekaveh-nnx[hub]`` extra is installed) to gain ``save_pretrained`` /
    ``push_to_hub`` / ``from_pretrained``. Without the extra installed,
    those three methods raise a clear ImportError pointing at the extra;
    no other NNModel functionality is affected.
    """

    def __init__(self, net_params: NNParams, params: NNModelParams):
        # NOTE: we deliberately do NOT call super().__init__() — the
        # PyTorchModelHubMixin base has no __init__ of its own (it's a
        # mixin that only contributes class-level methods), and even if
        # it grew one in a future hub release, the only side effect we'd
        # want is config-attribute initialization which we handle below.
        if net_params is None:
            raise ValueError("net_params must not be None")

        self.net_params = net_params
        self.params = params

        self.device = self.params.device()
        self.loss_fn = self.params.loss().to(self.device)
        self.net = self.params.net(params=net_params).to(self.device)

    def to_onnx(
        self,
        path: str,
        example_input: Union[torch.Tensor, tuple, np.ndarray],
        input_names: Optional[list[str]] = None,
        output_names: Optional[list[str]] = None,
        dynamic_batch: bool = True,
        opset_version: int = 17,
        dynamo: bool = False,
    ) -> str:
        """Export the underlying network to ONNX format.

        Args:
            path: output filename (e.g., "model.onnx").
            example_input: a tensor (or tuple of tensors for multi-input
                nets) with realistic shape/dtype used to trace the network.
            input_names: optional list of human-readable input port names.
            output_names: optional list of human-readable output port names.
            dynamic_batch: when True (default), marks dim 0 as dynamic so
                the exported model accepts any batch size at inference.
            opset_version: ONNX opset to target. 17 is broadly supported
                by current runtimes.
            dynamo: when False (default), uses the legacy TorchScript-based
                `torch.onnx.export` path — plain `pip install onnx` is
                enough. When True, dispatches to PyTorch's new
                `torch.export`-based exporter (default in torch>=2.9,
                supports >2 GB models via external data, faster). The
                dynamo path requires `onnxscript`; install via
                `pip install thekaveh-nnx[onnx-dynamo]`.

        Returns the path written. Network is put in eval mode for tracing.
        """
        if dynamo:
            # Lazy-import: keep `onnxscript` out of NNx's required deps so
            # plain `pip install thekaveh-nnx[onnx]` (legacy path) still works. If
            # the user opted in to `dynamo=True` without the extra, give
            # an error that names the install command instead of letting
            # torch surface a less actionable failure.
            try:
                import onnxscript  # noqa: F401
            except ImportError as e:
                raise ImportError(
                    "to_onnx(dynamo=True) requires the `onnxscript` package. "
                    "Install via `pip install thekaveh-nnx[onnx-dynamo]` (or `pip install onnxscript`)."
                ) from e

        # Normalize a single Tensor / np.ndarray to a length-1 tuple, then
        # coerce each element. Without the np.ndarray case in the singleton
        # check, a 2-D array like ``np.zeros((2, 4))`` falls into the
        # iterable branch and is unpacked row-by-row — torch.onnx.export
        # then sees a model with `N = first-dim` separate inputs instead of
        # the one input the caller meant.
        if isinstance(example_input, (torch.Tensor, np.ndarray)):
            example_input = (example_input,)
        example_input = tuple(
            (e.to(self.device) if isinstance(e, torch.Tensor) else torch.from_numpy(np.asarray(e)).to(self.device))
            for e in example_input
        )

        in_names = input_names or [f"input_{i}" for i in range(len(example_input))]
        out_names = output_names or ["output"]

        # Dynamic-shape spec is exporter-specific: the legacy TorchScript path
        # takes `dynamic_axes` (string-keyed dict of dim -> name); the dynamo
        # path takes `dynamic_shapes` (a pytree mirroring `example_input` whose
        # leaves are `{dim: torch.export.Dim(...)}`). Passing dynamic_axes with
        # dynamo=True triggers a UserWarning and on newer torch/onnxscript can
        # surface a hard `ConversionError` because dynamo emits `aten.sym_size`
        # ops that onnxscript can't always dispatch.
        dynamic_axes = None
        dynamic_shapes: Optional[tuple[dict[int, Any], ...]] = None
        if dynamic_batch:
            if dynamo:
                from torch.export import Dim

                batch = Dim("batch", min=1)
                dynamic_shapes = tuple({0: batch} for _ in example_input)
            else:
                dynamic_axes = {n: {0: "batch"} for n in in_names + out_names}

        # Snapshot training mode so the train → to_onnx → train-more pattern
        # doesn't silently strand the caller in .eval() (BatchNorm running-
        # stats / Dropout masking would then stay disabled on the next train
        # step). Matches the non-destructive contract every sibling inference
        # helper enforces (predict / evaluate / generate / diffusion.sample /
        # embed_texts / lr_finder / viz.activation_map / viz.attribute /
        # viz.netron_export). The bare `self.net.eval()` here was the lone
        # exception.
        was_training = self.net.training
        self.net.eval()
        # torch>=2.5 defaults torch.onnx.export to the dynamo-based exporter,
        # which requires `onnxscript`. We pass `dynamo` through explicitly so
        # the default (False) keeps the legacy tracing path regardless of
        # the installed torch version — plain `pip install onnx` is enough.
        try:
            export_accepts_dynamo = "dynamo" in inspect.signature(torch.onnx.export).parameters
            if dynamo:
                if not export_accepts_dynamo:
                    raise RuntimeError(
                        "to_onnx(dynamo=True) requires torch>=2.5 (the dynamo-based "
                        "ONNX exporter wasn't available before then). Upgrade torch or "
                        "call with dynamo=False."
                    )
                torch.onnx.export(
                    self.net,
                    example_input,
                    path,
                    input_names=in_names,
                    output_names=out_names,
                    dynamic_shapes=dynamic_shapes,
                    opset_version=opset_version,
                    dynamo=True,
                )
            elif export_accepts_dynamo:
                torch.onnx.export(
                    self.net,
                    example_input,
                    path,
                    input_names=in_names,
                    output_names=out_names,
                    dynamic_axes=dynamic_axes,
                    opset_version=opset_version,
                    dynamo=False,
                )
            else:
                torch.onnx.export(
                    self.net,
                    example_input,
                    path,
                    input_names=in_names,
                    output_names=out_names,
                    dynamic_axes=dynamic_axes,
                    opset_version=opset_version,
                )
        finally:
            if was_training:
                self.net.train()
        return path

    @staticmethod
    def from_checkpoint(checkpoint: NNCheckpoint) -> NNModel:
        model = NNModel(params=checkpoint.model_params, net_params=checkpoint.net_params)

        model.net.load_state_dict(checkpoint.net_state)

        return model

    # ------------------------------------------------------------------
    # HuggingFace Hub integration (inherited save_pretrained / push_to_hub /
    # from_pretrained from PyTorchModelHubMixin dispatch into these two
    # overrides).
    #
    # We override both because NNModel is NOT itself an nn.Module — its
    # weights live on `self.net`, and the default PyTorchModelHubMixin
    # implementation would call `self.state_dict()` and miss them. We
    # also need full control over config.json since the default mixin
    # auto-encoder hits the is_dataclass branch and emits asdict(NNParams)
    # which leaks the internal `_dims` cache and emits raw enums that
    # break JSON. Using the public NNParams.state() / NNModelParams.state()
    # round-trip keeps the on-Hub config compatible with NNRun's
    # hash-grouping form.
    # ------------------------------------------------------------------

    def _save_pretrained(self, save_directory) -> None:
        """Write the network weights + params config under ``save_directory``.

        Dispatched-to by :meth:`PyTorchModelHubMixin.save_pretrained`. The
        on-disk layout is the canonical Hub layout:

          - ``model.safetensors`` — ``self.net.state_dict()`` as safetensors.
          - ``config.json`` — ``{"net_params": <state>, "params": <state>}``,
            using the same ``.state()`` form NNRun uses for hashing.

        :meth:`PyTorchModelHubMixin.save_pretrained` additionally writes a
        ``README.md`` model card via ``generate_model_card()``; that's
        emitted on top of these two files by the base implementation.
        """
        from pathlib import Path

        try:
            from safetensors.torch import save_file
        except ImportError as e:  # pragma: no cover — gated by optional dep
            raise ImportError("save_pretrained requires the `hub` extra: `pip install thekaveh-nnx[hub]`.") from e

        save_dir = Path(save_directory)
        save_dir.mkdir(parents=True, exist_ok=True)

        # Detach + contiguous + clone matches the hygiene NNCheckpoint
        # applies on its safetensors path: drop autograd hooks, ensure
        # C-contiguous storage, and BREAK STORAGE SHARING — safetensors
        # rejects tied tensors (tok_embed/lm_head share storage on every
        # default TransformerNN), and .contiguous() is a no-op on an
        # already-contiguous shared view. load_state_dict reassembles the
        # tie on reload by copying both identical keys into the shared
        # parameter.
        tensors = {k: v.detach().contiguous().clone() for k, v in self.net.state_dict().items()}
        save_file(tensors, str(save_dir / _HUB_MODEL_FILENAME))

        config = {
            "net_params": self.net_params.state(),
            "params": self.params.state(),
        }
        # Explicit utf-8 — Hub config files round-trip through HuggingFace's
        # repo download path and can be read on any platform; relying on
        # the host locale's default encoding could mis-encode unicode
        # paths or non-ASCII tokenizer names round-tripped via
        # `params.state()`.
        with open(save_dir / _HUB_CONFIG_FILENAME, "w", encoding="utf-8") as f:
            json.dump(config, f, sort_keys=True, indent=2)

    @classmethod
    def _from_pretrained(
        cls,
        *,
        model_id: str,
        revision=None,
        cache_dir=None,
        force_download: bool = False,
        local_files_only: bool = False,
        token=None,
        map_location: str = "cpu",
        strict: bool = True,
        **model_kwargs,
    ) -> NNModel:
        """Rebuild an NNModel from a save_pretrained directory or Hub repo.

        Dispatched-to by :meth:`PyTorchModelHubMixin.from_pretrained`,
        which handles the remote-download path before calling this. Local
        paths skip the download. Either way, we read ``config.json`` to
        reconstruct the net params and ``NNModelParams`` via their public
        ``from_state`` constructors, then load ``model.safetensors`` into
        the freshly-built ``self.net``.

        ``map_location`` is forwarded as the safetensors ``device=`` (the
        net is then moved to ``self.device`` by ``NNModel.__init__``
        regardless). ``strict`` is forwarded to ``load_state_dict``; it
        defaults to True because the net is instantiated from the same
        config the weights were saved with, so any key mismatch indicates
        a corrupted or hand-edited artifact. Unrecognized ``model_kwargs``
        raise instead of being silently dropped — NNModel reconstructs
        entirely from ``config.json``.
        """
        # The mixin inspects NNModel.__init__'s signature and auto-injects
        # matching config.json entries ("net_params"/"params") as kwargs.
        # Both are rebuilt from config.json below via from_state — the
        # raw dicts are dropped knowingly. Anything else is a caller
        # error (a typo'd knob would otherwise vanish silently).
        model_kwargs.pop("net_params", None)
        model_kwargs.pop("params", None)
        if model_kwargs:
            raise TypeError(
                f"from_pretrained got unexpected model kwargs {sorted(model_kwargs)!r} — "
                "NNModel reconstructs entirely from the repo's config.json."
            )
        try:
            from safetensors.torch import load_file
        except ImportError as e:  # pragma: no cover — gated by optional dep
            raise ImportError("from_pretrained requires the `hub` extra: `pip install thekaveh-nnx[hub]`.") from e

        if os.path.isdir(model_id):
            config_path = os.path.join(model_id, _HUB_CONFIG_FILENAME)
            weights_path = os.path.join(model_id, _HUB_MODEL_FILENAME)
        else:
            from huggingface_hub import hf_hub_download

            config_path = hf_hub_download(
                repo_id=model_id,
                filename=_HUB_CONFIG_FILENAME,
                revision=revision,
                cache_dir=cache_dir,
                force_download=force_download,
                token=token,
                local_files_only=local_files_only,
            )
            weights_path = hf_hub_download(
                repo_id=model_id,
                filename=_HUB_MODEL_FILENAME,
                revision=revision,
                cache_dir=cache_dir,
                force_download=force_download,
                token=token,
                local_files_only=local_files_only,
            )

        with open(config_path, encoding="utf-8") as f:
            config = json.load(f)

        # We accept either form for back-compat with any future config
        # writer: nested under {"net_params": ..., "params": ...} (what
        # _save_pretrained writes today) or flat at the top level. The
        # nested form takes precedence.
        net_params_state = config.get("net_params", config)
        model_params_state = config.get("params", config)

        # resolve_from_state dispatches transformer configs to
        # NNTransformerParams so LM models round-trip through the Hub.
        net_params = NNParams.resolve_from_state(net_params_state)
        params = NNModelParams.from_state(model_params_state)

        model = cls(net_params=net_params, params=params)
        state_dict = load_file(weights_path, device=map_location)
        model.net.load_state_dict(state_dict, strict=strict)
        return model

    def freeze(self, *patterns: str) -> int:
        """Freeze parameters under ``self.net`` matching any of ``patterns``
        (fnmatch globs against the dotted parameter name). Returns the
        number of parameters newly frozen.

        Convenience wrapper around :func:`nnx.finetune.freezing.freeze`
        — use the standalone function when freezing a module that isn't
        ``self.net`` (e.g., a custom decoder hanging off this model).
        """
        from ..finetune.freezing import freeze as _freeze

        return _freeze(self.net, *patterns)

    def unfreeze(self, *patterns: str) -> int:
        """Mirror of :meth:`freeze` — set ``requires_grad=True`` on
        matching parameters."""
        from ..finetune.freezing import unfreeze as _unfreeze

        return _unfreeze(self.net, *patterns)

    def export_state_dict(self, path: str) -> str:
        """Save just ``self.net.state_dict()`` to ``path``.

        The file is a plain ``torch.save`` of a state-dict — loadable by
        any torch consumer without nnx installed, and by
        :func:`nnx.finetune.load_pretrained` for the fine-tuning round-trip.
        Companion to the NNCheckpoint format, which carries the params +
        idp wrapper alongside the weights; ``export_state_dict`` strips
        all of that and leaves just the weights.

        Returns ``path`` so calls can be chained.
        """
        torch.save(self.net.state_dict(), path)
        return path

    def train(
        self,
        params: NNTrainParams,
        callbacks: Optional[list[CallbackLike]] = None,
        train_step_fn: Optional[TrainStepFn] = None,
    ) -> NNRun:
        """Run the training loop and return the resulting NNRun.

        Args:
            params: dataloaders + optim + scheduler + epochs + seed. The
                train_loader is required; val_loader is optional (skips the
                per-epoch evaluation when absent).
            callbacks: optional list of `Callback` instances (or legacy
                `Callable[[List[IDP]], None]` for back-compat). Each hook
                runs at the documented lifecycle point (on_train_begin,
                on_epoch_begin/end, on_train_end).
            train_step_fn: optional override for the per-batch training
                step. When None (default), runs `default_train_step` —
                supervised forward → loss_fn(net(X), Y) → backward → step.
                Pass a `Callable[[TrainStepContext], NNEvaluationDataPoint]`
                for non-supervised paradigms (autoencoder, VAE, link
                prediction, recommendation, diffusion). The custom function
                is responsible for forward/backward/step and honoring the
                grad-clip/accumulation/AMP knobs the context carries. See
                `docs/concepts.md` and `examples/05_*.py`.

        Returns:
            An `NNRun` with per-iteration `idps`, persisted under
            `runs/<run.id>/` along with per-tag checkpoints. The same
            object is returned with the in-memory idps list attached.

        Raises:
            ValueError: if `params` is None, `params.train_loader` is
                None, or `params.optim` is invalid.
            FloatingPointError: from `default_train_step` if training
                loss becomes non-finite (custom `train_step_fn` hooks are
                responsible for their own divergence checks).
        """
        if params is None:
            raise ValueError("train params must be non-None")
        if params.train_loader is None:
            raise ValueError(
                "params.train_loader is required — set it directly or via with_train_loader(...) before train()."
            )
        if params.optim is None or not params.optim.is_valid():
            raise ValueError(f"train params has an invalid optim config: {params.optim!r}")
        if not any(p.requires_grad for p in self.net.parameters()):
            raise ValueError(
                "model has no trainable parameters — did you freeze('*')? Unfreeze something before train()."
            )

        # V1: seed every RNG before constructing the run so dataset shuffling,
        # weight init, dropout — anything stochastic — is reproducible. The
        # `seed` field only affects state() (and run.id) when explicitly set,
        # so back-compat for no-seed callers is preserved.
        if params.seed is not None:
            from ..seeding import set_seed

            set_seed(params.seed)

        validate: bool = params.val_loader is not None
        # Use self.net_params (always set in __init__) rather than
        # self.net.params: the latter is FeedFwdNN-specific and fails when
        # the caller substitutes a custom nn.Module post-construction
        # (the same idiom Trainer's GAN demo uses and that diffusion
        # demos use for DiffusionMLP). They're identical for the
        # standard supervised path; the rename is a back-compat-safe
        # robustness fix.
        run: NNRun = NNRun(train=params, model=self.params, net=self.net_params)

        optimizer = params.optim.name(
            net=self.net,
            lr_start=params.optim.max_lr,
            momentum=params.optim.momentum,
            weight_decay=params.optim.weight_decay,
            param_groups=params.optim.param_groups,
        )

        # Warm resume: load weights + optimizer state from a prior run's
        # checkpoint. The .opt.pt sidecar is best-effort — pre-resume
        # checkpoints don't have it, in which case we still load weights
        # but the optimizer starts fresh.
        if params.resume_from_run_id is not None:
            ckpt_type = Checkpoints(params.resume_from_checkpoint)
            resume_ckpt = NNCheckpoint.load(run=params.resume_from_run_id, type=ckpt_type)
            if resume_ckpt is None:
                raise ValueError(f"resume_from_run_id={params.resume_from_run_id!r}/{ckpt_type} not found on disk")
            self.net.load_state_dict(resume_ckpt.net_state)
            opt_state = NNCheckpoint.load_optimizer_state(
                run=params.resume_from_run_id,
                type=ckpt_type,
            )
            if opt_state is not None:
                optimizer.load_state_dict(opt_state)

        scheduler = self._build_scheduler(optimizer, params)
        scaler = self._build_grad_scaler()

        normalized_callbacks = self._normalize_callbacks(callbacks)

        idps: list[NNIterationDataPoint] = []
        # `len()` is not defined on iterable-style DataLoaders (IterableDataset).
        # Fall back to None so tqdm renders without a total instead of crashing.
        try:
            n_iter: Optional[int] = int(params.n_epochs * len(params.train_loader))
        except TypeError:
            n_iter = None
        best_checkpoint: Optional[NNCheckpoint] = NNCheckpoint.load(run=run.id, type=Checkpoints.BEST)

        Utils.print_table(header=False, title="Run Details...", data=Utils.flatten_dict(data=run.state()))

        ctx = _CallbackContext(model=self, run=run, optimizer=optimizer)
        for cb in normalized_callbacks:
            cb.on_train_begin(ctx)

        # Default to the standard supervised step when the caller doesn't
        # override. Custom step gets dispatched from inside the batch loop
        # below so the rest of train() (scheduler, callbacks, checkpoint
        # cadence, val loop, incremental save) is identical either way.
        # Explicit None check (not `or`) so a hypothetical callable that
        # happens to be falsy by __bool__ doesn't silently fall back.
        step_fn: TrainStepFn = default_train_step if train_step_fn is None else train_step_fn

        idx_iter = 0
        # Respect NNX_TQDM_DISABLE=1 in tests / CI / non-TTY environments so
        # the progress bar doesn't pollute output. Same env var works as
        # well in subprocess contexts where the user can't pass a flag.
        tqdm_disabled = os.environ.get("NNX_TQDM_DISABLE", "").lower() in {"1", "true", "yes"}
        with (
            torch.set_grad_enabled(True),
            tqdm(colour="blue", total=n_iter, desc="Training", disable=tqdm_disabled) as tqdm_bar,
            _CallbackFinalizer(normalized_callbacks, ctx),
        ):
            for idx_epoch in range(params.n_epochs):
                ctx.epoch = idx_epoch
                for cb in normalized_callbacks:
                    cb.on_epoch_begin(ctx)

                n_idps_before_epoch = len(idps)
                for idx_batch, batch in enumerate(params.train_loader):
                    step_ctx = TrainStepContext(
                        model=self,
                        batch=batch,
                        optimizer=optimizer,
                        scaler=scaler,
                        grad_clip_norm=params.optim.grad_clip_norm,
                        extra_metrics=params.extra_metrics,
                        accumulate_grad_batches=params.optim.accumulate_grad_batches,
                        batch_idx=idx_batch,
                        epoch_idx=idx_epoch,
                    )
                    train_edp = step_fn(step_ctx)

                    idps.append(
                        NNIterationDataPoint(
                            iter_idx=idx_iter,
                            epoch_idx=idx_epoch,
                            batch_idx=idx_batch,
                            train_edp=train_edp,
                            lr=optimizer.param_groups[0]["lr"],
                        )
                    )

                    idx_iter += 1
                    tqdm_bar.update(1)

                if len(idps) == n_idps_before_epoch:
                    # Zero batches this epoch: first epoch would crash on
                    # idps[-1] below; later epochs would silently attach
                    # this epoch's val_edp to the PREVIOUS epoch's last
                    # idp and reuse its stale train_edp.
                    raise ValueError(
                        f"train_loader yielded no batches in epoch {idx_epoch} — check batch_size vs "
                        "dataset size with drop_last=True, or whether the loader is a one-shot iterable."
                    )

                val_edp = (
                    self.evaluate(loader=params.val_loader, extra_metrics=params.extra_metrics) if validate else None
                )
                idps[-1] = idps[-1].with_val_edp(val_edp)

                checkpoint = self._save_checkpoints(
                    idp=idps[-1],
                    run_id=run.id,
                    idx_epoch=idx_epoch,
                    n_epochs=params.n_epochs,
                    best_checkpoint=best_checkpoint,
                    save_phase_checkpoints=params.save_phase_checkpoints,
                    optimizer=optimizer,
                )
                # In-memory best_checkpoint tracking must use the same
                # comparison as the on-disk BEST write inside
                # _save_checkpoints (val→train, error→loss, +inf fall-through).
                # Without this, val_loader=None runs would silently overwrite
                # best_checkpoint every epoch (because checkpoint.idp.val_edp
                # is None there) while the on-disk BEST tracks training error,
                # diverging the two views of "best".
                if best_checkpoint is None or _best_err(checkpoint) < _best_err(best_checkpoint):
                    best_checkpoint = checkpoint

                self._step_scheduler(scheduler, val_edp, train_edp)
                self._update_tqdm_postfix(tqdm_bar, optimizer, val_edp, train_edp)

                # Incremental persistence: write idps.csv + run.yaml after
                # every epoch. KeyboardInterrupt / OOM during training now
                # leaves a partial-but-loadable run on disk. The extra
                # writes are O(idps so far) per epoch — negligible vs the
                # checkpoint write that already happens.
                run.with_idps(idps).save()

                ctx.idp = idps[-1]
                ctx.idps = idps
                for cb in normalized_callbacks:
                    cb.on_epoch_end(ctx)

                if ctx.should_stop:
                    break

        saved = run.with_idps(idps).save()
        print()
        runs_root_path = os.path.join(os.getcwd(), "runs", run.id)
        print(f"Run saved to {runs_root_path}")
        return saved

    def evaluate(self, loader: DataLoader, extra_metrics=None) -> NNEvaluationDataPoint:
        """Aggregate predictions across all batches in `loader` and compute
        a single NNEvaluationDataPoint. Aggregating (rather than averaging
        per-batch metrics) gives correct sample-weighted f1/precision/recall
        when the final batch is short.

        Raises ValueError if the loader yields zero batches — previously
        produced NaN metrics silently from np.mean over an empty list.
        """
        # Ensure loss_fn lives on the same device as the model — guards
        # against callers reassigning self.device after construction.
        self.loss_fn = self.loss_fn.to(self.device)
        # Snapshot training-mode for non-destructive restore (matches the
        # convention already used by `nnx.viz.activation_map` and
        # `nnx.lr_finder`). Without this, a caller doing the common
        # train → evaluate → train-more pattern silently leaves the net
        # in `.eval()` mode after evaluate(); BatchNorm / Dropout layers
        # would behave incorrectly on the next batch unless the caller
        # remembered to call `self.net.train()` themselves.
        was_training = self.net.training
        self.net.eval()

        all_Y: list[np.ndarray] = []
        all_Y_hat: list[np.ndarray] = []
        loss_sum = 0.0
        n_samples = 0

        try:
            with torch.no_grad():
                for batch in loader:
                    _, Y, Y_hat_logits, Y_hat = self._fwd_pass(batch)
                    batch_n = int(Y.size(0))
                    # Aggregate predictions / labels across the entire loader so
                    # metrics are computed on the full eval set, not per-batch.
                    all_Y.append(Y.cpu().numpy())
                    all_Y_hat.append(Y_hat.cpu().numpy())
                    # Sum-weight the loss by samples; divide once at the end.
                    loss_sum += float(self.loss_fn(Y_hat_logits, Y).detach()) * batch_n
                    n_samples += batch_n
        finally:
            if was_training:
                self.net.train()

        if n_samples == 0:
            raise ValueError("evaluate() loader produced zero samples")

        Y_concat = np.concatenate(all_Y)
        Y_hat_concat = np.concatenate(all_Y_hat)

        return (
            NNEvaluationDataPoint.of(Y=Y_concat, Y_hat=Y_hat_concat, extra_metrics=extra_metrics)
            .with_loss(value=loss_sum / n_samples)
            .with_error(value=float(1 - (Y_concat == Y_hat_concat).sum() / n_samples))
        )

    def predict(self, X) -> PredictResult:
        """Run the network in eval mode and return logits + argmax classes.

        Accepts any of:

        - ``np.ndarray`` (single input tensor) — historical API.
        - ``tuple[np.ndarray, ...]`` — for multi-input networks.
        - ``torch.Tensor`` / ``tuple[torch.Tensor, ...]`` — skips the numpy
          conversion when callers already have tensors.
        - ``DataLoader`` — iterates the loader, runs predictions per batch,
          concatenates and returns the full result. Y labels in the batch
          (if present) are ignored.

        Returns a ``PredictResult`` (a ``NamedTuple`` of (logits, classes))
        that unpacks like the original 2-tuple.

        Non-destructive: ``self.net.training`` is snapshotted before
        switching to ``eval()`` and restored on exit (matches
        ``NNModel.evaluate``, ``nnx.viz.activation_map``, and
        ``nnx.lr_finder``). Without this, a caller doing the common
        train → predict → train-more pattern silently leaves the net
        in ``.eval()`` mode.
        """
        was_training = self.net.training
        self.net.eval()

        try:
            if isinstance(X, DataLoader):
                logits_chunks: list[np.ndarray] = []
                classes_chunks: list[np.ndarray] = []
                with torch.no_grad():
                    for batch in X:
                        # net.unpack_batch handles both (X, Y) tuples and PyG Data,
                        # returning the X-tuple. The label is discarded for predict.
                        X_in, _ = self.net.unpack_batch(batch)
                        X_in = tuple(x.to(self.device) for x in X_in)
                        logits = self.net(*X_in).cpu().numpy()
                        # NeighborLoader subgraphs: only the leading seed
                        # rows are this batch's nodes (see
                        # GraphNNBase.seed_count) — without the slice,
                        # predictions for sampled neighbors pollute the
                        # output and the row count exceeds the loader's
                        # node set.
                        seed_count = getattr(self.net, "seed_count", None)
                        if seed_count is not None:
                            n_seed = seed_count(batch)
                            if n_seed is not None:
                                logits = logits[:n_seed]
                        logits_chunks.append(logits)
                        classes_chunks.append(logits.argmax(axis=1))
                return PredictResult(
                    logits=np.concatenate(logits_chunks),
                    classes=np.concatenate(classes_chunks),
                )

            # Single input (any of: ndarray, Tensor, or a tuple thereof).
            if not isinstance(X, tuple):
                X = (X,)

            def _to_tensor(x):
                if isinstance(x, torch.Tensor):
                    return x.to(self.device)
                # Fall through to numpy → tensor for arrays and array-likes.
                return torch.from_numpy(np.asarray(x)).to(self.device)

            X_t = tuple(_to_tensor(x) for x in X)

            with torch.no_grad():
                Y_hat_logits = self.net(*X_t).cpu().numpy()
                Y_hat = Y_hat_logits.argmax(axis=1)
                return PredictResult(logits=Y_hat_logits, classes=Y_hat)
        finally:
            if was_training:
                self.net.train()

    def _fwd_pass(self, batch):
        """Standard supervised forward pass: unpack batch, move to device,
        run net, take argmax over class logits. Used by `default_train_step`
        and `evaluate()`; custom train_step_fn's may call this directly
        or roll their own forward pass."""
        X, Y = self.net.unpack_batch(batch)

        X = tuple(x.to(self.device) for x in X)
        Y = Y.to(self.device)

        Y_hat_logits = self.net(*X)
        # Graph nets score every node in the sampled subgraph, but only
        # the leading seed rows belong to this batch's split — see
        # GraphNNBase.seed_count for the leakage this prevents.
        seed_count = getattr(self.net, "seed_count", None)
        if seed_count is not None:
            n_seed = seed_count(batch)
            if n_seed is not None:
                Y_hat_logits = Y_hat_logits[:n_seed]
                Y = Y[:n_seed]
        Y_hat = Y_hat_logits.argmax(dim=1)

        return X, Y, Y_hat_logits, Y_hat

    def _train_step(
        self,
        batch,
        optimizer: torch.optim.Optimizer,
        scaler: Optional[torch.amp.GradScaler],
        grad_clip_norm: Optional[float] = None,
        extra_metrics=None,
        accumulate_grad_batches: int = 1,
        batch_idx: int = 0,
    ) -> NNEvaluationDataPoint:
        """Thin wrapper around :func:`default_train_step` kept for back-compat
        with any code that calls ``model._train_step(batch, ...)`` directly
        (e.g., a notebook that pre-dates the ``train_step_fn`` hook).

        **The :meth:`train` loop does NOT call this method.** It builds a
        :class:`TrainStepContext` and dispatches to
        ``train_step_fn or default_train_step`` directly. A subclass that
        overrides ``_train_step`` will therefore be ignored by ``train()`` —
        if you want a custom training step for ``train()``, pass it as the
        ``train_step_fn=`` kwarg instead.
        """
        return default_train_step(
            TrainStepContext(
                model=self,
                batch=batch,
                optimizer=optimizer,
                scaler=scaler,
                grad_clip_norm=grad_clip_norm,
                extra_metrics=extra_metrics,
                accumulate_grad_batches=accumulate_grad_batches,
                batch_idx=batch_idx,
                epoch_idx=0,
            )
        )

    def _build_scheduler(
        self,
        optimizer: torch.optim.Optimizer,
        params: NNTrainParams,
    ):
        # If params.scheduler has a `kind` attribute (set by the Schedulers
        # enum), dispatch on it; otherwise fall back to ReduceLROnPlateau
        # for backwards compatibility with existing notebook code.
        sched_params = params.scheduler
        kind = getattr(sched_params, "kind", None)

        if kind is None:
            return lr_scheduler.ReduceLROnPlateau(
                optimizer,
                mode="min",
                min_lr=sched_params.min_lr,
                factor=sched_params.factor,
                cooldown=sched_params.cooldown,
                patience=sched_params.patience,
                threshold=sched_params.threshold,
            )

        # When a `kind` is supplied, the params dataclass carries kind-specific
        # config. The enum's __call__ knows how to construct.
        return kind(optimizer=optimizer, params=sched_params, n_epochs=params.n_epochs)

    def _build_grad_scaler(self) -> Optional[torch.amp.GradScaler]:
        if getattr(self.params, "mixed_precision", False) and self.device.type == "cuda":
            return torch.amp.GradScaler("cuda")
        return None

    def _save_checkpoints(
        self,
        idp: NNIterationDataPoint,
        run_id: str,
        idx_epoch: int,
        n_epochs: int,
        best_checkpoint: Optional[NNCheckpoint],
        save_phase_checkpoints: bool = True,
        optimizer: Optional[torch.optim.Optimizer] = None,
    ) -> NNCheckpoint:
        checkpoint = NNCheckpoint(
            idp=idp, model_params=self.params, net_params=self.net_params, net_state=self.net.state_dict()
        )
        # The optimizer state-dict goes only into LAST and BEST sidecars
        # (resume points). Phase markers (FIRST/Q*) don't carry one.
        opt_state = optimizer.state_dict() if optimizer is not None else None

        # Phase markers at epoch boundaries — fractions are nominal (1/4, 2/4,
        # 3/4 of the planned epoch count); off-by-one allowed when n_epochs
        # isn't divisible by 4. See `phase_tag` for the small-`n_epochs`
        # caveat. Opt-out via NNTrainParams.save_phase_checkpoints.
        if save_phase_checkpoints:
            tag = phase_tag(idx_epoch, n_epochs)
            if tag is not None:
                checkpoint.save(run=run_id, type=tag)

        checkpoint.save(run=run_id, type=Checkpoints.LAST, optimizer_state=opt_state)

        # BEST tracking goes through the same _best_err helper used by
        # NNRun.save's cross-run comparison and by Trainer._save_checkpoint
        # — single source of truth for "what's the comparable error here"
        # (val→train, error→loss, +inf fall-through, tolerating None EDP
        # or None .error from custom train_step_fn paradigms).
        if best_checkpoint is None or _best_err(checkpoint) < _best_err(best_checkpoint):
            checkpoint.save(run=run_id, type=Checkpoints.BEST, optimizer_state=opt_state)

        return checkpoint

    def _step_scheduler(
        self,
        scheduler,
        val_edp: Optional[NNEvaluationDataPoint],
        train_edp: NNEvaluationDataPoint,
    ) -> None:
        # ReduceLROnPlateau wants a metric; other schedulers step on epoch index.
        if isinstance(scheduler, lr_scheduler.ReduceLROnPlateau):
            # Custom train_step_fn hooks may leave .error unset;
            # ReduceLROnPlateau.step(None) crashes inside float(). Use the
            # shared val→train, error→loss fallback resolver so the four
            # call sites (NNModel + Trainer × scheduler + tqdm) can't drift.
            metric = _resolve_metric(val_edp, train_edp)
            if metric is None:
                # No signal to feed the scheduler — skip the step. The user
                # picked a metric-driven scheduler without producing a metric;
                # better to no-op than to crash mid-train.
                return
            scheduler.step(metric)
        else:
            scheduler.step()

    def _update_tqdm_postfix(
        self,
        tqdm_bar,
        optimizer,
        val_edp: Optional[NNEvaluationDataPoint],
        train_edp: NNEvaluationDataPoint,
    ) -> None:
        lr = optimizer.param_groups[0]["lr"]
        # Custom train_step_fn hooks may leave .error unset — fall back to
        # .loss for display so the progress bar doesn't crash mid-train on
        # an `f"{None:.4f}"` format error. Same shared fallback resolver
        # used by _step_scheduler above.
        err = _resolve_metric(val_edp, train_edp)
        err_str = f"{err:.4f}" if err is not None else "n/a"
        tqdm_bar.set_postfix_str(f"error={err_str}, lr={lr:.4f}")

    @staticmethod
    def _normalize_callbacks(
        callbacks: Optional[list[CallbackLike]],
    ) -> list[Callback]:
        # Lazy import to keep nn_model.py importable before callbacks module exists.
        from .callbacks import Callback, _LegacyCallback

        if callbacks is None:
            return []
        out: list[Callback] = []
        for cb in callbacks:
            if isinstance(cb, Callback):
                out.append(cb)
            else:
                out.append(_LegacyCallback(cb))
        return out

evaluate(loader: DataLoader, extra_metrics=None) -> NNEvaluationDataPoint

Aggregate predictions across all batches in loader and compute a single NNEvaluationDataPoint. Aggregating (rather than averaging per-batch metrics) gives correct sample-weighted f1/precision/recall when the final batch is short.

Raises ValueError if the loader yields zero batches — previously produced NaN metrics silently from np.mean over an empty list.

Source code in nnx/nn/nn_model.py

def evaluate(self, loader: DataLoader, extra_metrics=None) -> NNEvaluationDataPoint:
    """Aggregate predictions across all batches in `loader` and compute
    a single NNEvaluationDataPoint. Aggregating (rather than averaging
    per-batch metrics) gives correct sample-weighted f1/precision/recall
    when the final batch is short.

    Raises ValueError if the loader yields zero batches — previously
    produced NaN metrics silently from np.mean over an empty list.
    """
    # Ensure loss_fn lives on the same device as the model — guards
    # against callers reassigning self.device after construction.
    self.loss_fn = self.loss_fn.to(self.device)
    # Snapshot training-mode for non-destructive restore (matches the
    # convention already used by `nnx.viz.activation_map` and
    # `nnx.lr_finder`). Without this, a caller doing the common
    # train → evaluate → train-more pattern silently leaves the net
    # in `.eval()` mode after evaluate(); BatchNorm / Dropout layers
    # would behave incorrectly on the next batch unless the caller
    # remembered to call `self.net.train()` themselves.
    was_training = self.net.training
    self.net.eval()

    all_Y: list[np.ndarray] = []
    all_Y_hat: list[np.ndarray] = []
    loss_sum = 0.0
    n_samples = 0

    try:
        with torch.no_grad():
            for batch in loader:
                _, Y, Y_hat_logits, Y_hat = self._fwd_pass(batch)
                batch_n = int(Y.size(0))
                # Aggregate predictions / labels across the entire loader so
                # metrics are computed on the full eval set, not per-batch.
                all_Y.append(Y.cpu().numpy())
                all_Y_hat.append(Y_hat.cpu().numpy())
                # Sum-weight the loss by samples; divide once at the end.
                loss_sum += float(self.loss_fn(Y_hat_logits, Y).detach()) * batch_n
                n_samples += batch_n
    finally:
        if was_training:
            self.net.train()

    if n_samples == 0:
        raise ValueError("evaluate() loader produced zero samples")

    Y_concat = np.concatenate(all_Y)
    Y_hat_concat = np.concatenate(all_Y_hat)

    return (
        NNEvaluationDataPoint.of(Y=Y_concat, Y_hat=Y_hat_concat, extra_metrics=extra_metrics)
        .with_loss(value=loss_sum / n_samples)
        .with_error(value=float(1 - (Y_concat == Y_hat_concat).sum() / n_samples))
    )

export_state_dict(path: str) -> str

Save just self.net.state_dict() to path.

The file is a plain torch.save of a state-dict — loadable by any torch consumer without nnx installed, and by :func:nnx.finetune.load_pretrained for the fine-tuning round-trip. Companion to the NNCheckpoint format, which carries the params + idp wrapper alongside the weights; export_state_dict strips all of that and leaves just the weights.

Returns path so calls can be chained.

Source code in nnx/nn/nn_model.py

def export_state_dict(self, path: str) -> str:
    """Save just ``self.net.state_dict()`` to ``path``.

    The file is a plain ``torch.save`` of a state-dict — loadable by
    any torch consumer without nnx installed, and by
    :func:`nnx.finetune.load_pretrained` for the fine-tuning round-trip.
    Companion to the NNCheckpoint format, which carries the params +
    idp wrapper alongside the weights; ``export_state_dict`` strips
    all of that and leaves just the weights.

    Returns ``path`` so calls can be chained.
    """
    torch.save(self.net.state_dict(), path)
    return path

freeze(*patterns: str) -> int

Freeze parameters under self.net matching any of patterns (fnmatch globs against the dotted parameter name). Returns the number of parameters newly frozen.

Convenience wrapper around :func:nnx.finetune.freezing.freeze — use the standalone function when freezing a module that isn't self.net (e.g., a custom decoder hanging off this model).

Source code in nnx/nn/nn_model.py

def freeze(self, *patterns: str) -> int:
    """Freeze parameters under ``self.net`` matching any of ``patterns``
    (fnmatch globs against the dotted parameter name). Returns the
    number of parameters newly frozen.

    Convenience wrapper around :func:`nnx.finetune.freezing.freeze`
    — use the standalone function when freezing a module that isn't
    ``self.net`` (e.g., a custom decoder hanging off this model).
    """
    from ..finetune.freezing import freeze as _freeze

    return _freeze(self.net, *patterns)

predict(X) -> PredictResult

Run the network in eval mode and return logits + argmax classes.

Accepts any of:

• np.ndarray (single input tensor) — historical API.
• tuple[np.ndarray, ...] — for multi-input networks.
• torch.Tensor / tuple[torch.Tensor, ...] — skips the numpy conversion when callers already have tensors.
• DataLoader — iterates the loader, runs predictions per batch, concatenates and returns the full result. Y labels in the batch (if present) are ignored.

Returns a PredictResult (a NamedTuple of (logits, classes)) that unpacks like the original 2-tuple.

Non-destructive: self.net.training is snapshotted before switching to eval() and restored on exit (matches NNModel.evaluate, nnx.viz.activation_map, and nnx.lr_finder). Without this, a caller doing the common train → predict → train-more pattern silently leaves the net in .eval() mode.

Source code in nnx/nn/nn_model.py

def predict(self, X) -> PredictResult:
    """Run the network in eval mode and return logits + argmax classes.

    Accepts any of:

    - ``np.ndarray`` (single input tensor) — historical API.
    - ``tuple[np.ndarray, ...]`` — for multi-input networks.
    - ``torch.Tensor`` / ``tuple[torch.Tensor, ...]`` — skips the numpy
      conversion when callers already have tensors.
    - ``DataLoader`` — iterates the loader, runs predictions per batch,
      concatenates and returns the full result. Y labels in the batch
      (if present) are ignored.

    Returns a ``PredictResult`` (a ``NamedTuple`` of (logits, classes))
    that unpacks like the original 2-tuple.

    Non-destructive: ``self.net.training`` is snapshotted before
    switching to ``eval()`` and restored on exit (matches
    ``NNModel.evaluate``, ``nnx.viz.activation_map``, and
    ``nnx.lr_finder``). Without this, a caller doing the common
    train → predict → train-more pattern silently leaves the net
    in ``.eval()`` mode.
    """
    was_training = self.net.training
    self.net.eval()

    try:
        if isinstance(X, DataLoader):
            logits_chunks: list[np.ndarray] = []
            classes_chunks: list[np.ndarray] = []
            with torch.no_grad():
                for batch in X:
                    # net.unpack_batch handles both (X, Y) tuples and PyG Data,
                    # returning the X-tuple. The label is discarded for predict.
                    X_in, _ = self.net.unpack_batch(batch)
                    X_in = tuple(x.to(self.device) for x in X_in)
                    logits = self.net(*X_in).cpu().numpy()
                    # NeighborLoader subgraphs: only the leading seed
                    # rows are this batch's nodes (see
                    # GraphNNBase.seed_count) — without the slice,
                    # predictions for sampled neighbors pollute the
                    # output and the row count exceeds the loader's
                    # node set.
                    seed_count = getattr(self.net, "seed_count", None)
                    if seed_count is not None:
                        n_seed = seed_count(batch)
                        if n_seed is not None:
                            logits = logits[:n_seed]
                    logits_chunks.append(logits)
                    classes_chunks.append(logits.argmax(axis=1))
            return PredictResult(
                logits=np.concatenate(logits_chunks),
                classes=np.concatenate(classes_chunks),
            )

        # Single input (any of: ndarray, Tensor, or a tuple thereof).
        if not isinstance(X, tuple):
            X = (X,)

        def _to_tensor(x):
            if isinstance(x, torch.Tensor):
                return x.to(self.device)
            # Fall through to numpy → tensor for arrays and array-likes.
            return torch.from_numpy(np.asarray(x)).to(self.device)

        X_t = tuple(_to_tensor(x) for x in X)

        with torch.no_grad():
            Y_hat_logits = self.net(*X_t).cpu().numpy()
            Y_hat = Y_hat_logits.argmax(axis=1)
            return PredictResult(logits=Y_hat_logits, classes=Y_hat)
    finally:
        if was_training:
            self.net.train()

to_onnx(path: str, example_input: Union[torch.Tensor, tuple, np.ndarray], input_names: Optional[list[str]] = None, output_names: Optional[list[str]] = None, dynamic_batch: bool = True, opset_version: int = 17, dynamo: bool = False) -> str

Export the underlying network to ONNX format.

Parameters:

Name	Type	Description	Default
path	str	output filename (e.g., "model.onnx").	required
example_input	Union[Tensor, tuple, ndarray]	a tensor (or tuple of tensors for multi-input nets) with realistic shape/dtype used to trace the network.	required
input_names	Optional[list[str]]	optional list of human-readable input port names.	None
output_names	Optional[list[str]]	optional list of human-readable output port names.	None
dynamic_batch	bool	when True (default), marks dim 0 as dynamic so the exported model accepts any batch size at inference.	True
opset_version	int	ONNX opset to target. 17 is broadly supported by current runtimes.	17
dynamo	bool	when False (default), uses the legacy TorchScript-based torch.onnx.export path — plain pip install onnx is enough. When True, dispatches to PyTorch's new torch.export-based exporter (default in torch>=2.9, supports >2 GB models via external data, faster). The dynamo path requires onnxscript; install via pip install thekaveh-nnx[onnx-dynamo].	False

Returns the path written. Network is put in eval mode for tracing.

Source code in nnx/nn/nn_model.py

def to_onnx(
    self,
    path: str,
    example_input: Union[torch.Tensor, tuple, np.ndarray],
    input_names: Optional[list[str]] = None,
    output_names: Optional[list[str]] = None,
    dynamic_batch: bool = True,
    opset_version: int = 17,
    dynamo: bool = False,
) -> str:
    """Export the underlying network to ONNX format.

    Args:
        path: output filename (e.g., "model.onnx").
        example_input: a tensor (or tuple of tensors for multi-input
            nets) with realistic shape/dtype used to trace the network.
        input_names: optional list of human-readable input port names.
        output_names: optional list of human-readable output port names.
        dynamic_batch: when True (default), marks dim 0 as dynamic so
            the exported model accepts any batch size at inference.
        opset_version: ONNX opset to target. 17 is broadly supported
            by current runtimes.
        dynamo: when False (default), uses the legacy TorchScript-based
            `torch.onnx.export` path — plain `pip install onnx` is
            enough. When True, dispatches to PyTorch's new
            `torch.export`-based exporter (default in torch>=2.9,
            supports >2 GB models via external data, faster). The
            dynamo path requires `onnxscript`; install via
            `pip install thekaveh-nnx[onnx-dynamo]`.

    Returns the path written. Network is put in eval mode for tracing.
    """
    if dynamo:
        # Lazy-import: keep `onnxscript` out of NNx's required deps so
        # plain `pip install thekaveh-nnx[onnx]` (legacy path) still works. If
        # the user opted in to `dynamo=True` without the extra, give
        # an error that names the install command instead of letting
        # torch surface a less actionable failure.
        try:
            import onnxscript  # noqa: F401
        except ImportError as e:
            raise ImportError(
                "to_onnx(dynamo=True) requires the `onnxscript` package. "
                "Install via `pip install thekaveh-nnx[onnx-dynamo]` (or `pip install onnxscript`)."
            ) from e

    # Normalize a single Tensor / np.ndarray to a length-1 tuple, then
    # coerce each element. Without the np.ndarray case in the singleton
    # check, a 2-D array like ``np.zeros((2, 4))`` falls into the
    # iterable branch and is unpacked row-by-row — torch.onnx.export
    # then sees a model with `N = first-dim` separate inputs instead of
    # the one input the caller meant.
    if isinstance(example_input, (torch.Tensor, np.ndarray)):
        example_input = (example_input,)
    example_input = tuple(
        (e.to(self.device) if isinstance(e, torch.Tensor) else torch.from_numpy(np.asarray(e)).to(self.device))
        for e in example_input
    )

    in_names = input_names or [f"input_{i}" for i in range(len(example_input))]
    out_names = output_names or ["output"]

    # Dynamic-shape spec is exporter-specific: the legacy TorchScript path
    # takes `dynamic_axes` (string-keyed dict of dim -> name); the dynamo
    # path takes `dynamic_shapes` (a pytree mirroring `example_input` whose
    # leaves are `{dim: torch.export.Dim(...)}`). Passing dynamic_axes with
    # dynamo=True triggers a UserWarning and on newer torch/onnxscript can
    # surface a hard `ConversionError` because dynamo emits `aten.sym_size`
    # ops that onnxscript can't always dispatch.
    dynamic_axes = None
    dynamic_shapes: Optional[tuple[dict[int, Any], ...]] = None
    if dynamic_batch:
        if dynamo:
            from torch.export import Dim

            batch = Dim("batch", min=1)
            dynamic_shapes = tuple({0: batch} for _ in example_input)
        else:
            dynamic_axes = {n: {0: "batch"} for n in in_names + out_names}

    # Snapshot training mode so the train → to_onnx → train-more pattern
    # doesn't silently strand the caller in .eval() (BatchNorm running-
    # stats / Dropout masking would then stay disabled on the next train
    # step). Matches the non-destructive contract every sibling inference
    # helper enforces (predict / evaluate / generate / diffusion.sample /
    # embed_texts / lr_finder / viz.activation_map / viz.attribute /
    # viz.netron_export). The bare `self.net.eval()` here was the lone
    # exception.
    was_training = self.net.training
    self.net.eval()
    # torch>=2.5 defaults torch.onnx.export to the dynamo-based exporter,
    # which requires `onnxscript`. We pass `dynamo` through explicitly so
    # the default (False) keeps the legacy tracing path regardless of
    # the installed torch version — plain `pip install onnx` is enough.
    try:
        export_accepts_dynamo = "dynamo" in inspect.signature(torch.onnx.export).parameters
        if dynamo:
            if not export_accepts_dynamo:
                raise RuntimeError(
                    "to_onnx(dynamo=True) requires torch>=2.5 (the dynamo-based "
                    "ONNX exporter wasn't available before then). Upgrade torch or "
                    "call with dynamo=False."
                )
            torch.onnx.export(
                self.net,
                example_input,
                path,
                input_names=in_names,
                output_names=out_names,
                dynamic_shapes=dynamic_shapes,
                opset_version=opset_version,
                dynamo=True,
            )
        elif export_accepts_dynamo:
            torch.onnx.export(
                self.net,
                example_input,
                path,
                input_names=in_names,
                output_names=out_names,
                dynamic_axes=dynamic_axes,
                opset_version=opset_version,
                dynamo=False,
            )
        else:
            torch.onnx.export(
                self.net,
                example_input,
                path,
                input_names=in_names,
                output_names=out_names,
                dynamic_axes=dynamic_axes,
                opset_version=opset_version,
            )
    finally:
        if was_training:
            self.net.train()
    return path

train(params: NNTrainParams, callbacks: Optional[list[CallbackLike]] = None, train_step_fn: Optional[TrainStepFn] = None) -> NNRun

Run the training loop and return the resulting NNRun.

Parameters:

Name	Type	Description	Default
params	NNTrainParams	dataloaders + optim + scheduler + epochs + seed. The train_loader is required; val_loader is optional (skips the per-epoch evaluation when absent).	required
callbacks	Optional[list[CallbackLike]]	optional list of Callback instances (or legacy Callable[[List[IDP]], None] for back-compat). Each hook runs at the documented lifecycle point (on_train_begin, on_epoch_begin/end, on_train_end).	None
train_step_fn	Optional[TrainStepFn]	optional override for the per-batch training step. When None (default), runs default_train_step — supervised forward → loss_fn(net(X), Y) → backward → step. Pass a Callable[[TrainStepContext], NNEvaluationDataPoint] for non-supervised paradigms (autoencoder, VAE, link prediction, recommendation, diffusion). The custom function is responsible for forward/backward/step and honoring the grad-clip/accumulation/AMP knobs the context carries. See docs/concepts.md and examples/05_*.py.	None

Returns:

Type	Description
NNRun	An NNRun with per-iteration idps, persisted under
NNRun	runs/<run.id>/ along with per-tag checkpoints. The same
NNRun	object is returned with the in-memory idps list attached.

Raises:

Type	Description
ValueError	if params is None, params.train_loader is None, or params.optim is invalid.
FloatingPointError	from default_train_step if training loss becomes non-finite (custom train_step_fn hooks are responsible for their own divergence checks).

Source code in nnx/nn/nn_model.py

def train(
    self,
    params: NNTrainParams,
    callbacks: Optional[list[CallbackLike]] = None,
    train_step_fn: Optional[TrainStepFn] = None,
) -> NNRun:
    """Run the training loop and return the resulting NNRun.

    Args:
        params: dataloaders + optim + scheduler + epochs + seed. The
            train_loader is required; val_loader is optional (skips the
            per-epoch evaluation when absent).
        callbacks: optional list of `Callback` instances (or legacy
            `Callable[[List[IDP]], None]` for back-compat). Each hook
            runs at the documented lifecycle point (on_train_begin,
            on_epoch_begin/end, on_train_end).
        train_step_fn: optional override for the per-batch training
            step. When None (default), runs `default_train_step` —
            supervised forward → loss_fn(net(X), Y) → backward → step.
            Pass a `Callable[[TrainStepContext], NNEvaluationDataPoint]`
            for non-supervised paradigms (autoencoder, VAE, link
            prediction, recommendation, diffusion). The custom function
            is responsible for forward/backward/step and honoring the
            grad-clip/accumulation/AMP knobs the context carries. See
            `docs/concepts.md` and `examples/05_*.py`.

    Returns:
        An `NNRun` with per-iteration `idps`, persisted under
        `runs/<run.id>/` along with per-tag checkpoints. The same
        object is returned with the in-memory idps list attached.

    Raises:
        ValueError: if `params` is None, `params.train_loader` is
            None, or `params.optim` is invalid.
        FloatingPointError: from `default_train_step` if training
            loss becomes non-finite (custom `train_step_fn` hooks are
            responsible for their own divergence checks).
    """
    if params is None:
        raise ValueError("train params must be non-None")
    if params.train_loader is None:
        raise ValueError(
            "params.train_loader is required — set it directly or via with_train_loader(...) before train()."
        )
    if params.optim is None or not params.optim.is_valid():
        raise ValueError(f"train params has an invalid optim config: {params.optim!r}")
    if not any(p.requires_grad for p in self.net.parameters()):
        raise ValueError(
            "model has no trainable parameters — did you freeze('*')? Unfreeze something before train()."
        )

    # V1: seed every RNG before constructing the run so dataset shuffling,
    # weight init, dropout — anything stochastic — is reproducible. The
    # `seed` field only affects state() (and run.id) when explicitly set,
    # so back-compat for no-seed callers is preserved.
    if params.seed is not None:
        from ..seeding import set_seed

        set_seed(params.seed)

    validate: bool = params.val_loader is not None
    # Use self.net_params (always set in __init__) rather than
    # self.net.params: the latter is FeedFwdNN-specific and fails when
    # the caller substitutes a custom nn.Module post-construction
    # (the same idiom Trainer's GAN demo uses and that diffusion
    # demos use for DiffusionMLP). They're identical for the
    # standard supervised path; the rename is a back-compat-safe
    # robustness fix.
    run: NNRun = NNRun(train=params, model=self.params, net=self.net_params)

    optimizer = params.optim.name(
        net=self.net,
        lr_start=params.optim.max_lr,
        momentum=params.optim.momentum,
        weight_decay=params.optim.weight_decay,
        param_groups=params.optim.param_groups,
    )

    # Warm resume: load weights + optimizer state from a prior run's
    # checkpoint. The .opt.pt sidecar is best-effort — pre-resume
    # checkpoints don't have it, in which case we still load weights
    # but the optimizer starts fresh.
    if params.resume_from_run_id is not None:
        ckpt_type = Checkpoints(params.resume_from_checkpoint)
        resume_ckpt = NNCheckpoint.load(run=params.resume_from_run_id, type=ckpt_type)
        if resume_ckpt is None:
            raise ValueError(f"resume_from_run_id={params.resume_from_run_id!r}/{ckpt_type} not found on disk")
        self.net.load_state_dict(resume_ckpt.net_state)
        opt_state = NNCheckpoint.load_optimizer_state(
            run=params.resume_from_run_id,
            type=ckpt_type,
        )
        if opt_state is not None:
            optimizer.load_state_dict(opt_state)

    scheduler = self._build_scheduler(optimizer, params)
    scaler = self._build_grad_scaler()

    normalized_callbacks = self._normalize_callbacks(callbacks)

    idps: list[NNIterationDataPoint] = []
    # `len()` is not defined on iterable-style DataLoaders (IterableDataset).
    # Fall back to None so tqdm renders without a total instead of crashing.
    try:
        n_iter: Optional[int] = int(params.n_epochs * len(params.train_loader))
    except TypeError:
        n_iter = None
    best_checkpoint: Optional[NNCheckpoint] = NNCheckpoint.load(run=run.id, type=Checkpoints.BEST)

    Utils.print_table(header=False, title="Run Details...", data=Utils.flatten_dict(data=run.state()))

    ctx = _CallbackContext(model=self, run=run, optimizer=optimizer)
    for cb in normalized_callbacks:
        cb.on_train_begin(ctx)

    # Default to the standard supervised step when the caller doesn't
    # override. Custom step gets dispatched from inside the batch loop
    # below so the rest of train() (scheduler, callbacks, checkpoint
    # cadence, val loop, incremental save) is identical either way.
    # Explicit None check (not `or`) so a hypothetical callable that
    # happens to be falsy by __bool__ doesn't silently fall back.
    step_fn: TrainStepFn = default_train_step if train_step_fn is None else train_step_fn

    idx_iter = 0
    # Respect NNX_TQDM_DISABLE=1 in tests / CI / non-TTY environments so
    # the progress bar doesn't pollute output. Same env var works as
    # well in subprocess contexts where the user can't pass a flag.
    tqdm_disabled = os.environ.get("NNX_TQDM_DISABLE", "").lower() in {"1", "true", "yes"}
    with (
        torch.set_grad_enabled(True),
        tqdm(colour="blue", total=n_iter, desc="Training", disable=tqdm_disabled) as tqdm_bar,
        _CallbackFinalizer(normalized_callbacks, ctx),
    ):
        for idx_epoch in range(params.n_epochs):
            ctx.epoch = idx_epoch
            for cb in normalized_callbacks:
                cb.on_epoch_begin(ctx)

            n_idps_before_epoch = len(idps)
            for idx_batch, batch in enumerate(params.train_loader):
                step_ctx = TrainStepContext(
                    model=self,
                    batch=batch,
                    optimizer=optimizer,
                    scaler=scaler,
                    grad_clip_norm=params.optim.grad_clip_norm,
                    extra_metrics=params.extra_metrics,
                    accumulate_grad_batches=params.optim.accumulate_grad_batches,
                    batch_idx=idx_batch,
                    epoch_idx=idx_epoch,
                )
                train_edp = step_fn(step_ctx)

                idps.append(
                    NNIterationDataPoint(
                        iter_idx=idx_iter,
                        epoch_idx=idx_epoch,
                        batch_idx=idx_batch,
                        train_edp=train_edp,
                        lr=optimizer.param_groups[0]["lr"],
                    )
                )

                idx_iter += 1
                tqdm_bar.update(1)

            if len(idps) == n_idps_before_epoch:
                # Zero batches this epoch: first epoch would crash on
                # idps[-1] below; later epochs would silently attach
                # this epoch's val_edp to the PREVIOUS epoch's last
                # idp and reuse its stale train_edp.
                raise ValueError(
                    f"train_loader yielded no batches in epoch {idx_epoch} — check batch_size vs "
                    "dataset size with drop_last=True, or whether the loader is a one-shot iterable."
                )

            val_edp = (
                self.evaluate(loader=params.val_loader, extra_metrics=params.extra_metrics) if validate else None
            )
            idps[-1] = idps[-1].with_val_edp(val_edp)

            checkpoint = self._save_checkpoints(
                idp=idps[-1],
                run_id=run.id,
                idx_epoch=idx_epoch,
                n_epochs=params.n_epochs,
                best_checkpoint=best_checkpoint,
                save_phase_checkpoints=params.save_phase_checkpoints,
                optimizer=optimizer,
            )
            # In-memory best_checkpoint tracking must use the same
            # comparison as the on-disk BEST write inside
            # _save_checkpoints (val→train, error→loss, +inf fall-through).
            # Without this, val_loader=None runs would silently overwrite
            # best_checkpoint every epoch (because checkpoint.idp.val_edp
            # is None there) while the on-disk BEST tracks training error,
            # diverging the two views of "best".
            if best_checkpoint is None or _best_err(checkpoint) < _best_err(best_checkpoint):
                best_checkpoint = checkpoint

            self._step_scheduler(scheduler, val_edp, train_edp)
            self._update_tqdm_postfix(tqdm_bar, optimizer, val_edp, train_edp)

            # Incremental persistence: write idps.csv + run.yaml after
            # every epoch. KeyboardInterrupt / OOM during training now
            # leaves a partial-but-loadable run on disk. The extra
            # writes are O(idps so far) per epoch — negligible vs the
            # checkpoint write that already happens.
            run.with_idps(idps).save()

            ctx.idp = idps[-1]
            ctx.idps = idps
            for cb in normalized_callbacks:
                cb.on_epoch_end(ctx)

            if ctx.should_stop:
                break

    saved = run.with_idps(idps).save()
    print()
    runs_root_path = os.path.join(os.getcwd(), "runs", run.id)
    print(f"Run saved to {runs_root_path}")
    return saved

unfreeze(*patterns: str) -> int

Mirror of :meth:freeze — set requires_grad=True on matching parameters.

Source code in nnx/nn/nn_model.py

def unfreeze(self, *patterns: str) -> int:
    """Mirror of :meth:`freeze` — set ``requires_grad=True`` on
    matching parameters."""
    from ..finetune.freezing import unfreeze as _unfreeze

    return _unfreeze(self.net, *patterns)

nnx.nn.nn_model.PredictResult

Bases: NamedTuple

Structured result of NNModel.predict().

Unpacks positionally as (logits, classes) so callers doing log, hat = model.predict(X) keep working after the upgrade from the original 2-tuple. Field access (result.logits, result.classes) is preferred for new code.

Source code in nnx/nn/nn_model.py

class PredictResult(NamedTuple):
    """Structured result of NNModel.predict().

    Unpacks positionally as ``(logits, classes)`` so callers doing
    ``log, hat = model.predict(X)`` keep working after the upgrade from
    the original 2-tuple. Field access (``result.logits``, ``result.classes``)
    is preferred for new code.
    """

    logits: np.ndarray
    classes: np.ndarray

nnx.nn.nn_model.TrainStepContext dataclass

Frozen bundle of state passed into a training-step function.

The default default_train_step runs the standard supervised forward/backward/step. Users can pass their own train_step_fn: Callable[[TrainStepContext], NNEvaluationDataPoint] to NNModel.train() for non-supervised paradigms (autoencoder, VAE, link prediction, recommendation, diffusion, etc.). The custom step is fully responsible for forward, backward, optimizer.step, gradient accumulation, AMP scale/unscale, grad clipping, and the NaN/Inf guard — the context tells it what knobs are set; honoring them is on the caller.

Source code in nnx/nn/nn_model.py

@dataclass(frozen=True, slots=True)
class TrainStepContext:
    """Frozen bundle of state passed into a training-step function.

    The default `default_train_step` runs the standard supervised
    forward/backward/step. Users can pass their own
    `train_step_fn: Callable[[TrainStepContext], NNEvaluationDataPoint]`
    to NNModel.train() for non-supervised paradigms (autoencoder, VAE,
    link prediction, recommendation, diffusion, etc.). The custom step
    is fully responsible for forward, backward, optimizer.step,
    gradient accumulation, AMP scale/unscale, grad clipping, and the
    NaN/Inf guard — the context tells it what knobs are set; honoring
    them is on the caller.
    """

    model: NNModel
    batch: Any
    optimizer: torch.optim.Optimizer
    scaler: Optional[torch.amp.GradScaler]
    grad_clip_norm: Optional[float]
    extra_metrics: Optional[Mapping[str, Callable]]
    accumulate_grad_batches: int
    batch_idx: int
    epoch_idx: int

nnx.nn.nn_model.TrainStepFn = Callable[[TrainStepContext], NNEvaluationDataPoint] module-attribute

nnx.nn.nn_model.default_train_step(ctx: TrainStepContext) -> NNEvaluationDataPoint

Standard supervised training step: forward → loss → backward → step.

This is the body that NNModel.train() runs when no custom train_step_fn is supplied. It honors: - gradient accumulation (zero_grad at cycle start, step at cycle end). Caveat: when an epoch's batch count isn't a multiple of accumulate_grad_batches, the trailing partial cycle's grads are zeroed at the next epoch's first batch (or dropped at the final epoch's end) without an optimizer step — size your loader or accumulation factor accordingly. - AMP (unscales before grad clip; scaler.step + update at cycle end) - grad clipping by L2 norm - the NaN/Inf guard (raises FloatingPointError on divergent loss) - extra_metrics injection on the returned NNEvaluationDataPoint

Custom training-step functions can call this directly to layer on behavior (e.g., extra logging) without reimplementing the standard forward/backward dance.

Source code in nnx/nn/nn_model.py

def default_train_step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
    """Standard supervised training step: forward → loss → backward → step.

    This is the body that `NNModel.train()` runs when no custom
    `train_step_fn` is supplied. It honors:
      - gradient accumulation (zero_grad at cycle start, step at cycle
        end). Caveat: when an epoch's batch count isn't a multiple of
        ``accumulate_grad_batches``, the trailing partial cycle's grads
        are zeroed at the next epoch's first batch (or dropped at the
        final epoch's end) without an optimizer step — size your loader
        or accumulation factor accordingly.
      - AMP (unscales before grad clip; scaler.step + update at cycle end)
      - grad clipping by L2 norm
      - the NaN/Inf guard (raises FloatingPointError on divergent loss)
      - extra_metrics injection on the returned NNEvaluationDataPoint

    Custom training-step functions can call this directly to layer on
    behavior (e.g., extra logging) without reimplementing the standard
    forward/backward dance.
    """
    model = ctx.model
    model.net.train()

    # Gradient accumulation: only zero grads at the start of a fresh
    # accumulation cycle, and only step the optimizer at the end of one.
    accumulate_grad_batches = ctx.accumulate_grad_batches
    is_cycle_start = (ctx.batch_idx % accumulate_grad_batches) == 0
    is_cycle_end = ((ctx.batch_idx + 1) % accumulate_grad_batches) == 0
    if is_cycle_start:
        model.net.zero_grad()

    # Mixed precision is opt-in via NNModelParams.mixed_precision; only
    # takes effect on CUDA where autocast + GradScaler are meaningful.
    scaler = ctx.scaler
    amp_enabled = scaler is not None and model.device.type == "cuda"

    if amp_enabled:
        with torch.amp.autocast(device_type="cuda"):
            X, Y, Y_hat_logits, Y_hat = model._fwd_pass(ctx.batch)
            train_loss = model.loss_fn(Y_hat_logits, Y)
        loss_value = _finite_training_loss_value(train_loss)
        # Scale loss by 1/N so accumulated grads = mean across batches.
        scaler.scale(train_loss / accumulate_grad_batches).backward()
        if is_cycle_end:
            if ctx.grad_clip_norm is not None:
                # Unscale before clipping so the clip threshold applies
                # in the original gradient space, not the scaled one.
                scaler.unscale_(ctx.optimizer)
                torch.nn.utils.clip_grad_norm_(model.net.parameters(), ctx.grad_clip_norm)
            scaler.step(ctx.optimizer)
            scaler.update()
    else:
        X, Y, Y_hat_logits, Y_hat = model._fwd_pass(ctx.batch)
        train_loss = model.loss_fn(Y_hat_logits, Y)
        loss_value = _finite_training_loss_value(train_loss)
        (train_loss / accumulate_grad_batches).backward()
        if is_cycle_end:
            if ctx.grad_clip_norm is not None:
                torch.nn.utils.clip_grad_norm_(model.net.parameters(), ctx.grad_clip_norm)
            ctx.optimizer.step()

    return classification_edp(
        Y=Y,
        Y_hat=Y_hat,
        loss=loss_value,
        extra_metrics=ctx.extra_metrics,
    )

2.2. GenerativeNNModel — decoder-only LM orchestrator

nnx.nn.generative_nn_model.GenerativeNNModel

Bases: NNModel

Language model with an autoregressive generate() method.

tokenizer is held as a regular instance attribute (not a constructor-arg of NNModel) so existing NNModel callers don't have to know about it. It's required for generate() but optional at construction — train-time you can build the model first and attach the tokenizer later.

Source code in nnx/nn/generative_nn_model.py

class GenerativeNNModel(NNModel):
    """Language model with an autoregressive ``generate()`` method.

    ``tokenizer`` is held as a regular instance attribute (not a
    constructor-arg of NNModel) so existing NNModel callers don't
    have to know about it. It's required for ``generate()`` but
    optional at construction — train-time you can build the model
    first and attach the tokenizer later.
    """

    def __init__(
        self,
        net_params: NNParams,
        params: NNModelParams,
        tokenizer: Optional[NNTokenizerParams] = None,
    ):
        super().__init__(net_params=net_params, params=params)
        self.tokenizer = tokenizer

    # ---------- generate ----------

    def generate(
        self,
        prompt: str,
        *,
        max_new_tokens: int = 64,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        repetition_penalty: float = 1.0,
        stop: Optional[list[str]] = None,
        seed: Optional[int] = None,
        use_cache: bool = True,
        logits_chain: Optional[LogitsChain] = None,
    ) -> str:
        """Autoregressive decode from ``prompt``.

        Args:
            prompt: input text. Encoded via ``self.tokenizer``.
            max_new_tokens: hard cap on new tokens emitted. Generation
                also stops if the context window (max_seq_len) would be
                exceeded and the model can't shrink the window further,
                or if a ``stop`` string is decoded.
            temperature: 0 means greedy (argmax). Higher values produce
                more diverse output. Routes through TemperatureScaling.
            top_k: keep only the top-k logits. None disables.
            top_p: nucleus (top-p) cutoff. None disables.
            repetition_penalty: divide previously-seen tokens' positive
                logits by this. 1.0 is no-op (default).
            stop: list of stop strings — generation halts once any of
                them appears in the decoded CONTINUATION (the prompt
                itself is not searched, so a prompt containing a stop
                string doesn't halt generation immediately; a stop
                string straddling the prompt/continuation boundary is
                likewise not detected — matching the generated-text-only
                convention HF uses).
            seed: when set, sampling is reproducible — two calls with
                the same seed + prompt + model produce identical output.
            use_cache: when True (default), uses an incremental KV
                cache — each new token only re-runs attention on the
                last position, not the whole prefix. When False, falls
                back to the full-recompute path (kept for regression
                testing). Both paths produce the same tokens for greedy
                decoding (sampling paths agree given the same seed).
            logits_chain: optional pre-built ``LogitsChain`` (see
                ``nnx.LogitsChain.builder()``). When provided, the
                inline chain construction from ``temperature`` /
                ``top_k`` / ``top_p`` / ``repetition_penalty`` kwargs
                is skipped — the supplied chain is used as-is.
                Power-user path for custom logit processors (e.g.,
                logit-bias for forbidden tokens). When ``None`` (the
                default), behavior is unchanged.

        Returns:
            The full decoded string (prompt + generated continuation).

        Non-destructive: ``self.net.training`` is snapshotted before
        switching to ``eval()`` and restored on exit (including the
        exception path via ``try/finally``). Matches the convention
        used by ``NNModel.predict`` / ``NNModel.evaluate``,
        ``nnx.diffusion.sample``, ``nnx.embeddings.embed_texts``,
        ``nnx.viz.activation_map``, and ``nnx.lr_finder``.
        """
        if self.tokenizer is None:
            raise ValueError(
                "GenerativeNNModel.generate requires a tokenizer. "
                "Construct with `GenerativeNNModel(..., tokenizer=NNTokenizerParams.of(tk, path))`."
            )
        prompt_ids = self.tokenizer.encode(prompt)
        if not prompt_ids:
            # Use a single space as a non-empty seed when the prompt
            # encodes to nothing (e.g., an empty string with no BOS).
            prompt_ids = self.tokenizer.encode(" ") or [0]

        max_seq_len = getattr(self.net_params, "max_seq_len", None)
        if max_seq_len is None:
            raise ValueError("net_params must expose `max_seq_len` (got a non-Transformer net?)")
        # Wrappers that consume window slots (PromptTuner's soft prompt)
        # advertise a smaller effective window — without this, the
        # sliding window overflows the wrapped model mid-generation.
        max_seq_len = getattr(self.net, "effective_max_seq_len", max_seq_len)

        # Build the processor chain. Two paths:
        # (a) If the caller supplied a pre-built ``logits_chain``, use
        #     its processors as-is — they've already been ordered.
        # (b) Otherwise build the standard chain from kwargs in NNx's
        #     canonical order: repetition penalty first (raw logits),
        #     then top-k / top-p (filtering), then temperature
        #     (scaling — deliberately last: temperature=0's ±inf greedy
        #     markers must not be re-filtered). Temperature is always applied;
        #     temperature=0 is the greedy short-circuit (still routed
        #     through TemperatureScaling so the sampler path is
        #     uniform).
        processors: list[LogitsProcessor]
        if logits_chain is not None:
            processors = list(logits_chain.processors)
        else:
            processors = []
            if repetition_penalty != 1.0:
                processors.append(RepetitionPenalty(penalty=repetition_penalty))
            if top_k is not None:
                processors.append(TopKFilter(top_k=top_k))
            if top_p is not None:
                processors.append(TopPFilter(top_p=top_p))
            processors.append(TemperatureScaling(temperature=temperature))

        # Optional seeded torch.Generator for reproducible sampling.
        # We pin it to the model's device so multinomial doesn't fall
        # back to CPU silently on a CUDA model.
        gen = None
        if seed is not None:
            gen = torch.Generator(device=self.device)
            gen.manual_seed(int(seed))

        # The KV-cache path needs ``forward_with_cache`` on the net.
        # Plain LSTM/MLP nets (or older Transformer forks) don't have
        # it — fall back transparently rather than crashing.
        if use_cache and not hasattr(self.net, "forward_with_cache"):
            use_cache = False

        generated: list[int] = list(prompt_ids)
        # Snapshot training-mode for non-destructive restore on exit
        # (matches NNModel.predict / evaluate / nnx.viz.activation_map /
        # nnx.lr_finder). Deliberately the LAST thing before the
        # try/finally: an exception in the validation / chain-building
        # code above would otherwise strand the net in eval() with no
        # finally to restore it.
        was_training = self.net.training
        self.net.eval()
        try:
            with torch.no_grad():
                if use_cache:
                    self._generate_with_cache(
                        generated=generated,
                        n_prompt=len(prompt_ids),
                        max_new_tokens=max_new_tokens,
                        max_seq_len=max_seq_len,
                        processors=processors,
                        gen=gen,
                        stop=stop,
                    )
                else:
                    self._generate_no_cache(
                        generated=generated,
                        n_prompt=len(prompt_ids),
                        max_new_tokens=max_new_tokens,
                        max_seq_len=max_seq_len,
                        processors=processors,
                        gen=gen,
                        stop=stop,
                    )
        finally:
            if was_training:
                self.net.train()

        return self.tokenizer.decode(generated)

    # ---------- generate helpers ----------

    def _generate_no_cache(
        self,
        *,
        generated: list[int],
        n_prompt: int,
        max_new_tokens: int,
        max_seq_len: int,
        processors: list[LogitsProcessor],
        gen: Optional[torch.Generator],
        stop: Optional[list[str]],
    ) -> None:
        """Full-recompute path: every step re-runs the model on the
        last ``max_seq_len`` tokens. O(T^2) attention cost; kept for
        regression-testing parity against the cached path."""
        for _ in range(max_new_tokens):
            # Truncate context to max_seq_len from the right so the
            # most recent tokens stay in the window. Sliding window —
            # the simple production-ready approach.
            context_ids = generated[-max_seq_len:]
            ctx = torch.tensor([context_ids], dtype=torch.long, device=self.device)
            logits = self.net(ctx)  # (1, T, vocab)
            next_logits = logits[:, -1, :]  # (1, vocab) — last token's logits
            adjusted = apply_chain(next_logits, token_history=generated, processors=processors)
            next_id = sample_next_token(adjusted, generator=gen)
            generated.append(next_id)

            # Optional stop-string check, scoped to the CONTINUATION —
            # a stop string already present in the prompt must not halt
            # generation after one token. Only decodes when stops are
            # configured, keeping the hot loop cheap otherwise.
            if stop:
                continuation = self.tokenizer.decode(generated[n_prompt:])
                if any(s in continuation for s in stop):
                    break

    def _generate_with_cache(
        self,
        *,
        generated: list[int],
        n_prompt: int,
        max_new_tokens: int,
        max_seq_len: int,
        processors: list[LogitsProcessor],
        gen: Optional[torch.Generator],
        stop: Optional[list[str]],
    ) -> None:
        """KV-cache path. Runs one prefill pass on the
        truncated prompt, then per new token re-runs only the last
        position's attention against the cached prefix. O(T) per
        step instead of O(T^2).

        Sliding-window safety: once appending another token would
        overflow ``max_seq_len``, the cache is rebuilt from the current
        window. Cached k/v are RoPE-stamped at the absolute position
        they were written at, so merely dropping the oldest entry would
        pin every later token's offset at ``max_seq_len - 1`` and
        corrupt the relative position geometry (logits drift vs the
        no-cache path). Rebuilding re-rotates the window to positions
        ``0..max_seq_len-1`` — exactly what the no-cache path computes,
        so greedy/seeded parity holds across overflow. Post-overflow
        steps therefore cost one full window forward, same as the
        no-cache path; the O(T) win applies within the window.
        """
        if max_new_tokens <= 0:
            # Hard cap honored on this path too: the prefill below
            # always samples one token, which would emit 1 instead of 0.
            return

        # ----- Prefill pass on the prompt (sliding window). -----
        context_ids = generated[-max_seq_len:]
        ctx = torch.tensor([context_ids], dtype=torch.long, device=self.device)
        logits, past_kvs = self.net.forward_with_cache(ctx, past_kvs=None)
        next_logits = logits[:, -1, :]
        adjusted = apply_chain(next_logits, token_history=generated, processors=processors)
        next_id = sample_next_token(adjusted, generator=gen)
        generated.append(next_id)
        if stop:
            # Continuation-scoped, like the loop check below.
            continuation = self.tokenizer.decode(generated[n_prompt:])
            if any(s in continuation for s in stop):
                return

        # ----- Incremental decode loop. -----
        for _ in range(max_new_tokens - 1):
            cached_len = past_kvs[0][0].size(-2) if past_kvs and past_kvs[0] is not None else 0
            if cached_len + 1 > max_seq_len:
                # Window full: rebuild the cache from the current
                # window (see the docstring — dropping the oldest k/v
                # would corrupt RoPE relative positions). The window
                # includes generated[-1], so this forward both refills
                # the cache and yields the next token's logits.
                ctx = torch.tensor([generated[-max_seq_len:]], dtype=torch.long, device=self.device)
                logits, past_kvs = self.net.forward_with_cache(ctx, past_kvs=None)
            else:
                ctx = torch.tensor([[generated[-1]]], dtype=torch.long, device=self.device)
                logits, past_kvs = self.net.forward_with_cache(ctx, past_kvs=past_kvs)
            next_logits = logits[:, -1, :]
            adjusted = apply_chain(next_logits, token_history=generated, processors=processors)
            next_id = sample_next_token(adjusted, generator=gen)
            generated.append(next_id)

            if stop:
                # Continuation-scoped — see _generate_no_cache.
                continuation = self.tokenizer.decode(generated[n_prompt:])
                if any(s in continuation for s in stop):
                    break

generate(prompt: str, *, max_new_tokens: int = 64, temperature: float = 1.0, top_k: Optional[int] = None, top_p: Optional[float] = None, repetition_penalty: float = 1.0, stop: Optional[list[str]] = None, seed: Optional[int] = None, use_cache: bool = True, logits_chain: Optional[LogitsChain] = None) -> str

Autoregressive decode from prompt.

Parameters:

Name	Type	Description	Default
prompt	str	input text. Encoded via self.tokenizer.	required
max_new_tokens	int	hard cap on new tokens emitted. Generation also stops if the context window (max_seq_len) would be exceeded and the model can't shrink the window further, or if a stop string is decoded.	64
temperature	float	0 means greedy (argmax). Higher values produce more diverse output. Routes through TemperatureScaling.	1.0
top_k	Optional[int]	keep only the top-k logits. None disables.	None
top_p	Optional[float]	nucleus (top-p) cutoff. None disables.	None
repetition_penalty	float	divide previously-seen tokens' positive logits by this. 1.0 is no-op (default).	1.0
stop	Optional[list[str]]	list of stop strings — generation halts once any of them appears in the decoded CONTINUATION (the prompt itself is not searched, so a prompt containing a stop string doesn't halt generation immediately; a stop string straddling the prompt/continuation boundary is likewise not detected — matching the generated-text-only convention HF uses).	None
seed	Optional[int]	when set, sampling is reproducible — two calls with the same seed + prompt + model produce identical output.	None
use_cache	bool	when True (default), uses an incremental KV cache — each new token only re-runs attention on the last position, not the whole prefix. When False, falls back to the full-recompute path (kept for regression testing). Both paths produce the same tokens for greedy decoding (sampling paths agree given the same seed).	True
logits_chain	Optional[LogitsChain]	optional pre-built LogitsChain (see nnx.LogitsChain.builder()). When provided, the inline chain construction from temperature / top_k / top_p / repetition_penalty kwargs is skipped — the supplied chain is used as-is. Power-user path for custom logit processors (e.g., logit-bias for forbidden tokens). When None (the default), behavior is unchanged.	None

Returns:

Type	Description
str	The full decoded string (prompt + generated continuation).

Non-destructive: self.net.training is snapshotted before switching to eval() and restored on exit (including the exception path via try/finally). Matches the convention used by NNModel.predict / NNModel.evaluate, nnx.diffusion.sample, nnx.embeddings.embed_texts, nnx.viz.activation_map, and nnx.lr_finder.

Source code in nnx/nn/generative_nn_model.py

def generate(
    self,
    prompt: str,
    *,
    max_new_tokens: int = 64,
    temperature: float = 1.0,
    top_k: Optional[int] = None,
    top_p: Optional[float] = None,
    repetition_penalty: float = 1.0,
    stop: Optional[list[str]] = None,
    seed: Optional[int] = None,
    use_cache: bool = True,
    logits_chain: Optional[LogitsChain] = None,
) -> str:
    """Autoregressive decode from ``prompt``.

    Args:
        prompt: input text. Encoded via ``self.tokenizer``.
        max_new_tokens: hard cap on new tokens emitted. Generation
            also stops if the context window (max_seq_len) would be
            exceeded and the model can't shrink the window further,
            or if a ``stop`` string is decoded.
        temperature: 0 means greedy (argmax). Higher values produce
            more diverse output. Routes through TemperatureScaling.
        top_k: keep only the top-k logits. None disables.
        top_p: nucleus (top-p) cutoff. None disables.
        repetition_penalty: divide previously-seen tokens' positive
            logits by this. 1.0 is no-op (default).
        stop: list of stop strings — generation halts once any of
            them appears in the decoded CONTINUATION (the prompt
            itself is not searched, so a prompt containing a stop
            string doesn't halt generation immediately; a stop
            string straddling the prompt/continuation boundary is
            likewise not detected — matching the generated-text-only
            convention HF uses).
        seed: when set, sampling is reproducible — two calls with
            the same seed + prompt + model produce identical output.
        use_cache: when True (default), uses an incremental KV
            cache — each new token only re-runs attention on the
            last position, not the whole prefix. When False, falls
            back to the full-recompute path (kept for regression
            testing). Both paths produce the same tokens for greedy
            decoding (sampling paths agree given the same seed).
        logits_chain: optional pre-built ``LogitsChain`` (see
            ``nnx.LogitsChain.builder()``). When provided, the
            inline chain construction from ``temperature`` /
            ``top_k`` / ``top_p`` / ``repetition_penalty`` kwargs
            is skipped — the supplied chain is used as-is.
            Power-user path for custom logit processors (e.g.,
            logit-bias for forbidden tokens). When ``None`` (the
            default), behavior is unchanged.

    Returns:
        The full decoded string (prompt + generated continuation).

    Non-destructive: ``self.net.training`` is snapshotted before
    switching to ``eval()`` and restored on exit (including the
    exception path via ``try/finally``). Matches the convention
    used by ``NNModel.predict`` / ``NNModel.evaluate``,
    ``nnx.diffusion.sample``, ``nnx.embeddings.embed_texts``,
    ``nnx.viz.activation_map``, and ``nnx.lr_finder``.
    """
    if self.tokenizer is None:
        raise ValueError(
            "GenerativeNNModel.generate requires a tokenizer. "
            "Construct with `GenerativeNNModel(..., tokenizer=NNTokenizerParams.of(tk, path))`."
        )
    prompt_ids = self.tokenizer.encode(prompt)
    if not prompt_ids:
        # Use a single space as a non-empty seed when the prompt
        # encodes to nothing (e.g., an empty string with no BOS).
        prompt_ids = self.tokenizer.encode(" ") or [0]

    max_seq_len = getattr(self.net_params, "max_seq_len", None)
    if max_seq_len is None:
        raise ValueError("net_params must expose `max_seq_len` (got a non-Transformer net?)")
    # Wrappers that consume window slots (PromptTuner's soft prompt)
    # advertise a smaller effective window — without this, the
    # sliding window overflows the wrapped model mid-generation.
    max_seq_len = getattr(self.net, "effective_max_seq_len", max_seq_len)

    # Build the processor chain. Two paths:
    # (a) If the caller supplied a pre-built ``logits_chain``, use
    #     its processors as-is — they've already been ordered.
    # (b) Otherwise build the standard chain from kwargs in NNx's
    #     canonical order: repetition penalty first (raw logits),
    #     then top-k / top-p (filtering), then temperature
    #     (scaling — deliberately last: temperature=0's ±inf greedy
    #     markers must not be re-filtered). Temperature is always applied;
    #     temperature=0 is the greedy short-circuit (still routed
    #     through TemperatureScaling so the sampler path is
    #     uniform).
    processors: list[LogitsProcessor]
    if logits_chain is not None:
        processors = list(logits_chain.processors)
    else:
        processors = []
        if repetition_penalty != 1.0:
            processors.append(RepetitionPenalty(penalty=repetition_penalty))
        if top_k is not None:
            processors.append(TopKFilter(top_k=top_k))
        if top_p is not None:
            processors.append(TopPFilter(top_p=top_p))
        processors.append(TemperatureScaling(temperature=temperature))

    # Optional seeded torch.Generator for reproducible sampling.
    # We pin it to the model's device so multinomial doesn't fall
    # back to CPU silently on a CUDA model.
    gen = None
    if seed is not None:
        gen = torch.Generator(device=self.device)
        gen.manual_seed(int(seed))

    # The KV-cache path needs ``forward_with_cache`` on the net.
    # Plain LSTM/MLP nets (or older Transformer forks) don't have
    # it — fall back transparently rather than crashing.
    if use_cache and not hasattr(self.net, "forward_with_cache"):
        use_cache = False

    generated: list[int] = list(prompt_ids)
    # Snapshot training-mode for non-destructive restore on exit
    # (matches NNModel.predict / evaluate / nnx.viz.activation_map /
    # nnx.lr_finder). Deliberately the LAST thing before the
    # try/finally: an exception in the validation / chain-building
    # code above would otherwise strand the net in eval() with no
    # finally to restore it.
    was_training = self.net.training
    self.net.eval()
    try:
        with torch.no_grad():
            if use_cache:
                self._generate_with_cache(
                    generated=generated,
                    n_prompt=len(prompt_ids),
                    max_new_tokens=max_new_tokens,
                    max_seq_len=max_seq_len,
                    processors=processors,
                    gen=gen,
                    stop=stop,
                )
            else:
                self._generate_no_cache(
                    generated=generated,
                    n_prompt=len(prompt_ids),
                    max_new_tokens=max_new_tokens,
                    max_seq_len=max_seq_len,
                    processors=processors,
                    gen=gen,
                    stop=stop,
                )
    finally:
        if was_training:
            self.net.train()

    return self.tokenizer.decode(generated)

2.3. Trainer — multi-optimizer orchestrator

nnx.trainer.trainer.Trainer

Multi-optimizer training orchestrator.

Constructed around a single NNModel. At train() time, builds one torch.optim.Optimizer per entry in NNTrainerParams.optims (each scoped to its sub-net via NNOptimParams.param_groups) and invokes the user-supplied trainer_step_fn for each batch.

Same NNRun + per-tag NNCheckpoint cadence as NNModel.train(), with the extra trainer block on NNRun preserving the multi-optim configuration on disk.

Source code in nnx/trainer/trainer.py

class Trainer:
    """Multi-optimizer training orchestrator.

    Constructed around a single NNModel. At train() time, builds one
    torch.optim.Optimizer per entry in NNTrainerParams.optims (each
    scoped to its sub-net via NNOptimParams.param_groups) and invokes
    the user-supplied trainer_step_fn for each batch.

    Same NNRun + per-tag NNCheckpoint cadence as NNModel.train(),
    with the extra `trainer` block on NNRun preserving the multi-optim
    configuration on disk.
    """

    def __init__(self, model: NNModel):
        if model is None:
            raise ValueError("Trainer requires a non-None model")
        self.model = model

    def train(
        self,
        params: NNTrainerParams,
        trainer_step_fn: TrainerStepFn,
        callbacks: Optional[list[CallbackLike]] = None,
    ) -> NNRun:
        """Run the multi-optimizer training loop and return the resulting NNRun.

        Args:
            params: NNTrainerParams — train_loader + n_epochs + optims dict +
                (optional) schedulers dict + (optional) val_loader, seed,
                save_phase_checkpoints, extra_metrics.
            trainer_step_fn: required. `Callable[[TrainerStepContext],
                NNEvaluationDataPoint]`. The function owns the entire per-batch
                update — including which optimizers to step, in what order, and
                with what loss(es). There is no supervised fallback.
            callbacks: optional list of Callback instances. The callback
                context exposes `ctx.optimizer` (primary, sorted-first), plus
                a `ctx.optimizers` dict and `ctx.trainer` reference for
                trainer-aware callbacks.

        Returns:
            NNRun with per-iteration idps, persisted under runs/<run.id>/
            alongside the standard FIRST/Q1/Q2/Q3/LAST/BEST checkpoints.

        Raises:
            ValueError: when params is None, params.train_loader is None,
                trainer_step_fn is None, or any optim's
                NNOptimParams.is_valid() returns False.
        """
        if params is None:
            raise ValueError("trainer params must not be None")
        if params.train_loader is None:
            raise ValueError(
                "params.train_loader is required — set it directly or via with_train_loader(...) before train()."
            )
        if trainer_step_fn is None:
            raise ValueError(
                "trainer_step_fn is required — Trainer has no default "
                "supervised step because multi-optim updates are inherently "
                "scenario-specific."
            )
        for name, opt_params in params.optims.items():
            if not opt_params.is_valid():
                raise ValueError(f"optim {name!r} has invalid config: {opt_params}")
        if not any(p.requires_grad for p in self.model.net.parameters()):
            raise ValueError(
                "model has no trainable parameters — did you freeze('*')? Unfreeze something before train()."
            )

        if params.seed is not None:
            from ..seeding import set_seed

            set_seed(params.seed)

        validate = params.val_loader is not None

        run = NNRun(
            train=_representative_train_params(params),
            trainer=params,
            model=self.model.params,
            # Use the model's stored NNParams rather than self.model.net.params
            # so callers who substitute a custom nn.Module post-construction
            # (the GAN composite idiom) still produce a saveable run.
            net=self.model.net_params,
        )

        # `strict_param_groups=True` is the multi-optim contract: each
        # optimizer owns only the parameters its specs explicitly match,
        # not also the default-bucket leftovers. Without this, opt_G
        # would also hold D's params (unmatched by G's specs) and the
        # two optimizers would silently fight over the same gradients.
        #
        # That contract only bites when each optimizer HAS specs: an optimizer
        # with `param_groups=None` routes to all `net.parameters()`, so two such
        # optimizers would each step the full set — silently double-stepping
        # every parameter. Require explicit scoping up front (the params object
        # itself stays constructible for serialization / builder round-trips).
        unscoped = sorted(n for n, op in params.optims.items() if op.param_groups is None)
        if len(params.optims) >= 2 and unscoped:
            raise ValueError(
                "Trainer with multiple optimizers requires every optimizer to scope its "
                "parameters via `param_groups` (NNParamGroupSpec); these have none and would "
                f"each grab all net parameters, double-stepping them: {unscoped}."
            )
        optimizers = {
            name: opt_params.name(
                net=self.model.net,
                lr_start=opt_params.max_lr,
                momentum=opt_params.momentum,
                weight_decay=opt_params.weight_decay,
                param_groups=opt_params.param_groups,
                strict_param_groups=True,
            )
            for name, opt_params in params.optims.items()
        }
        schedulers = {
            name: _build_scheduler(
                opt=optimizers[name],
                sched_params=params.schedulers.get(name, _DEFAULT_SCHEDULER_PARAMS),
                n_epochs=params.n_epochs,
            )
            for name in optimizers
        }

        normalized_callbacks = NNModel._normalize_callbacks(callbacks)

        primary = _primary_name(optimizers.keys())
        ctx = _CallbackContext(
            model=self.model,
            run=run,
            optimizer=optimizers[primary],
        )
        # Trainer-aware extensions: existing callbacks reading ctx.optimizer
        # see the primary; new callbacks can downcast through ctx.optimizers
        # or ctx.trainer for the multi-optim view.
        ctx.optimizers = optimizers
        ctx.trainer = self

        idps: list[NNIterationDataPoint] = []
        # `len()` is not defined on iterable-style DataLoaders (IterableDataset).
        # Fall back to None so tqdm renders without a total instead of crashing.
        try:
            n_iter: Optional[int] = int(params.n_epochs * len(params.train_loader))
        except TypeError:
            n_iter = None
        best_checkpoint: Optional[NNCheckpoint] = NNCheckpoint.load(
            run=run.id,
            type=Checkpoints.BEST,
        )

        Utils.print_table(
            header=False,
            title="Trainer Run Details...",
            data=Utils.flatten_dict(data=run.state()),
        )

        for cb in normalized_callbacks:
            cb.on_train_begin(ctx)

        idx_iter = 0
        tqdm_disabled = os.environ.get("NNX_TQDM_DISABLE", "").lower() in {"1", "true", "yes"}
        with (
            torch.set_grad_enabled(True),
            tqdm(colour="blue", total=n_iter, desc="Training", disable=tqdm_disabled) as tqdm_bar,
            _CallbackFinalizer(normalized_callbacks, ctx),
        ):
            for idx_epoch in range(params.n_epochs):
                ctx.epoch = idx_epoch
                for cb in normalized_callbacks:
                    cb.on_epoch_begin(ctx)

                n_idps_before_epoch = len(idps)
                for idx_batch, batch in enumerate(params.train_loader):
                    step_ctx = TrainerStepContext(
                        model=self.model,
                        batch=batch,
                        optimizers=optimizers,
                        schedulers=schedulers,
                        extra_metrics=params.extra_metrics,
                        batch_idx=idx_batch,
                        epoch_idx=idx_epoch,
                    )
                    train_edp = trainer_step_fn(step_ctx)

                    idps.append(
                        NNIterationDataPoint(
                            iter_idx=idx_iter,
                            epoch_idx=idx_epoch,
                            batch_idx=idx_batch,
                            train_edp=train_edp,
                            lr=optimizers[primary].param_groups[0]["lr"],
                        )
                    )
                    idx_iter += 1
                    tqdm_bar.update(1)

                if len(idps) == n_idps_before_epoch:
                    # Same guard as NNModel.train: zero batches would
                    # crash on idps[-1] (first epoch) or corrupt the
                    # previous epoch's logged metrics (later epochs).
                    raise ValueError(
                        f"train_loader yielded no batches in epoch {idx_epoch} — check batch_size vs "
                        "dataset size with drop_last=True, or whether the loader is a one-shot iterable."
                    )

                val_edp = (
                    self.model.evaluate(
                        loader=params.val_loader,
                        extra_metrics=params.extra_metrics,
                    )
                    if validate
                    else None
                )
                idps[-1] = idps[-1].with_val_edp(val_edp)

                checkpoint = self._save_checkpoint(
                    idp=idps[-1],
                    run_id=run.id,
                    idx_epoch=idx_epoch,
                    n_epochs=params.n_epochs,
                    best_checkpoint=best_checkpoint,
                    save_phase_checkpoints=params.save_phase_checkpoints,
                )
                if best_checkpoint is None or _best_err(checkpoint) < _best_err(best_checkpoint):
                    best_checkpoint = checkpoint

                # Each scheduler steps on its own optimizer's signal.
                # We feed the SAME (val_edp, train_edp) pair to all of
                # them because IDPs aggregate over all optims — separating
                # per-optim metrics would require the step fn to return
                # multiple EDPs, which complicates the contract without
                # clear benefit. Users who need per-optim scheduler
                # signals can step their scheduler directly in the step fn.
                for sched in schedulers.values():
                    _step_scheduler(sched, val_edp, train_edp)

                # Verbatim-shared with the NNModel.train loop — one
                # implementation, so the postfix format can't drift.
                self.model._update_tqdm_postfix(tqdm_bar, optimizers[primary], val_edp, train_edp)

                run.with_idps(idps).save()

                ctx.idp = idps[-1]
                ctx.idps = idps
                for cb in normalized_callbacks:
                    cb.on_epoch_end(ctx)

                if ctx.should_stop:
                    break

        saved = run.with_idps(idps).save()
        print()
        runs_root_path = os.path.join(os.getcwd(), "runs", run.id)
        print(f"Run saved to {runs_root_path}")
        return saved

    def _save_checkpoint(
        self,
        idp: NNIterationDataPoint,
        run_id: str,
        idx_epoch: int,
        n_epochs: int,
        best_checkpoint: Optional[NNCheckpoint],
        save_phase_checkpoints: bool,
    ) -> NNCheckpoint:
        """Delegates to NNModel._save_checkpoints with optimizer=None —
        same FIRST/Q1/Q2/Q3/LAST/BEST cadence, minus the optimizer-state
        sidecar (a single sidecar can't represent the multi-optim dict).
        Trainer-mode warm-resume is a follow-up — saving one sidecar per
        optim (`<tag>.opt.<name>.pt`) is the natural extension when that
        lands."""
        return self.model._save_checkpoints(
            idp=idp,
            run_id=run_id,
            idx_epoch=idx_epoch,
            n_epochs=n_epochs,
            best_checkpoint=best_checkpoint,
            save_phase_checkpoints=save_phase_checkpoints,
            optimizer=None,
        )

train(params: NNTrainerParams, trainer_step_fn: TrainerStepFn, callbacks: Optional[list[CallbackLike]] = None) -> NNRun

Run the multi-optimizer training loop and return the resulting NNRun.

Parameters:

Name	Type	Description	Default
params	NNTrainerParams	NNTrainerParams — train_loader + n_epochs + optims dict + (optional) schedulers dict + (optional) val_loader, seed, save_phase_checkpoints, extra_metrics.	required
trainer_step_fn	TrainerStepFn	required. Callable[[TrainerStepContext], NNEvaluationDataPoint]. The function owns the entire per-batch update — including which optimizers to step, in what order, and with what loss(es). There is no supervised fallback.	required
callbacks	Optional[list[CallbackLike]]	optional list of Callback instances. The callback context exposes ctx.optimizer (primary, sorted-first), plus a ctx.optimizers dict and ctx.trainer reference for trainer-aware callbacks.	None

Returns:

Type	Description
NNRun	NNRun with per-iteration idps, persisted under runs//
NNRun	alongside the standard FIRST/Q1/Q2/Q3/LAST/BEST checkpoints.

Raises:

Type	Description
ValueError	when params is None, params.train_loader is None, trainer_step_fn is None, or any optim's NNOptimParams.is_valid() returns False.

Source code in nnx/trainer/trainer.py

def train(
    self,
    params: NNTrainerParams,
    trainer_step_fn: TrainerStepFn,
    callbacks: Optional[list[CallbackLike]] = None,
) -> NNRun:
    """Run the multi-optimizer training loop and return the resulting NNRun.

    Args:
        params: NNTrainerParams — train_loader + n_epochs + optims dict +
            (optional) schedulers dict + (optional) val_loader, seed,
            save_phase_checkpoints, extra_metrics.
        trainer_step_fn: required. `Callable[[TrainerStepContext],
            NNEvaluationDataPoint]`. The function owns the entire per-batch
            update — including which optimizers to step, in what order, and
            with what loss(es). There is no supervised fallback.
        callbacks: optional list of Callback instances. The callback
            context exposes `ctx.optimizer` (primary, sorted-first), plus
            a `ctx.optimizers` dict and `ctx.trainer` reference for
            trainer-aware callbacks.

    Returns:
        NNRun with per-iteration idps, persisted under runs/<run.id>/
        alongside the standard FIRST/Q1/Q2/Q3/LAST/BEST checkpoints.

    Raises:
        ValueError: when params is None, params.train_loader is None,
            trainer_step_fn is None, or any optim's
            NNOptimParams.is_valid() returns False.
    """
    if params is None:
        raise ValueError("trainer params must not be None")
    if params.train_loader is None:
        raise ValueError(
            "params.train_loader is required — set it directly or via with_train_loader(...) before train()."
        )
    if trainer_step_fn is None:
        raise ValueError(
            "trainer_step_fn is required — Trainer has no default "
            "supervised step because multi-optim updates are inherently "
            "scenario-specific."
        )
    for name, opt_params in params.optims.items():
        if not opt_params.is_valid():
            raise ValueError(f"optim {name!r} has invalid config: {opt_params}")
    if not any(p.requires_grad for p in self.model.net.parameters()):
        raise ValueError(
            "model has no trainable parameters — did you freeze('*')? Unfreeze something before train()."
        )

    if params.seed is not None:
        from ..seeding import set_seed

        set_seed(params.seed)

    validate = params.val_loader is not None

    run = NNRun(
        train=_representative_train_params(params),
        trainer=params,
        model=self.model.params,
        # Use the model's stored NNParams rather than self.model.net.params
        # so callers who substitute a custom nn.Module post-construction
        # (the GAN composite idiom) still produce a saveable run.
        net=self.model.net_params,
    )

    # `strict_param_groups=True` is the multi-optim contract: each
    # optimizer owns only the parameters its specs explicitly match,
    # not also the default-bucket leftovers. Without this, opt_G
    # would also hold D's params (unmatched by G's specs) and the
    # two optimizers would silently fight over the same gradients.
    #
    # That contract only bites when each optimizer HAS specs: an optimizer
    # with `param_groups=None` routes to all `net.parameters()`, so two such
    # optimizers would each step the full set — silently double-stepping
    # every parameter. Require explicit scoping up front (the params object
    # itself stays constructible for serialization / builder round-trips).
    unscoped = sorted(n for n, op in params.optims.items() if op.param_groups is None)
    if len(params.optims) >= 2 and unscoped:
        raise ValueError(
            "Trainer with multiple optimizers requires every optimizer to scope its "
            "parameters via `param_groups` (NNParamGroupSpec); these have none and would "
            f"each grab all net parameters, double-stepping them: {unscoped}."
        )
    optimizers = {
        name: opt_params.name(
            net=self.model.net,
            lr_start=opt_params.max_lr,
            momentum=opt_params.momentum,
            weight_decay=opt_params.weight_decay,
            param_groups=opt_params.param_groups,
            strict_param_groups=True,
        )
        for name, opt_params in params.optims.items()
    }
    schedulers = {
        name: _build_scheduler(
            opt=optimizers[name],
            sched_params=params.schedulers.get(name, _DEFAULT_SCHEDULER_PARAMS),
            n_epochs=params.n_epochs,
        )
        for name in optimizers
    }

    normalized_callbacks = NNModel._normalize_callbacks(callbacks)

    primary = _primary_name(optimizers.keys())
    ctx = _CallbackContext(
        model=self.model,
        run=run,
        optimizer=optimizers[primary],
    )
    # Trainer-aware extensions: existing callbacks reading ctx.optimizer
    # see the primary; new callbacks can downcast through ctx.optimizers
    # or ctx.trainer for the multi-optim view.
    ctx.optimizers = optimizers
    ctx.trainer = self

    idps: list[NNIterationDataPoint] = []
    # `len()` is not defined on iterable-style DataLoaders (IterableDataset).
    # Fall back to None so tqdm renders without a total instead of crashing.
    try:
        n_iter: Optional[int] = int(params.n_epochs * len(params.train_loader))
    except TypeError:
        n_iter = None
    best_checkpoint: Optional[NNCheckpoint] = NNCheckpoint.load(
        run=run.id,
        type=Checkpoints.BEST,
    )

    Utils.print_table(
        header=False,
        title="Trainer Run Details...",
        data=Utils.flatten_dict(data=run.state()),
    )

    for cb in normalized_callbacks:
        cb.on_train_begin(ctx)

    idx_iter = 0
    tqdm_disabled = os.environ.get("NNX_TQDM_DISABLE", "").lower() in {"1", "true", "yes"}
    with (
        torch.set_grad_enabled(True),
        tqdm(colour="blue", total=n_iter, desc="Training", disable=tqdm_disabled) as tqdm_bar,
        _CallbackFinalizer(normalized_callbacks, ctx),
    ):
        for idx_epoch in range(params.n_epochs):
            ctx.epoch = idx_epoch
            for cb in normalized_callbacks:
                cb.on_epoch_begin(ctx)

            n_idps_before_epoch = len(idps)
            for idx_batch, batch in enumerate(params.train_loader):
                step_ctx = TrainerStepContext(
                    model=self.model,
                    batch=batch,
                    optimizers=optimizers,
                    schedulers=schedulers,
                    extra_metrics=params.extra_metrics,
                    batch_idx=idx_batch,
                    epoch_idx=idx_epoch,
                )
                train_edp = trainer_step_fn(step_ctx)

                idps.append(
                    NNIterationDataPoint(
                        iter_idx=idx_iter,
                        epoch_idx=idx_epoch,
                        batch_idx=idx_batch,
                        train_edp=train_edp,
                        lr=optimizers[primary].param_groups[0]["lr"],
                    )
                )
                idx_iter += 1
                tqdm_bar.update(1)

            if len(idps) == n_idps_before_epoch:
                # Same guard as NNModel.train: zero batches would
                # crash on idps[-1] (first epoch) or corrupt the
                # previous epoch's logged metrics (later epochs).
                raise ValueError(
                    f"train_loader yielded no batches in epoch {idx_epoch} — check batch_size vs "
                    "dataset size with drop_last=True, or whether the loader is a one-shot iterable."
                )

            val_edp = (
                self.model.evaluate(
                    loader=params.val_loader,
                    extra_metrics=params.extra_metrics,
                )
                if validate
                else None
            )
            idps[-1] = idps[-1].with_val_edp(val_edp)

            checkpoint = self._save_checkpoint(
                idp=idps[-1],
                run_id=run.id,
                idx_epoch=idx_epoch,
                n_epochs=params.n_epochs,
                best_checkpoint=best_checkpoint,
                save_phase_checkpoints=params.save_phase_checkpoints,
            )
            if best_checkpoint is None or _best_err(checkpoint) < _best_err(best_checkpoint):
                best_checkpoint = checkpoint

            # Each scheduler steps on its own optimizer's signal.
            # We feed the SAME (val_edp, train_edp) pair to all of
            # them because IDPs aggregate over all optims — separating
            # per-optim metrics would require the step fn to return
            # multiple EDPs, which complicates the contract without
            # clear benefit. Users who need per-optim scheduler
            # signals can step their scheduler directly in the step fn.
            for sched in schedulers.values():
                _step_scheduler(sched, val_edp, train_edp)

            # Verbatim-shared with the NNModel.train loop — one
            # implementation, so the postfix format can't drift.
            self.model._update_tqdm_postfix(tqdm_bar, optimizers[primary], val_edp, train_edp)

            run.with_idps(idps).save()

            ctx.idp = idps[-1]
            ctx.idps = idps
            for cb in normalized_callbacks:
                cb.on_epoch_end(ctx)

            if ctx.should_stop:
                break

    saved = run.with_idps(idps).save()
    print()
    runs_root_path = os.path.join(os.getcwd(), "runs", run.id)
    print(f"Run saved to {runs_root_path}")
    return saved

nnx.trainer.trainer.TrainerStepContext dataclass

Per-batch state passed into a trainer_step_fn.

Mirrors TrainStepContext from NNModel.train() but with optimizer (singular) replaced by optimizers (name-keyed dict) and schedulers threaded through alongside — multi-optim hooks may want to reach into either at any point during a step (e.g., for warmup logic).

model is the single NNModel the Trainer was constructed with; model.net carries the actual nn.Module (which may itself be a composite, e.g., a GAN-style wrapper exposing G and D as submodules).

Source code in nnx/trainer/trainer.py

@dataclass(frozen=True, slots=True)
class TrainerStepContext:
    """Per-batch state passed into a trainer_step_fn.

    Mirrors TrainStepContext from NNModel.train() but with `optimizer`
    (singular) replaced by `optimizers` (name-keyed dict) and `schedulers`
    threaded through alongside — multi-optim hooks may want to reach
    into either at any point during a step (e.g., for warmup logic).

    `model` is the single NNModel the Trainer was constructed with;
    `model.net` carries the actual nn.Module (which may itself be a
    composite, e.g., a GAN-style wrapper exposing G and D as submodules).
    """

    model: NNModel
    batch: Any
    optimizers: Mapping[str, torch.optim.Optimizer]
    schedulers: Mapping[str, Any]
    extra_metrics: Optional[Mapping[str, Callable]]
    batch_idx: int
    epoch_idx: int

nnx.trainer.trainer.TrainerStepFn = Callable[[TrainerStepContext], NNEvaluationDataPoint] module-attribute

nnx.trainer.params.NNTrainerParams dataclass

Configuration for Trainer.train() — the multi-optimizer parallel to NNModel.train() / NNTrainParams.

optims is a name-keyed mapping of NNOptimParams; each entry produces a distinct torch Optimizer. Use NNOptimParams.param_groups on each entry (the fine-tuning hook from :mod:nnx.finetune) to scope an optimizer to a subset of the model's parameters — e.g., one optim for the generator sub-net (name_pattern="G.*"), one for the discriminator (name_pattern="D.*") inside a single combined NNModel.

schedulers is similarly keyed and indexes the same names. Missing entries default to ReduceLROnPlateau with the same defaults NNTrainParams uses, so callers only have to populate schedulers for the optims they want to customize.

seed, save_phase_checkpoints, extra_metrics, train_loader, val_loader mirror NNTrainParams exactly — the orchestration around a single user-supplied trainer_step_fn is otherwise identical.

Source code in nnx/trainer/params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNTrainerParams:
    """Configuration for `Trainer.train()` — the multi-optimizer parallel
    to `NNModel.train()` / `NNTrainParams`.

    `optims` is a name-keyed mapping of NNOptimParams; each entry
    produces a distinct torch Optimizer. Use `NNOptimParams.param_groups`
    on each entry (the fine-tuning hook from :mod:`nnx.finetune`) to scope an optimizer
    to a subset of the model's parameters — e.g., one optim for the
    generator sub-net (`name_pattern="G.*"`), one for the discriminator
    (`name_pattern="D.*"`) inside a single combined NNModel.

    `schedulers` is similarly keyed and indexes the same names. Missing
    entries default to ReduceLROnPlateau with the same defaults
    NNTrainParams uses, so callers only have to populate schedulers for
    the optims they want to customize.

    `seed`, `save_phase_checkpoints`, `extra_metrics`, `train_loader`,
    `val_loader` mirror NNTrainParams exactly — the orchestration around
    a single user-supplied `trainer_step_fn` is otherwise identical.
    """

    n_epochs: int
    optims: Mapping[str, NNOptimParams]
    schedulers: Mapping[str, NNSchedulerParams] = field(default_factory=dict)

    seed: Optional[int] = None
    save_phase_checkpoints: bool = True

    train_loader: Optional[DataLoader] = field(repr=False, default=None)
    val_loader: Optional[DataLoader] = field(repr=False, default=None)

    extra_metrics: Optional[Mapping[str, Callable]] = field(repr=False, default=None)

    def __post_init__(self):
        # Fail-fast: `n_epochs` drives `range(params.n_epochs)` in Trainer.train,
        # so a value < 1 silently makes training a no-op. Symmetric with
        # NNTrainParams.__post_init__.
        if self.n_epochs < 1:
            raise ValueError(f"NNTrainerParams requires n_epochs >= 1, got {self.n_epochs}")
        if not self.optims:
            raise ValueError(
                "NNTrainerParams.optims must have at least one entry — the Trainer constructs one Optimizer per name."
            )
        unknown = set(self.schedulers.keys()) - set(self.optims.keys())
        if unknown:
            raise ValueError(
                "NNTrainerParams.schedulers has keys not present in optims: "
                f"{sorted(unknown)} (known optim names: {sorted(self.optims.keys())})"
            )

    def with_train_loader(self, value: DataLoader) -> NNTrainerParams:
        return replace(self, train_loader=value)

    def with_val_loader(self, value: DataLoader) -> NNTrainerParams:
        return replace(self, val_loader=value)

    def __str__(self):
        return f"Trainer={{n_epochs={self.n_epochs}, optims={sorted(self.optims.keys())}, seed={self.seed}}}"

    def state(self):
        # Keys are sorted so dict insertion order doesn't affect run.id.
        d = dict(
            n_epochs=self.n_epochs,
            optims={k: self.optims[k].state() for k in sorted(self.optims.keys())},
        )
        # Match NNTrainParams: emit `schedulers` / `seed` /
        # `save_phase_checkpoints` only when set to a non-default value, so a
        # trainer run with the defaults hashes stably across versions and
        # follows the project-wide omit-when-default convention.
        if self.schedulers:
            d["schedulers"] = {k: self.schedulers[k].state() for k in sorted(self.schedulers.keys())}
        if self.seed is not None:
            d["seed"] = self.seed
        if self.save_phase_checkpoints is not True:
            d["save_phase_checkpoints"] = self.save_phase_checkpoints
        return d

    @staticmethod
    def from_state(state: dict) -> NNTrainerParams:
        return NNTrainerParams(
            n_epochs=state["n_epochs"],
            optims={k: NNOptimParams.from_state(v) for k, v in state["optims"].items()},
            schedulers={k: NNSchedulerParams.from_state(v) for k, v in state.get("schedulers", {}).items()},
            seed=state.get("seed"),
            save_phase_checkpoints=state.get("save_phase_checkpoints", True),
        )

    @classmethod
    def builder(cls) -> NNTrainerParamsBuilder:
        """Return a composite multi-optim builder. See
        `NNTrainerParamsBuilder`. Composes
        `NNOptimParams.builder()` + `NNSchedulerParams.builder()`.
        """
        from .params_builder import NNTrainerParamsBuilder

        return NNTrainerParamsBuilder()

builder() -> NNTrainerParamsBuilder classmethod

Return a composite multi-optim builder. See NNTrainerParamsBuilder. Composes NNOptimParams.builder() + NNSchedulerParams.builder().

Source code in nnx/trainer/params.py

@classmethod
def builder(cls) -> NNTrainerParamsBuilder:
    """Return a composite multi-optim builder. See
    `NNTrainerParamsBuilder`. Composes
    `NNOptimParams.builder()` + `NNSchedulerParams.builder()`.
    """
    from .params_builder import NNTrainerParamsBuilder

    return NNTrainerParamsBuilder()

nnx.trainer.params_builder.NNTrainerParamsBuilder

Composite builder for NNTrainerParams.

Reach via NNTrainerParams.builder(). The required setter is .n_epochs(N); at least one .optimizer(name, params) call is also required (NNTrainerParams.__post_init__ rejects empty optims). Schedulers, seed, loaders, etc. are all chained optionals.

Source code in nnx/trainer/params_builder.py

class NNTrainerParamsBuilder:
    """Composite builder for `NNTrainerParams`.

    Reach via `NNTrainerParams.builder()`. The required setter is
    `.n_epochs(N)`; at least one `.optimizer(name, params)` call is
    also required (`NNTrainerParams.__post_init__` rejects empty
    optims). Schedulers, seed, loaders, etc. are all chained optionals.
    """

    def __init__(self) -> None:
        self._fields: dict[str, Any] = {}
        self._optims: dict[str, NNOptimParams] = {}
        self._schedulers: dict[str, NNSchedulerParams] = {}

    def n_epochs(self, n: int) -> NNTrainerParamsBuilder:
        """Number of training epochs. Required."""
        self._fields["n_epochs"] = n
        return self

    def optimizer(self, name: str, params: NNOptimParams) -> NNTrainerParamsBuilder:
        """Register one optimizer under `name`. Each name gets its
        own torch.optim.Optimizer at Trainer.train() time. Use
        `NNOptimParams.builder()` (Plan 2) to construct `params`."""
        self._optims[name] = params
        return self

    def scheduler(self, name: str, params: NNSchedulerParams) -> NNTrainerParamsBuilder:
        """Register one scheduler under `name`. The name must match a
        previously-registered `.optimizer(name, ...)` call — `.build()`
        enforces the subset invariant."""
        self._schedulers[name] = params
        return self

    def seed(self, value: int) -> NNTrainerParamsBuilder:
        """Seed for reproducibility. None at default (no seeding via
        params; the caller's `set_seed()` is the only path)."""
        self._fields["seed"] = value
        return self

    def save_phase_checkpoints(self, value: bool) -> NNTrainerParamsBuilder:
        """Whether to write phase checkpoints (FIRST / Q1 / Q2 / Q3 /
        LAST / BEST). Default True. The fluent contract is "last call
        wins" — a prior `.save_phase_checkpoints(False)` followed by
        `.save_phase_checkpoints(True)` leaves the dataclass at the
        default (which `state()` then omits)."""
        self._fields["save_phase_checkpoints"] = value
        return self

    def train_loader(self, loader: DataLoader) -> NNTrainerParamsBuilder:
        """Training DataLoader. Optional at Builder time (can be wired
        later via NNTrainerParams.with_train_loader)."""
        self._fields["train_loader"] = loader
        return self

    def val_loader(self, loader: DataLoader) -> NNTrainerParamsBuilder:
        """Validation DataLoader. Optional at Builder time (can be wired
        later via NNTrainerParams.with_val_loader)."""
        self._fields["val_loader"] = loader
        return self

    def extra_metrics(self, metrics: Mapping[str, Callable]) -> NNTrainerParamsBuilder:
        """Extra metrics callables, name-keyed. Each is called with
        (y_pred, y_true) at every validation step."""
        self._fields["extra_metrics"] = metrics
        return self

    def build(self) -> NNTrainerParams:
        """Validate the key-subset invariant, then construct the dataclass.

        `schedulers.keys() ⊆ optims.keys()` is the contract
        `NNTrainerParams.__post_init__` enforces. We check here so the
        user sees the violation at the Builder boundary — e.g., they
        called `.scheduler("d", ...)` without first calling
        `.optimizer("d", ...)` — rather than at the dataclass ctor.

        `n_epochs` has no meaningful default — call `.n_epochs(N)` before
        `.build()`. Caught here too, for the same Builder-boundary reason.

        Raises:
            ValueError: if `.n_epochs(N)` was not called before
                `.build()`, OR if a `.scheduler(name, ...)` was
                attached for a name that has no corresponding
                `.optimizer(name, ...)`. Both messages name the
                Builder methods to call so the user can fix the chain
                without consulting the dataclass schema.
        """
        if "n_epochs" not in self._fields:
            raise ValueError(
                "NNTrainerParamsBuilder.n_epochs() must be called before .build() — "
                "n_epochs has no meaningful default. Example: "
                ".n_epochs(50).optimizer('main', NNOptimParams(...)).build()"
            )
        unknown = set(self._schedulers.keys()) - set(self._optims.keys())
        if unknown:
            raise ValueError(
                "NNTrainerParamsBuilder.scheduler() called with names not present in optims: "
                f"{sorted(unknown)}. Call .optimizer({sorted(unknown)[0]!r}, ...) before scheduling. "
                f"Known optim names so far: {sorted(self._optims.keys())}"
            )
        # Always include optims (required by NNTrainerParams.__post_init__).
        # Only include schedulers if non-empty (preserves omit-when-default in state()).
        fields = dict(self._fields)
        fields["optims"] = self._optims
        if self._schedulers:
            fields["schedulers"] = self._schedulers
        return NNTrainerParams(**fields)

build() -> NNTrainerParams

Validate the key-subset invariant, then construct the dataclass.

schedulers.keys() ⊆ optims.keys() is the contract NNTrainerParams.__post_init__ enforces. We check here so the user sees the violation at the Builder boundary — e.g., they called .scheduler("d", ...) without first calling .optimizer("d", ...) — rather than at the dataclass ctor.

n_epochs has no meaningful default — call .n_epochs(N) before .build(). Caught here too, for the same Builder-boundary reason.

Raises:

Type	Description
ValueError	if .n_epochs(N) was not called before .build(), OR if a .scheduler(name, ...) was attached for a name that has no corresponding .optimizer(name, ...). Both messages name the Builder methods to call so the user can fix the chain without consulting the dataclass schema.

Source code in nnx/trainer/params_builder.py

def build(self) -> NNTrainerParams:
    """Validate the key-subset invariant, then construct the dataclass.

    `schedulers.keys() ⊆ optims.keys()` is the contract
    `NNTrainerParams.__post_init__` enforces. We check here so the
    user sees the violation at the Builder boundary — e.g., they
    called `.scheduler("d", ...)` without first calling
    `.optimizer("d", ...)` — rather than at the dataclass ctor.

    `n_epochs` has no meaningful default — call `.n_epochs(N)` before
    `.build()`. Caught here too, for the same Builder-boundary reason.

    Raises:
        ValueError: if `.n_epochs(N)` was not called before
            `.build()`, OR if a `.scheduler(name, ...)` was
            attached for a name that has no corresponding
            `.optimizer(name, ...)`. Both messages name the
            Builder methods to call so the user can fix the chain
            without consulting the dataclass schema.
    """
    if "n_epochs" not in self._fields:
        raise ValueError(
            "NNTrainerParamsBuilder.n_epochs() must be called before .build() — "
            "n_epochs has no meaningful default. Example: "
            ".n_epochs(50).optimizer('main', NNOptimParams(...)).build()"
        )
    unknown = set(self._schedulers.keys()) - set(self._optims.keys())
    if unknown:
        raise ValueError(
            "NNTrainerParamsBuilder.scheduler() called with names not present in optims: "
            f"{sorted(unknown)}. Call .optimizer({sorted(unknown)[0]!r}, ...) before scheduling. "
            f"Known optim names so far: {sorted(self._optims.keys())}"
        )
    # Always include optims (required by NNTrainerParams.__post_init__).
    # Only include schedulers if non-empty (preserves omit-when-default in state()).
    fields = dict(self._fields)
    fields["optims"] = self._optims
    if self._schedulers:
        fields["schedulers"] = self._schedulers
    return NNTrainerParams(**fields)

extra_metrics(metrics: Mapping[str, Callable]) -> NNTrainerParamsBuilder

Extra metrics callables, name-keyed. Each is called with (y_pred, y_true) at every validation step.

Source code in nnx/trainer/params_builder.py

def extra_metrics(self, metrics: Mapping[str, Callable]) -> NNTrainerParamsBuilder:
    """Extra metrics callables, name-keyed. Each is called with
    (y_pred, y_true) at every validation step."""
    self._fields["extra_metrics"] = metrics
    return self

n_epochs(n: int) -> NNTrainerParamsBuilder

Number of training epochs. Required.

Source code in nnx/trainer/params_builder.py

def n_epochs(self, n: int) -> NNTrainerParamsBuilder:
    """Number of training epochs. Required."""
    self._fields["n_epochs"] = n
    return self

optimizer(name: str, params: NNOptimParams) -> NNTrainerParamsBuilder

Register one optimizer under name. Each name gets its own torch.optim.Optimizer at Trainer.train() time. Use NNOptimParams.builder() (Plan 2) to construct params.

Source code in nnx/trainer/params_builder.py

def optimizer(self, name: str, params: NNOptimParams) -> NNTrainerParamsBuilder:
    """Register one optimizer under `name`. Each name gets its
    own torch.optim.Optimizer at Trainer.train() time. Use
    `NNOptimParams.builder()` (Plan 2) to construct `params`."""
    self._optims[name] = params
    return self

save_phase_checkpoints(value: bool) -> NNTrainerParamsBuilder

Whether to write phase checkpoints (FIRST / Q1 / Q2 / Q3 / LAST / BEST). Default True. The fluent contract is "last call wins" — a prior .save_phase_checkpoints(False) followed by .save_phase_checkpoints(True) leaves the dataclass at the default (which state() then omits).

Source code in nnx/trainer/params_builder.py

def save_phase_checkpoints(self, value: bool) -> NNTrainerParamsBuilder:
    """Whether to write phase checkpoints (FIRST / Q1 / Q2 / Q3 /
    LAST / BEST). Default True. The fluent contract is "last call
    wins" — a prior `.save_phase_checkpoints(False)` followed by
    `.save_phase_checkpoints(True)` leaves the dataclass at the
    default (which `state()` then omits)."""
    self._fields["save_phase_checkpoints"] = value
    return self

scheduler(name: str, params: NNSchedulerParams) -> NNTrainerParamsBuilder

Register one scheduler under name. The name must match a previously-registered .optimizer(name, ...) call — .build() enforces the subset invariant.

Source code in nnx/trainer/params_builder.py

def scheduler(self, name: str, params: NNSchedulerParams) -> NNTrainerParamsBuilder:
    """Register one scheduler under `name`. The name must match a
    previously-registered `.optimizer(name, ...)` call — `.build()`
    enforces the subset invariant."""
    self._schedulers[name] = params
    return self

seed(value: int) -> NNTrainerParamsBuilder

Seed for reproducibility. None at default (no seeding via params; the caller's set_seed() is the only path).

Source code in nnx/trainer/params_builder.py

def seed(self, value: int) -> NNTrainerParamsBuilder:
    """Seed for reproducibility. None at default (no seeding via
    params; the caller's `set_seed()` is the only path)."""
    self._fields["seed"] = value
    return self

train_loader(loader: DataLoader) -> NNTrainerParamsBuilder

Training DataLoader. Optional at Builder time (can be wired later via NNTrainerParams.with_train_loader).

Source code in nnx/trainer/params_builder.py

def train_loader(self, loader: DataLoader) -> NNTrainerParamsBuilder:
    """Training DataLoader. Optional at Builder time (can be wired
    later via NNTrainerParams.with_train_loader)."""
    self._fields["train_loader"] = loader
    return self

val_loader(loader: DataLoader) -> NNTrainerParamsBuilder

Validation DataLoader. Optional at Builder time (can be wired later via NNTrainerParams.with_val_loader).

Source code in nnx/trainer/params_builder.py

def val_loader(self, loader: DataLoader) -> NNTrainerParamsBuilder:
    """Validation DataLoader. Optional at Builder time (can be wired
    later via NNTrainerParams.with_val_loader)."""
    self._fields["val_loader"] = loader
    return self

3. Params

nnx.nn.params.nn_params.NNParams dataclass

Source code in nnx/nn/params/nn_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNParams:
    dropout_prob: float
    n_heads: Optional[int] = field(default=None)
    activation: Optional[Activations] = field(default=Activations.LEAKY_RELU)

    input_dim: int = field(repr=False)
    output_dim: int = field(repr=False)
    hidden_dims: Optional[list[int]] = field(repr=False, default=None)
    _dims: Optional[list[int]] = field(repr=False, init=False, default=None)

    @property
    def dims(self) -> list[int]:
        # `_dims` is set unconditionally in __post_init__, so the
        # Optional type on the field is purely a dataclass artifact
        # (no init=True default for slotted frozen subclasses). The
        # assert documents that contract for both readers and the
        # type-checker (pyright can't model __post_init__-set fields
        # via `object.__setattr__`).
        assert self._dims is not None
        return self._dims

    def __post_init__(self):
        # Fail-fast on out-of-range numeric fields at construction time rather
        # than deep inside layer building. nn.Dropout / nn.Linear would raise
        # eventually for these, but far from the origin — surfacing the error
        # here keeps the [[params-boundary-validation]] contract consistent
        # with the rest of the params hierarchy. None of these touch state().
        if not 0.0 <= self.dropout_prob <= 1.0:
            raise ValueError(f"NNParams requires 0.0 <= dropout_prob <= 1.0, got {self.dropout_prob}")
        if self.input_dim <= 0:
            raise ValueError(f"NNParams requires input_dim > 0, got {self.input_dim}")
        if self.output_dim <= 0:
            raise ValueError(f"NNParams requires output_dim > 0, got {self.output_dim}")
        if self.hidden_dims is not None and not all(d > 0 for d in self.hidden_dims):
            raise ValueError(f"NNParams requires all hidden_dims > 0, got {self.hidden_dims}")

        dims = [self.input_dim]
        dims += self.hidden_dims if self.hidden_dims is not None else []
        dims += [self.output_dim]

        object.__setattr__(self, "_dims", dims)

    def state(self) -> dict:
        ret = dict(
            input_dim=self.input_dim,
            output_dim=self.output_dim,
            dropout_prob=self.dropout_prob,
            hidden_dims=str(self.hidden_dims),
            # None stays None (yaml null / json null) rather than the
            # string "None", which from_state could never parse back
            # into the Activations enum.
            activation=str(self.activation) if self.activation is not None else None,
        )

        if self.n_heads is not None:
            ret["n_heads"] = self.n_heads

        return ret

    @staticmethod
    def from_state(state: dict) -> NNParams:
        raw_activation = state["activation"]
        return NNParams(
            input_dim=state["input_dim"],
            output_dim=state["output_dim"],
            dropout_prob=state["dropout_prob"],
            activation=Activations(raw_activation) if raw_activation is not None else None,
            hidden_dims=ast.literal_eval(state["hidden_dims"]),
            n_heads=state["n_heads"] if "n_heads" in state else None,
        )

    @staticmethod
    def resolve_from_state(state: dict) -> NNParams:
        """Dispatch to the params subclass that wrote ``state``.

        ``NNTransformerParams.state()`` always emits its required
        architectural keys (``vocab_size`` among them); base
        ``NNParams.state()`` never does. Without this dispatch a
        transformer state is silently downgraded to base ``NNParams`` —
        the subclass keys are dropped, the reloaded run re-hashes to a
        different id, and net rebuilding crashes. Every loader
        (``NNRun.load``, the ``NNCheckpoint`` readers, hub
        ``from_pretrained``) resolves through here.
        """
        if "vocab_size" in state:
            # Local import: nn_transformer_params imports this module.
            from .nn_transformer_params import NNTransformerParams

            return NNTransformerParams.from_state(state)
        return NNParams.from_state(state)

resolve_from_state(state: dict) -> NNParams staticmethod

Dispatch to the params subclass that wrote state.

NNTransformerParams.state() always emits its required architectural keys (vocab_size among them); base NNParams.state() never does. Without this dispatch a transformer state is silently downgraded to base NNParams — the subclass keys are dropped, the reloaded run re-hashes to a different id, and net rebuilding crashes. Every loader (NNRun.load, the NNCheckpoint readers, hub from_pretrained) resolves through here.

Source code in nnx/nn/params/nn_params.py

@staticmethod
def resolve_from_state(state: dict) -> NNParams:
    """Dispatch to the params subclass that wrote ``state``.

    ``NNTransformerParams.state()`` always emits its required
    architectural keys (``vocab_size`` among them); base
    ``NNParams.state()`` never does. Without this dispatch a
    transformer state is silently downgraded to base ``NNParams`` —
    the subclass keys are dropped, the reloaded run re-hashes to a
    different id, and net rebuilding crashes. Every loader
    (``NNRun.load``, the ``NNCheckpoint`` readers, hub
    ``from_pretrained``) resolves through here.
    """
    if "vocab_size" in state:
        # Local import: nn_transformer_params imports this module.
        from .nn_transformer_params import NNTransformerParams

        return NNTransformerParams.from_state(state)
    return NNParams.from_state(state)

nnx.nn.params.nn_model_params.NNModelParams dataclass

Source code in nnx/nn/params/nn_model_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNModelParams:
    net: Nets
    device: Devices = Devices.CPU
    loss: Losses = Losses.CROSS_ENTROPY

    # Opt-in fp16 autocast + GradScaler in train(). Only effective on CUDA;
    # silently bypassed on CPU/MPS where torch.cuda.amp is a no-op or unavailable.
    mixed_precision: bool = False

    def __str__(self) -> str:
        return f"[net={self.net}, device={self.device}, loss={self.loss}, mixed_precision={self.mixed_precision}]"

    def is_valid(self) -> bool:
        return self.net is not None and self.device is not None and self.loss is not None

    def state(self) -> dict:
        d = dict(
            net=str(self.net),
            loss=str(self.loss),
            device=str(self.device),
        )
        # `mixed_precision` is omitted from state() when False so a
        # NNModelParams without AMP enabled hashes to the same run.id
        # as before this field existed. Same omit-when-default invariant
        # as NNTrainParams.seed / NNOptimParams.param_groups.
        if self.mixed_precision:
            d["mixed_precision"] = True
        return d

    @staticmethod
    def from_state(state: dict) -> NNModelParams:
        return NNModelParams(
            net=Nets(state["net"]),
            loss=Losses(state["loss"]),
            device=Devices(state["device"]),
            mixed_precision=state.get("mixed_precision", False),
        )

nnx.nn.params.nn_train_params.NNTrainParams dataclass

Training configuration.

seed pins every RNG that affects training (Python random, NumPy, torch CPU+CUDA, cuDNN) when NNModel.train() runs. None disables seeding (default).

To preserve back-compat with previously-saved runs, seed is included in state() ONLY when set — so existing runs with no seed continue to hash to the same run.id.

Source code in nnx/nn/params/nn_train_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNTrainParams:
    """Training configuration.

    `seed` pins every RNG that affects training (Python random, NumPy,
    torch CPU+CUDA, cuDNN) when NNModel.train() runs. None disables
    seeding (default).

    To preserve back-compat with previously-saved runs, `seed` is included
    in state() ONLY when set — so existing runs with no seed continue to
    hash to the same `run.id`.
    """

    n_epochs: int
    scheduler: NNSchedulerParams = NNSchedulerParams(patience=8, cooldown=2, factor=95e-2, threshold=1e-3, min_lr=1e-7)
    optim: NNOptimParams = NNOptimParams(name=Optims.ADAM, max_lr=1e-2, momentum=(0.9, 0.999), weight_decay=5e-5)

    seed: Optional[int] = None

    # When True (default, back-compat), train() saves FIRST + Q1 + Q2 + Q3
    # phase checkpoints in addition to LAST + BEST. Set False to skip the
    # FIRST/Q* writes — useful for tiny experiments or huge models where
    # per-epoch checkpoint I/O dominates wall-clock time.
    save_phase_checkpoints: bool = True

    train_loader: Optional[DataLoader] = field(repr=False, default=None)
    val_loader: Optional[DataLoader] = field(repr=False, default=None)

    # Custom metrics: name -> callable(Y_true, Y_pred) -> float. Runtime-only
    # (functions don't round-trip through YAML), so this lives outside
    # state() / from_state() — like train_loader/val_loader. Each is invoked
    # on every train batch and on every evaluate() aggregate.
    extra_metrics: Optional[Mapping[str, Callable]] = field(repr=False, default=None)

    # Resume control. When `resume_from_run_id` is set, train() loads that
    # run's checkpoint of the named type and warm-restarts training from
    # its model weights AND optimizer state (when an .opt.pt sidecar exists).
    # Runtime-only (knowing the prior run id isn't part of *this* run's
    # configuration identity). Consequence: a resumed run with otherwise
    # identical params hashes to the SAME run.id as the original, so its
    # incremental saves overwrite the original run's run.yaml/idps.csv
    # with this session's history only. Vary a state()-bearing field
    # (e.g. n_epochs, seed) if the original history must be kept.
    resume_from_run_id: Optional[str] = field(repr=False, default=None)
    resume_from_checkpoint: Optional[str] = field(repr=False, default="last")

    def __post_init__(self):
        # Fail-fast: `n_epochs` drives `range(params.n_epochs)` in the train
        # loop, so a value < 1 silently makes training a no-op (empty idps, a
        # degenerate saved run, no BEST checkpoint) rather than erroring.
        # `n_epochs` is always emitted into state(), so this never shifts a
        # run.id for any valid config.
        if self.n_epochs < 1:
            raise ValueError(f"NNTrainParams requires n_epochs >= 1, got {self.n_epochs}")

    def with_train_loader(self, value: DataLoader) -> NNTrainParams:
        return replace(self, train_loader=value)

    def with_val_loader(self, value: DataLoader) -> NNTrainParams:
        return replace(self, val_loader=value)

    def __str__(self):
        return f"Train={{n_epochs={self.n_epochs}, seed={self.seed}, Optim={self.optim}, Scheduler={self.scheduler}}}"

    def state(self):
        d = dict(n_epochs=self.n_epochs, optim=self.optim.state(), scheduler=self.scheduler.state())
        # Only emit `seed` / `save_phase_checkpoints` into state() when they
        # diverge from their defaults so a NNTrainParams created without
        # them hashes to the same run.id as before these fields existed.
        # Existing on-disk runs without these keys are loadable via .get()
        # defaults in from_state below.
        if self.seed is not None:
            d["seed"] = self.seed
        if self.save_phase_checkpoints is not True:
            d["save_phase_checkpoints"] = self.save_phase_checkpoints
        return d

    @staticmethod
    def from_state(state: dict) -> NNTrainParams:
        return NNTrainParams(
            n_epochs=state["n_epochs"],
            optim=NNOptimParams.from_state(state["optim"]),
            scheduler=NNSchedulerParams.from_state(state["scheduler"]),
            seed=state.get("seed"),
            save_phase_checkpoints=state.get("save_phase_checkpoints", True),
        )

nnx.nn.params.nn_optim_params.NNOptimParams dataclass

Optimizer config.

momentum is overloaded by optimizer kind: - For SGD / SGD_NESTEROV: a single float, the SGD momentum coefficient. - For ADAM / ADAM_AMSGRAD: a (beta1, beta2) tuple, passed as the Adam betas= argument. The name is retained for backwards compatibility — is_valid() enforces the per-optim shape.

grad_clip_norm clips gradients by global L2 norm before optimizer.step(). None = no clipping (back-compat default). Typical values: 1.0 for transformers, 5.0 for RNNs.

accumulate_grad_batches enables gradient accumulation — the effective batch size becomes batch_size * accumulate_grad_batches. The loss is scaled by 1/N so the accumulated gradient is the mean across N batches. Default 1 (back-compat: step every batch).

param_groups enables per-layer-group LR / weight_decay overrides — the fine-tuning idiom of "small LR on the backbone, large LR on the head." None = single-group behavior (every parameter at max_lr / weight_decay). When set, the optimizer factory dispatches via :func:nnx.finetune.param_groups.build_param_groups to construct per-group dicts.

Source code in nnx/nn/params/nn_optim_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNOptimParams:
    """Optimizer config.

    `momentum` is overloaded by optimizer kind:
      - For SGD / SGD_NESTEROV: a single float, the SGD momentum coefficient.
      - For ADAM / ADAM_AMSGRAD: a (beta1, beta2) tuple, passed as the
        Adam `betas=` argument. The name is retained for backwards
        compatibility — `is_valid()` enforces the per-optim shape.

    `grad_clip_norm` clips gradients by global L2 norm before optimizer.step().
    None = no clipping (back-compat default). Typical values: 1.0 for
    transformers, 5.0 for RNNs.

    `accumulate_grad_batches` enables gradient accumulation — the effective
    batch size becomes batch_size * accumulate_grad_batches. The loss is
    scaled by 1/N so the accumulated gradient is the mean across N batches.
    Default 1 (back-compat: step every batch).

    `param_groups` enables per-layer-group LR / weight_decay overrides — the
    fine-tuning idiom of "small LR on the backbone, large LR on the head."
    None = single-group behavior (every parameter at `max_lr` / `weight_decay`).
    When set, the optimizer factory dispatches via
    :func:`nnx.finetune.param_groups.build_param_groups` to construct
    per-group dicts.
    """

    name: Optims
    max_lr: float
    weight_decay: float
    momentum: Union[float, tuple[float, float]]

    grad_clip_norm: Optional[float] = None
    accumulate_grad_batches: int = 1
    param_groups: Optional[list[NNParamGroupSpec]] = field(default=None)

    def __post_init__(self):
        # Fail fast on out-of-range scalars — both construct fine but
        # misbehave silently/obscurely deep in the train loop:
        #   * accumulate_grad_batches < 1: =0 dies mid-training with
        #     `ZeroDivisionError` on `batch_idx % accumulate_grad_batches`
        #     (AFTER printing the whole run-config table); <0 scales the
        #     loss by 1/N < 0 and silently performs gradient *ascent*.
        #   * grad_clip_norm <= 0 (when not the None "off" sentinel): 0.0
        #     passes the `is not None` clip-enable check and zeros every
        #     gradient, so training runs to completion making no progress.
        if self.accumulate_grad_batches < 1:
            raise ValueError(
                f"accumulate_grad_batches must be >= 1, got {self.accumulate_grad_batches} "
                "(1 = step every batch; N = step every N batches)."
            )
        if self.grad_clip_norm is not None and self.grad_clip_norm <= 0:
            raise ValueError(
                f"grad_clip_norm must be > 0 when set, got {self.grad_clip_norm} (use None to disable clipping, not 0)."
            )
        #   * max_lr < 0: a negative LR performs gradient *ascent*. max_lr=0
        #     is allowed — it is an explicit "freeze updates" choice (used as a
        #     no-update idiom, e.g. probing a loss without changing weights);
        #     unlike the knobs above, 0 is intended, not a silent footgun.
        #   * weight_decay < 0: grows weights instead of decaying them (matches
        #     the per-group NNParamGroupSpec guard).
        if self.max_lr < 0:
            raise ValueError(
                f"max_lr must be non-negative, got {self.max_lr} "
                "(a negative LR performs gradient ascent; 0 freezes updates)."
            )
        if self.weight_decay < 0:
            raise ValueError(f"weight_decay must be non-negative, got {self.weight_decay} (use 0 to disable).")
        # Fail fast on plain dicts: they construct fine but crash much
        # later inside state() during NNRun hashing with an opaque
        # AttributeError. Same construction-time convention as the
        # dataset classes.
        if self.param_groups is not None:
            if not isinstance(self.param_groups, (list, tuple)):
                # A generator would be silently EXHAUSTED by the
                # validation loop below — state() would then emit an
                # empty param_groups and training would run single-group
                # with a shifted run.id.
                raise TypeError(
                    f"param_groups must be a list/tuple of NNParamGroupSpec, got {type(self.param_groups).__name__}"
                )
            # Lazy import — keeps this low-level dataclass importable
            # without eagerly loading the finetune subpackage (no cycle
            # today; same deferral style as from_state below).
            from ...finetune.param_groups import NNParamGroupSpec

            for i, g in enumerate(self.param_groups):
                if not isinstance(g, NNParamGroupSpec):
                    raise TypeError(
                        f"param_groups[{i}] must be an NNParamGroupSpec, got {type(g).__name__} — "
                        "wrap it: NNParamGroupSpec(name_pattern=..., lr=...)."
                    )

    def __str__(self):
        return f"[name={self.name}, max_lr={self.max_lr:1.0e}, weight_decay={self.weight_decay:1.0e}, momentum={self.momentum}, grad_clip={self.grad_clip_norm}, accum={self.accumulate_grad_batches}]"

    def state(self):
        d = dict(max_lr=self.max_lr, momentum=str(self.momentum), name=str(self.name), weight_decay=self.weight_decay)
        # grad_clip_norm / accumulate_grad_batches / param_groups: only emit
        # when set to a non-default value, so a NNOptimParams with none of
        # them set hashes to the same run.id as before these fields existed.
        # Existing on-disk YAML without these keys is loadable via the .get()
        # defaults below. (This invariant was broken once before for
        # grad_clip_norm — every existing run.id shifted. Same omit-when-
        # default pattern is now enforced on every params dataclass; see
        # test_params_round_trip.py for the regression tests.)
        if self.grad_clip_norm is not None:
            d["grad_clip_norm"] = self.grad_clip_norm
        if self.accumulate_grad_batches != 1:
            d["accumulate_grad_batches"] = self.accumulate_grad_batches
        if self.param_groups is not None:
            d["param_groups"] = [g.state() for g in self.param_groups]
        return d

    @staticmethod
    def from_state(state: dict) -> NNOptimParams:
        # Lazy import — defers the finetune subpackage so this
        # low-level dataclass stays light at import time (no actual
        # cycle today: param_groups.py imports only stdlib + torch).
        from ...finetune.param_groups import NNParamGroupSpec

        raw_pg = state.get("param_groups")
        param_groups = [NNParamGroupSpec.from_state(g) for g in raw_pg] if raw_pg is not None else None
        return NNOptimParams(
            max_lr=state["max_lr"],
            name=Optims(state["name"]),
            weight_decay=state["weight_decay"],
            momentum=ast.literal_eval(state["momentum"]),
            # .get() preserves back-compat with older YAML that predates
            # grad_clip_norm / accumulate_grad_batches / param_groups.
            grad_clip_norm=state.get("grad_clip_norm"),
            accumulate_grad_batches=state.get("accumulate_grad_batches", 1),
            param_groups=param_groups,
        )

    def is_valid(self) -> bool:
        if self.name == Optims.SGD or self.name == Optims.SGD_NESTEROV:
            return isinstance(self.momentum, float)
        if self.name == Optims.ADAM or self.name == Optims.ADAM_AMSGRAD:
            return (
                isinstance(self.momentum, tuple)
                and len(self.momentum) == 2
                and all(isinstance(x, float) for x in self.momentum)
            )
        # Unknown enum variant — refuse rather than implicitly returning None
        # (which would short-circuit `not params.optim.is_valid()` in train()).
        return False

    @classmethod
    def builder(cls) -> NNOptimParamsBuilder:
        """Return a variant-aware builder. See `NNOptimParamsBuilder`."""
        from .nn_optim_params_builder import NNOptimParamsBuilder

        return NNOptimParamsBuilder()

builder() -> NNOptimParamsBuilder classmethod

Return a variant-aware builder. See NNOptimParamsBuilder.

Source code in nnx/nn/params/nn_optim_params.py

@classmethod
def builder(cls) -> NNOptimParamsBuilder:
    """Return a variant-aware builder. See `NNOptimParamsBuilder`."""
    from .nn_optim_params_builder import NNOptimParamsBuilder

    return NNOptimParamsBuilder()

nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder

Variant-aware builder for NNOptimParams.

Reach via NNOptimParams.builder(). Pick exactly one variant method (adam, adam_amsgrad, sgd, sgd_nesterov), then chain optional methods (grad_clip, accumulate_grad, param_groups), then .build(). Method-call order is independent — a modifier called before a variant survives the variant call, and the last variant always wins.

Source code in nnx/nn/params/nn_optim_params_builder.py

class NNOptimParamsBuilder:
    """Variant-aware builder for `NNOptimParams`.

    Reach via `NNOptimParams.builder()`. Pick exactly one variant
    method (`adam`, `adam_amsgrad`, `sgd`, `sgd_nesterov`), then chain
    optional methods (`grad_clip`, `accumulate_grad`, `param_groups`),
    then `.build()`. Method-call order is independent — a modifier
    called before a variant survives the variant call, and the last
    variant always wins.
    """

    # Fields that a variant method owns. `_set_variant` drops these
    # from `self._fields` before applying the new variant so a second
    # variant call cleanly replaces the first AND any modifier-set
    # keys (grad_clip_norm / accumulate_grad_batches / param_groups)
    # survive.
    _VARIANT_KEYS: ClassVar[tuple[str, ...]] = ("name", "max_lr", "momentum", "weight_decay")

    def __init__(self) -> None:
        self._fields: dict[str, Any] = {}

    def _set_variant(self, **fields: Any) -> None:
        for k in self._VARIANT_KEYS:
            self._fields.pop(k, None)
        self._fields.update(fields)

    # ---------- variant methods ----------

    def adam(
        self,
        *,
        max_lr: float,
        betas: tuple[float, float] = (0.9, 0.999),
        weight_decay: float = 0.0,
    ) -> NNOptimParamsBuilder:
        """torch.optim.Adam. `betas` is PyTorch's name for the
        (beta1, beta2) tuple; the Builder maps it onto the underlying
        `NNOptimParams.momentum` field (which holds the tuple for Adam
        variants).
        """
        self._set_variant(
            name=Optims.ADAM,
            max_lr=max_lr,
            momentum=betas,
            weight_decay=weight_decay,
        )
        return self

    def adam_amsgrad(
        self,
        *,
        max_lr: float,
        betas: tuple[float, float] = (0.9, 0.999),
        weight_decay: float = 0.0,
    ) -> NNOptimParamsBuilder:
        """torch.optim.Adam with `amsgrad=True`. Same `betas` mapping
        as `adam()`.
        """
        self._set_variant(
            name=Optims.ADAM_AMSGRAD,
            max_lr=max_lr,
            momentum=betas,
            weight_decay=weight_decay,
        )
        return self

    def sgd(
        self,
        *,
        max_lr: float,
        momentum: float = 0.9,
        weight_decay: float = 0.0,
    ) -> NNOptimParamsBuilder:
        """torch.optim.SGD. The float momentum stays as `momentum`
        (no rename) — `betas` is an Adam-family term.
        """
        self._set_variant(
            name=Optims.SGD,
            max_lr=max_lr,
            momentum=momentum,
            weight_decay=weight_decay,
        )
        return self

    def sgd_nesterov(
        self,
        *,
        max_lr: float,
        momentum: float = 0.9,
        weight_decay: float = 0.0,
    ) -> NNOptimParamsBuilder:
        """torch.optim.SGD with `nesterov=True`. Same momentum shape
        as `sgd()`.
        """
        self._set_variant(
            name=Optims.SGD_NESTEROV,
            max_lr=max_lr,
            momentum=momentum,
            weight_decay=weight_decay,
        )
        return self

    # ---------- optional modifiers (chain after variant) ----------

    def grad_clip(self, norm: float) -> NNOptimParamsBuilder:
        """Global-L2 gradient-norm clipping. None = no clipping (the
        dataclass default; this method is the opt-in path).
        """
        self._fields["grad_clip_norm"] = norm
        return self

    def accumulate_grad(self, batches: int) -> NNOptimParamsBuilder:
        """Accumulate gradients over `batches` mini-batches before
        stepping. Default (no call) leaves the dataclass at 1.
        """
        self._fields["accumulate_grad_batches"] = batches
        return self

    def param_groups(self, groups: list[NNParamGroupSpec]) -> NNOptimParamsBuilder:
        """Per-layer-group LR / weight_decay overrides (the fine-tuning
        idiom). Default (no call) leaves the dataclass at None
        (single-group behavior).
        """
        self._fields["param_groups"] = groups
        return self

    # ---------- terminator ----------

    def build(self) -> NNOptimParams:
        """Construct the dataclass from the fields the user touched.

        Pre-empts the dataclass's missing-required-argument TypeError
        with an actionable Builder-level ValueError naming the variant
        methods — matches the [[builder-pattern-shape]] §11b convention
        that PR #52 established on NNTrainerParamsBuilder.

        Forwards only the keys present in `self._fields` so the
        dataclass defaults govern every untouched optional field —
        that's what preserves the omit-when-default state() invariant.

        Raises:
            ValueError: if no variant method (`.adam`, `.adam_amsgrad`,
                `.sgd`, `.sgd_nesterov`) was called before `.build()`.
                The message names the four methods so the user can
                fix the chain without consulting the dataclass schema.
        """
        if "name" not in self._fields:
            raise ValueError(
                "NNOptimParamsBuilder: call one of .adam(...), "
                ".adam_amsgrad(...), .sgd(...), or .sgd_nesterov(...) "
                "before .build() — a variant selects the optimizer kind "
                "and sets the required name/max_lr/momentum/weight_decay fields."
            )
        return NNOptimParams(**self._fields)

accumulate_grad(batches: int) -> NNOptimParamsBuilder

Accumulate gradients over batches mini-batches before stepping. Default (no call) leaves the dataclass at 1.

Source code in nnx/nn/params/nn_optim_params_builder.py

def accumulate_grad(self, batches: int) -> NNOptimParamsBuilder:
    """Accumulate gradients over `batches` mini-batches before
    stepping. Default (no call) leaves the dataclass at 1.
    """
    self._fields["accumulate_grad_batches"] = batches
    return self

adam(*, max_lr: float, betas: tuple[float, float] = (0.9, 0.999), weight_decay: float = 0.0) -> NNOptimParamsBuilder

torch.optim.Adam. betas is PyTorch's name for the (beta1, beta2) tuple; the Builder maps it onto the underlying NNOptimParams.momentum field (which holds the tuple for Adam variants).

Source code in nnx/nn/params/nn_optim_params_builder.py

def adam(
    self,
    *,
    max_lr: float,
    betas: tuple[float, float] = (0.9, 0.999),
    weight_decay: float = 0.0,
) -> NNOptimParamsBuilder:
    """torch.optim.Adam. `betas` is PyTorch's name for the
    (beta1, beta2) tuple; the Builder maps it onto the underlying
    `NNOptimParams.momentum` field (which holds the tuple for Adam
    variants).
    """
    self._set_variant(
        name=Optims.ADAM,
        max_lr=max_lr,
        momentum=betas,
        weight_decay=weight_decay,
    )
    return self

adam_amsgrad(*, max_lr: float, betas: tuple[float, float] = (0.9, 0.999), weight_decay: float = 0.0) -> NNOptimParamsBuilder

torch.optim.Adam with amsgrad=True. Same betas mapping as adam().

Source code in nnx/nn/params/nn_optim_params_builder.py

def adam_amsgrad(
    self,
    *,
    max_lr: float,
    betas: tuple[float, float] = (0.9, 0.999),
    weight_decay: float = 0.0,
) -> NNOptimParamsBuilder:
    """torch.optim.Adam with `amsgrad=True`. Same `betas` mapping
    as `adam()`.
    """
    self._set_variant(
        name=Optims.ADAM_AMSGRAD,
        max_lr=max_lr,
        momentum=betas,
        weight_decay=weight_decay,
    )
    return self

build() -> NNOptimParams

Construct the dataclass from the fields the user touched.

Pre-empts the dataclass's missing-required-argument TypeError with an actionable Builder-level ValueError naming the variant methods — matches the [[builder-pattern-shape]] §11b convention that PR #52 established on NNTrainerParamsBuilder.

Forwards only the keys present in self._fields so the dataclass defaults govern every untouched optional field — that's what preserves the omit-when-default state() invariant.

Raises:

Type	Description
ValueError	if no variant method (.adam, .adam_amsgrad, .sgd, .sgd_nesterov) was called before .build(). The message names the four methods so the user can fix the chain without consulting the dataclass schema.

Source code in nnx/nn/params/nn_optim_params_builder.py

def build(self) -> NNOptimParams:
    """Construct the dataclass from the fields the user touched.

    Pre-empts the dataclass's missing-required-argument TypeError
    with an actionable Builder-level ValueError naming the variant
    methods — matches the [[builder-pattern-shape]] §11b convention
    that PR #52 established on NNTrainerParamsBuilder.

    Forwards only the keys present in `self._fields` so the
    dataclass defaults govern every untouched optional field —
    that's what preserves the omit-when-default state() invariant.

    Raises:
        ValueError: if no variant method (`.adam`, `.adam_amsgrad`,
            `.sgd`, `.sgd_nesterov`) was called before `.build()`.
            The message names the four methods so the user can
            fix the chain without consulting the dataclass schema.
    """
    if "name" not in self._fields:
        raise ValueError(
            "NNOptimParamsBuilder: call one of .adam(...), "
            ".adam_amsgrad(...), .sgd(...), or .sgd_nesterov(...) "
            "before .build() — a variant selects the optimizer kind "
            "and sets the required name/max_lr/momentum/weight_decay fields."
        )
    return NNOptimParams(**self._fields)

grad_clip(norm: float) -> NNOptimParamsBuilder

Global-L2 gradient-norm clipping. None = no clipping (the dataclass default; this method is the opt-in path).

Source code in nnx/nn/params/nn_optim_params_builder.py

def grad_clip(self, norm: float) -> NNOptimParamsBuilder:
    """Global-L2 gradient-norm clipping. None = no clipping (the
    dataclass default; this method is the opt-in path).
    """
    self._fields["grad_clip_norm"] = norm
    return self

param_groups(groups: list[NNParamGroupSpec]) -> NNOptimParamsBuilder

Per-layer-group LR / weight_decay overrides (the fine-tuning idiom). Default (no call) leaves the dataclass at None (single-group behavior).

Source code in nnx/nn/params/nn_optim_params_builder.py

def param_groups(self, groups: list[NNParamGroupSpec]) -> NNOptimParamsBuilder:
    """Per-layer-group LR / weight_decay overrides (the fine-tuning
    idiom). Default (no call) leaves the dataclass at None
    (single-group behavior).
    """
    self._fields["param_groups"] = groups
    return self

sgd(*, max_lr: float, momentum: float = 0.9, weight_decay: float = 0.0) -> NNOptimParamsBuilder

torch.optim.SGD. The float momentum stays as momentum (no rename) — betas is an Adam-family term.

Source code in nnx/nn/params/nn_optim_params_builder.py

def sgd(
    self,
    *,
    max_lr: float,
    momentum: float = 0.9,
    weight_decay: float = 0.0,
) -> NNOptimParamsBuilder:
    """torch.optim.SGD. The float momentum stays as `momentum`
    (no rename) — `betas` is an Adam-family term.
    """
    self._set_variant(
        name=Optims.SGD,
        max_lr=max_lr,
        momentum=momentum,
        weight_decay=weight_decay,
    )
    return self

sgd_nesterov(*, max_lr: float, momentum: float = 0.9, weight_decay: float = 0.0) -> NNOptimParamsBuilder

torch.optim.SGD with nesterov=True. Same momentum shape as sgd().

Source code in nnx/nn/params/nn_optim_params_builder.py

def sgd_nesterov(
    self,
    *,
    max_lr: float,
    momentum: float = 0.9,
    weight_decay: float = 0.0,
) -> NNOptimParamsBuilder:
    """torch.optim.SGD with `nesterov=True`. Same momentum shape
    as `sgd()`.
    """
    self._set_variant(
        name=Optims.SGD_NESTEROV,
        max_lr=max_lr,
        momentum=momentum,
        weight_decay=weight_decay,
    )
    return self

nnx.nn.params.nn_scheduler_params.NNSchedulerParams dataclass

Source code in nnx/nn/params/nn_scheduler_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNSchedulerParams:
    min_lr: float
    factor: float
    patience: int
    cooldown: int
    threshold: float

    # Optional. Default None preserves ReduceLROnPlateau behavior (the only
    # scheduler before the Schedulers enum was added).
    kind: Optional[Schedulers] = None

    # Kind-specific config. Each variant uses a subset; unused fields stay None.
    step_size: Optional[int] = None  # STEP
    T_max: Optional[int] = None  # COSINE_ANNEALING
    max_lr: Optional[float] = None  # ONE_CYCLE
    total_steps: Optional[int] = None  # ONE_CYCLE, LINEAR_WARMUP_DECAY
    warmup_steps: Optional[int] = None  # LINEAR_WARMUP_DECAY

    def __post_init__(self):
        # Fail-fast on out-of-range numeric fields. None of these are emitted
        # into state() when at their defaults, so validation never shifts a
        # run.id — same [[params-boundary-validation]] contract as the other
        # params dataclasses. `factor` is the LR multiplier (the torch
        # ReduceLROnPlateau ctor additionally requires factor < 1, which it
        # enforces itself for that variant); `patience`/`cooldown` are epoch
        # counts; `min_lr`/`threshold` are non-negative bounds. Present
        # (non-None) variant knobs are positive step/length counts.
        if self.factor <= 0:
            raise ValueError(f"NNSchedulerParams requires factor > 0, got {self.factor}")
        if self.min_lr < 0:
            raise ValueError(f"NNSchedulerParams requires min_lr >= 0, got {self.min_lr}")
        if self.threshold < 0:
            raise ValueError(f"NNSchedulerParams requires threshold >= 0, got {self.threshold}")
        if self.patience < 0:
            raise ValueError(f"NNSchedulerParams requires patience >= 0, got {self.patience}")
        if self.cooldown < 0:
            raise ValueError(f"NNSchedulerParams requires cooldown >= 0, got {self.cooldown}")
        for name in ("step_size", "T_max", "max_lr", "total_steps", "warmup_steps"):
            value = getattr(self, name)
            if value is not None and value <= 0:
                raise ValueError(f"NNSchedulerParams requires {name} > 0 when set, got {value}")

    def __str__(self) -> str:
        if self.kind is None:
            return (
                f"[plateau, patience={self.patience}, cooldown={self.cooldown}, "
                f"factor={self.factor:1.0e}, threshold={self.threshold:1.0e}, "
                f"min_lr={self.min_lr:1.0e}]"
            )
        return f"[{self.kind}, factor={self.factor:1.0e}, min_lr={self.min_lr:1.0e}]"

    def state(self) -> dict:
        d = dict(
            min_lr=self.min_lr,
            factor=self.factor,
            cooldown=self.cooldown,
            patience=self.patience,
            threshold=self.threshold,
        )
        # `kind` and its variant-specific knobs (step_size, T_max, max_lr,
        # total_steps, warmup_steps) are omitted from state() when at
        # their defaults so a plain ReduceLROnPlateau NNSchedulerParams
        # hashes to the same run.id as before the Schedulers enum
        # existed. Same omit-when-default invariant as NNTrainParams.seed
        # / NNModelParams.mixed_precision / NNOptimParams.param_groups.
        if self.kind is not None:
            d["kind"] = self.kind.value
        if self.step_size is not None:
            d["step_size"] = self.step_size
        if self.T_max is not None:
            d["T_max"] = self.T_max
        if self.max_lr is not None:
            d["max_lr"] = self.max_lr
        if self.total_steps is not None:
            d["total_steps"] = self.total_steps
        if self.warmup_steps is not None:
            d["warmup_steps"] = self.warmup_steps
        return d

    @staticmethod
    def from_state(state: dict) -> NNSchedulerParams:
        kind_str = state.get("kind")
        return NNSchedulerParams(
            min_lr=state["min_lr"],
            factor=state["factor"],
            patience=state["patience"],
            cooldown=state["cooldown"],
            threshold=state["threshold"],
            kind=Schedulers(kind_str) if kind_str else None,
            step_size=state.get("step_size"),
            T_max=state.get("T_max"),
            max_lr=state.get("max_lr"),
            total_steps=state.get("total_steps"),
            warmup_steps=state.get("warmup_steps"),
        )

    @classmethod
    def builder(cls) -> NNSchedulerParamsBuilder:
        """Return a variant-aware builder. See `NNSchedulerParamsBuilder`."""
        from .nn_scheduler_params_builder import NNSchedulerParamsBuilder

        return NNSchedulerParamsBuilder()

builder() -> NNSchedulerParamsBuilder classmethod

Return a variant-aware builder. See NNSchedulerParamsBuilder.

Source code in nnx/nn/params/nn_scheduler_params.py

@classmethod
def builder(cls) -> NNSchedulerParamsBuilder:
    """Return a variant-aware builder. See `NNSchedulerParamsBuilder`."""
    from .nn_scheduler_params_builder import NNSchedulerParamsBuilder

    return NNSchedulerParamsBuilder()

nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder

Variant-aware builder for NNSchedulerParams.

Reach this via NNSchedulerParams.builder(). Each variant method is self-contained — the user calls exactly one of them per builder instance. Calling a second variant overwrites the first (last write wins); .build() produces the dataclass.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

class NNSchedulerParamsBuilder:
    """Variant-aware builder for `NNSchedulerParams`.

    Reach this via `NNSchedulerParams.builder()`. Each variant method
    is self-contained — the user calls exactly one of them per builder
    instance. Calling a second variant overwrites the first (last
    write wins); `.build()` produces the dataclass.
    """

    def __init__(self) -> None:
        self._fields: dict[str, Any] = {}

    def reduce_on_plateau(
        self,
        *,
        min_lr: float,
        factor: float,
        patience: int,
        cooldown: int,
        threshold: float,
    ) -> NNSchedulerParamsBuilder:
        """ReduceLROnPlateau — the default scheduler.

        Sets the five plateau fields. `kind` is left at None (the
        dataclass default), which preserves the omit-when-default
        state() invariant for callers who used the original pre-enum
        config.
        """
        self._fields = {
            "min_lr": min_lr,
            "factor": factor,
            "patience": patience,
            "cooldown": cooldown,
            "threshold": threshold,
        }
        return self

    def step(
        self,
        *,
        step_size: int,
        min_lr: float,
        factor: float,
        patience: int,
        cooldown: int,
        threshold: float,
    ) -> NNSchedulerParamsBuilder:
        """torch.optim.lr_scheduler.StepLR — decay LR by `factor` every
        `step_size` epochs. The plateau-shape fields (`min_lr`,
        `patience`, `cooldown`, `threshold`) are not consumed by
        StepLR but are required by the underlying NNSchedulerParams
        dataclass and serialised for back-compat.
        """
        self._fields = {
            "kind": Schedulers.STEP,
            "step_size": step_size,
            "min_lr": min_lr,
            "factor": factor,
            "patience": patience,
            "cooldown": cooldown,
            "threshold": threshold,
        }
        return self

    def cosine_annealing(
        self,
        *,
        T_max: int,
        min_lr: float,
        factor: float,
        patience: int,
        cooldown: int,
        threshold: float,
    ) -> NNSchedulerParamsBuilder:
        """torch.optim.lr_scheduler.CosineAnnealingLR — anneal LR over
        `T_max` steps.
        """
        self._fields = {
            "kind": Schedulers.COSINE_ANNEALING,
            "T_max": T_max,
            "min_lr": min_lr,
            "factor": factor,
            "patience": patience,
            "cooldown": cooldown,
            "threshold": threshold,
        }
        return self

    def one_cycle(
        self,
        *,
        max_lr: float,
        total_steps: int,
        min_lr: float,
        factor: float,
        patience: int,
        cooldown: int,
        threshold: float,
    ) -> NNSchedulerParamsBuilder:
        """torch.optim.lr_scheduler.OneCycleLR — Smith one-cycle schedule
        with peak LR `max_lr` over `total_steps` steps.
        """
        self._fields = {
            "kind": Schedulers.ONE_CYCLE,
            "max_lr": max_lr,
            "total_steps": total_steps,
            "min_lr": min_lr,
            "factor": factor,
            "patience": patience,
            "cooldown": cooldown,
            "threshold": threshold,
        }
        return self

    def linear_warmup_decay(
        self,
        *,
        warmup_steps: int,
        total_steps: int,
        min_lr: float,
        factor: float,
        patience: int,
        cooldown: int,
        threshold: float,
    ) -> NNSchedulerParamsBuilder:
        """Linear warm-up to `max_lr` over `warmup_steps`, linear decay
        to 0 over the remaining `total_steps - warmup_steps`. Used by
        most transformer training recipes.
        """
        self._fields = {
            "kind": Schedulers.LINEAR_WARMUP_DECAY,
            "warmup_steps": warmup_steps,
            "total_steps": total_steps,
            "min_lr": min_lr,
            "factor": factor,
            "patience": patience,
            "cooldown": cooldown,
            "threshold": threshold,
        }
        return self

    def build(self) -> NNSchedulerParams:
        """Construct the dataclass from the fields the user touched.

        Pre-empts the dataclass's missing-required-argument TypeError
        with an actionable Builder-level ValueError naming the variant
        methods — matches the [[builder-pattern-shape]] §11b convention
        that PR #52 established on NNTrainerParamsBuilder.

        Forwards only the keys present in `self._fields` so the
        dataclass defaults govern every untouched field — that's what
        preserves the omit-when-default state() invariant.

        Raises:
            ValueError: if no variant method (`.reduce_on_plateau`,
                `.step`, `.cosine_annealing`, `.one_cycle`,
                `.linear_warmup_decay`) was called before `.build()`.
                The message names the five methods so the user can
                fix the chain without consulting the dataclass schema.
        """
        # Every variant fills the 5 plateau-shape fields. An empty
        # _fields means no variant was called.
        if "min_lr" not in self._fields:
            raise ValueError(
                "NNSchedulerParamsBuilder: call one of .reduce_on_plateau(...), "
                ".step(...), .cosine_annealing(...), .one_cycle(...), or "
                ".linear_warmup_decay(...) before .build() — a variant selects "
                "the scheduler kind and sets the required plateau-shape fields."
            )
        return NNSchedulerParams(**self._fields)

build() -> NNSchedulerParams

Construct the dataclass from the fields the user touched.

Pre-empts the dataclass's missing-required-argument TypeError with an actionable Builder-level ValueError naming the variant methods — matches the [[builder-pattern-shape]] §11b convention that PR #52 established on NNTrainerParamsBuilder.

Forwards only the keys present in self._fields so the dataclass defaults govern every untouched field — that's what preserves the omit-when-default state() invariant.

Raises:

Type	Description
ValueError	if no variant method (.reduce_on_plateau, .step, .cosine_annealing, .one_cycle, .linear_warmup_decay) was called before .build(). The message names the five methods so the user can fix the chain without consulting the dataclass schema.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

def build(self) -> NNSchedulerParams:
    """Construct the dataclass from the fields the user touched.

    Pre-empts the dataclass's missing-required-argument TypeError
    with an actionable Builder-level ValueError naming the variant
    methods — matches the [[builder-pattern-shape]] §11b convention
    that PR #52 established on NNTrainerParamsBuilder.

    Forwards only the keys present in `self._fields` so the
    dataclass defaults govern every untouched field — that's what
    preserves the omit-when-default state() invariant.

    Raises:
        ValueError: if no variant method (`.reduce_on_plateau`,
            `.step`, `.cosine_annealing`, `.one_cycle`,
            `.linear_warmup_decay`) was called before `.build()`.
            The message names the five methods so the user can
            fix the chain without consulting the dataclass schema.
    """
    # Every variant fills the 5 plateau-shape fields. An empty
    # _fields means no variant was called.
    if "min_lr" not in self._fields:
        raise ValueError(
            "NNSchedulerParamsBuilder: call one of .reduce_on_plateau(...), "
            ".step(...), .cosine_annealing(...), .one_cycle(...), or "
            ".linear_warmup_decay(...) before .build() — a variant selects "
            "the scheduler kind and sets the required plateau-shape fields."
        )
    return NNSchedulerParams(**self._fields)

cosine_annealing(*, T_max: int, min_lr: float, factor: float, patience: int, cooldown: int, threshold: float) -> NNSchedulerParamsBuilder

torch.optim.lr_scheduler.CosineAnnealingLR — anneal LR over T_max steps.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

def cosine_annealing(
    self,
    *,
    T_max: int,
    min_lr: float,
    factor: float,
    patience: int,
    cooldown: int,
    threshold: float,
) -> NNSchedulerParamsBuilder:
    """torch.optim.lr_scheduler.CosineAnnealingLR — anneal LR over
    `T_max` steps.
    """
    self._fields = {
        "kind": Schedulers.COSINE_ANNEALING,
        "T_max": T_max,
        "min_lr": min_lr,
        "factor": factor,
        "patience": patience,
        "cooldown": cooldown,
        "threshold": threshold,
    }
    return self

linear_warmup_decay(*, warmup_steps: int, total_steps: int, min_lr: float, factor: float, patience: int, cooldown: int, threshold: float) -> NNSchedulerParamsBuilder

Linear warm-up to max_lr over warmup_steps, linear decay to 0 over the remaining total_steps - warmup_steps. Used by most transformer training recipes.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

def linear_warmup_decay(
    self,
    *,
    warmup_steps: int,
    total_steps: int,
    min_lr: float,
    factor: float,
    patience: int,
    cooldown: int,
    threshold: float,
) -> NNSchedulerParamsBuilder:
    """Linear warm-up to `max_lr` over `warmup_steps`, linear decay
    to 0 over the remaining `total_steps - warmup_steps`. Used by
    most transformer training recipes.
    """
    self._fields = {
        "kind": Schedulers.LINEAR_WARMUP_DECAY,
        "warmup_steps": warmup_steps,
        "total_steps": total_steps,
        "min_lr": min_lr,
        "factor": factor,
        "patience": patience,
        "cooldown": cooldown,
        "threshold": threshold,
    }
    return self

one_cycle(*, max_lr: float, total_steps: int, min_lr: float, factor: float, patience: int, cooldown: int, threshold: float) -> NNSchedulerParamsBuilder

torch.optim.lr_scheduler.OneCycleLR — Smith one-cycle schedule with peak LR max_lr over total_steps steps.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

def one_cycle(
    self,
    *,
    max_lr: float,
    total_steps: int,
    min_lr: float,
    factor: float,
    patience: int,
    cooldown: int,
    threshold: float,
) -> NNSchedulerParamsBuilder:
    """torch.optim.lr_scheduler.OneCycleLR — Smith one-cycle schedule
    with peak LR `max_lr` over `total_steps` steps.
    """
    self._fields = {
        "kind": Schedulers.ONE_CYCLE,
        "max_lr": max_lr,
        "total_steps": total_steps,
        "min_lr": min_lr,
        "factor": factor,
        "patience": patience,
        "cooldown": cooldown,
        "threshold": threshold,
    }
    return self

reduce_on_plateau(*, min_lr: float, factor: float, patience: int, cooldown: int, threshold: float) -> NNSchedulerParamsBuilder

ReduceLROnPlateau — the default scheduler.

Sets the five plateau fields. kind is left at None (the dataclass default), which preserves the omit-when-default state() invariant for callers who used the original pre-enum config.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

def reduce_on_plateau(
    self,
    *,
    min_lr: float,
    factor: float,
    patience: int,
    cooldown: int,
    threshold: float,
) -> NNSchedulerParamsBuilder:
    """ReduceLROnPlateau — the default scheduler.

    Sets the five plateau fields. `kind` is left at None (the
    dataclass default), which preserves the omit-when-default
    state() invariant for callers who used the original pre-enum
    config.
    """
    self._fields = {
        "min_lr": min_lr,
        "factor": factor,
        "patience": patience,
        "cooldown": cooldown,
        "threshold": threshold,
    }
    return self

step(*, step_size: int, min_lr: float, factor: float, patience: int, cooldown: int, threshold: float) -> NNSchedulerParamsBuilder

torch.optim.lr_scheduler.StepLR — decay LR by factor every step_size epochs. The plateau-shape fields (min_lr, patience, cooldown, threshold) are not consumed by StepLR but are required by the underlying NNSchedulerParams dataclass and serialised for back-compat.

Source code in nnx/nn/params/nn_scheduler_params_builder.py

def step(
    self,
    *,
    step_size: int,
    min_lr: float,
    factor: float,
    patience: int,
    cooldown: int,
    threshold: float,
) -> NNSchedulerParamsBuilder:
    """torch.optim.lr_scheduler.StepLR — decay LR by `factor` every
    `step_size` epochs. The plateau-shape fields (`min_lr`,
    `patience`, `cooldown`, `threshold`) are not consumed by
    StepLR but are required by the underlying NNSchedulerParams
    dataclass and serialised for back-compat.
    """
    self._fields = {
        "kind": Schedulers.STEP,
        "step_size": step_size,
        "min_lr": min_lr,
        "factor": factor,
        "patience": patience,
        "cooldown": cooldown,
        "threshold": threshold,
    }
    return self

nnx.nn.params.nn_transformer_params.NNTransformerParams dataclass

Bases: NNParams

Source code in nnx/nn/params/nn_transformer_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNTransformerParams(NNParams):
    # Required architectural knobs. There's no meaningful default for
    # "vocabulary size" or "depth", so these are required.
    vocab_size: int
    n_layers: int
    d_model: int
    max_seq_len: int

    # `n_heads` is lifted from NNParams (where it was an Optional[int]
    # used only by GraphAttNN). For TransformerNN it's required; we
    # validate at __post_init__.

    # Defaulted knobs — all of these omit themselves from state() when
    # equal to the default value.
    ffn_mult: int = 4
    rope_base: float = 10000.0
    tie_embeddings: bool = True
    # Separate attention dropout vs. residual dropout — both default to
    # 0.0 (matching modern LLM training where regularization comes from
    # data scale, not dropout). Omitted from state() at default.
    attn_dropout: float = 0.0
    resid_dropout: float = 0.0

    def __post_init__(self):
        # Initialize the base-class _dims slot first. We call the
        # parent's __post_init__ explicitly rather than via super()
        # because `slots=True` dataclass subclassing interacts oddly
        # with the cooperative super() chain (the explicit unbound call
        # is what the dataclasses docs recommend for slotted hierarchies).
        NNParams.__post_init__(self)
        # Transformer-specific validation. n_heads must be set and must
        # divide d_model evenly — checked here so an invalid config
        # fails loudly at params-construction time rather than during
        # the forward pass.
        if self.n_heads is None or self.n_heads <= 0:
            raise ValueError(f"NNTransformerParams requires n_heads > 0, got {self.n_heads!r}")
        # Required positive architectural dimensions. `d_model` must be
        # validated BEFORE the divisibility test below: `d_model=0` passes
        # `0 % n_heads == 0` but then yields `head_dim = d_model // n_heads = 0`
        # (a zero attention-scale divisor and zero-width FFN) downstream — a
        # silent-failure footgun, not a loud one. `n_layers<=0` likewise builds
        # a degenerate embed→norm→head model with no error at all.
        for name in ("vocab_size", "n_layers", "d_model", "max_seq_len", "ffn_mult"):
            value = getattr(self, name)
            if value <= 0:
                raise ValueError(f"NNTransformerParams requires {name} > 0, got {value}")
        # Probability fields validated alongside the base-class `dropout_prob`
        # (checked by NNParams.__post_init__ above) so all three dropout knobs
        # fail fast at construction rather than at the first training forward
        # (F.dropout / scaled_dot_product_attention reject p>1 only when
        # training=True, far from the config's origin).
        for name in ("attn_dropout", "resid_dropout"):
            value = getattr(self, name)
            if not 0.0 <= value <= 1.0:
                raise ValueError(f"NNTransformerParams requires 0.0 <= {name} <= 1.0, got {value}")
        if self.d_model % self.n_heads != 0:
            raise ValueError(f"d_model={self.d_model} must be divisible by n_heads={self.n_heads}")

    def state(self) -> dict:
        # Start with the base-class state (input_dim/output_dim/etc.)
        # so an LM run is still legible the same way every other run
        # is — pretty-printed yaml carries the dims for grep-friendliness.
        # Explicit unbound call: same reason as __post_init__ above —
        # slots=True dataclass subclassing doesn't play nicely with
        # cooperative super().
        d = NNParams.state(self)
        d.update(
            n_layers=self.n_layers,
            d_model=self.d_model,
            vocab_size=self.vocab_size,
            max_seq_len=self.max_seq_len,
        )
        # `n_heads` is always emitted by NNParams.state() — for the
        # transformer path it's required, so it's never None.

        # Omit-when-default invariants (see module docstring). These
        # exist so a "vanilla" TRANSFORMER config hashes to a stable
        # run.id even as we add more knobs over time.
        if self.ffn_mult != 4:
            d["ffn_mult"] = self.ffn_mult
        if self.rope_base != 10000.0:
            d["rope_base"] = self.rope_base
        if self.tie_embeddings is not True:
            d["tie_embeddings"] = self.tie_embeddings
        if self.attn_dropout != 0.0:
            d["attn_dropout"] = self.attn_dropout
        if self.resid_dropout != 0.0:
            d["resid_dropout"] = self.resid_dropout
        return d

    @staticmethod
    def from_state(state: dict) -> NNTransformerParams:
        # `hidden_dims` may be stringified (NNParams convention) or
        # absent for LM configs where it isn't meaningful.
        hidden = state.get("hidden_dims")
        if isinstance(hidden, str):
            hidden = ast.literal_eval(hidden)

        # Base-class fields — fall back to the LM convention where the
        # output dimension equals vocab_size so the NNParams.dims
        # invariant is satisfiable.
        vocab_size = state["vocab_size"]
        input_dim = state.get("input_dim", vocab_size)
        output_dim = state.get("output_dim", vocab_size)

        if "activation" in state and state["activation"] is not None:
            activation = Activations(state["activation"])
        elif "activation" in state:
            # Explicit null — the config was built with activation=None.
            activation = None
        else:
            # Key absent entirely — legacy LM configs omitted it.
            activation = Activations.LEAKY_RELU

        return NNTransformerParams(
            input_dim=input_dim,
            output_dim=output_dim,
            dropout_prob=state.get("dropout_prob", 0.0),
            hidden_dims=hidden,
            activation=activation,
            n_heads=state.get("n_heads"),
            vocab_size=vocab_size,
            n_layers=state["n_layers"],
            d_model=state["d_model"],
            max_seq_len=state["max_seq_len"],
            ffn_mult=state.get("ffn_mult", 4),
            rope_base=state.get("rope_base", 10000.0),
            tie_embeddings=state.get("tie_embeddings", True),
            attn_dropout=state.get("attn_dropout", 0.0),
            resid_dropout=state.get("resid_dropout", 0.0),
        )

    def __str__(self) -> str:
        return (
            f"[transformer, vocab={self.vocab_size}, layers={self.n_layers}, "
            f"d_model={self.d_model}, heads={self.n_heads}, seq={self.max_seq_len}]"
        )

    @classmethod
    def builder(cls) -> NNTransformerParamsBuilder:
        """Return a fluent LM-path builder. See `NNTransformerParamsBuilder`."""
        from .nn_transformer_params_builder import NNTransformerParamsBuilder

        return NNTransformerParamsBuilder()

builder() -> NNTransformerParamsBuilder classmethod

Return a fluent LM-path builder. See NNTransformerParamsBuilder.

Source code in nnx/nn/params/nn_transformer_params.py

@classmethod
def builder(cls) -> NNTransformerParamsBuilder:
    """Return a fluent LM-path builder. See `NNTransformerParamsBuilder`."""
    from .nn_transformer_params_builder import NNTransformerParamsBuilder

    return NNTransformerParamsBuilder()

nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder

Builder for NNTransformerParams.

Reach via NNTransformerParams.builder(). The six methods can be chained in any order; .build() collects them, fills in the LM-path defaults for the dead parent-NNParams fields, and constructs the dataclass.

Source code in nnx/nn/params/nn_transformer_params_builder.py

class NNTransformerParamsBuilder:
    """Builder for `NNTransformerParams`.

    Reach via `NNTransformerParams.builder()`. The six methods can be
    chained in any order; `.build()` collects them, fills in the
    LM-path defaults for the dead parent-NNParams fields, and
    constructs the dataclass.
    """

    def __init__(self) -> None:
        self._fields: dict[str, Any] = {}

    def vocab(self, size: int) -> NNTransformerParamsBuilder:
        """Set the vocabulary size. Mirrors into both `input_dim` and
        `output_dim` on the parent NNParams (the LM convention)."""
        self._fields["vocab_size"] = size
        self._fields["input_dim"] = size
        self._fields["output_dim"] = size
        return self

    def layers(
        self,
        *,
        n: int,
        heads: int,
        d_model: int,
    ) -> NNTransformerParamsBuilder:
        """Set depth (`n_layers`), attention head count (`n_heads`),
        and hidden dimension (`d_model`). Enforces
        `d_model % heads == 0` immediately — this is the Builder's
        safety value-add over the direct-kwarg ctor, which only
        catches the mismatch at __post_init__ time after all kwargs
        have already been typed."""
        if d_model % heads != 0:
            raise ValueError(
                f"d_model={d_model} must be divisible by heads={heads} "
                "(transformer attention requires d_model / n_heads to be integral)"
            )
        self._fields["n_layers"] = n
        self._fields["n_heads"] = heads
        self._fields["d_model"] = d_model
        return self

    def ffn(self, *, mult: int) -> NNTransformerParamsBuilder:
        """FFN expansion ratio. Default is 4 (the SwiGLU-friendly
        ratio); only call this method to override."""
        self._fields["ffn_mult"] = mult
        return self

    def context(
        self,
        *,
        max_seq_len: int,
        rope_base: Optional[float] = None,
    ) -> NNTransformerParamsBuilder:
        """Context-length and RoPE base. `max_seq_len` is required;
        `rope_base=None` is the sentinel for "use the dataclass default
        (10000.0, the LLaMA / GPT convention)". The fluent contract is
        "last call wins": `.context(rope_base=500000.0).context(max_seq_len=128)`
        resets `rope_base` to the default — the second call's implicit
        `rope_base=None` drops the prior override."""
        self._fields["max_seq_len"] = max_seq_len
        if rope_base is None:
            # Drop any prior override so the dataclass default governs
            # at build time; this is what makes last-call-wins work for
            # the None sentinel.
            self._fields.pop("rope_base", None)
        else:
            self._fields["rope_base"] = rope_base
        return self

    def dropout(
        self,
        *,
        attn: float = 0.0,
        resid: float = 0.0,
    ) -> NNTransformerParamsBuilder:
        """Attention and residual dropout rates. Defaults are both
        0.0 (modern LLM convention; regularization comes from data
        scale, not dropout).

        Like `.context()`, a `dropout()` call specifies BOTH rates
        together — each call fully replaces the pair, and a rate left
        at its 0.0 default is reset, not carried over from a prior
        call. So `.dropout(resid=0.3).dropout(attn=0.5)` yields
        `attn=0.5, resid=0.0` (the second call's implicit `resid=0.0`
        drops the prior override); call `.dropout(attn=0.5, resid=0.3)`
        once to set both. Same-field last-call-wins still holds:
        `.dropout(attn=0.5).dropout(attn=0.0)` resets to `attn=0.0`.
        The dataclass's omit-when-default `state()` then handles
        run.id stability automatically."""
        self._fields["attn_dropout"] = attn
        self._fields["resid_dropout"] = resid
        return self

    def tied_embeddings(self, value: bool) -> NNTransformerParamsBuilder:
        """Toggle weight-tying between input embeddings and LM head.
        Default is True. The fluent contract is "last call wins" — a
        prior `.tied_embeddings(False)` followed by `.tied_embeddings(True)`
        leaves the dataclass at the default (which `state()` then omits)."""
        self._fields["tie_embeddings"] = value
        return self

    def build(self) -> NNTransformerParams:
        """Construct the dataclass.

        Pre-empts the dataclass's missing-required-argument TypeError
        with an actionable Builder-level ValueError naming the setter
        methods that haven't been called yet — matches the
        [[builder-pattern-shape]] §11b convention that PR #52
        established on NNTrainerParamsBuilder.

        Fills in the dead parent-NNParams fields the TransformerNN
        net never reads but the parent dataclass requires at
        construction. `activation` mirrors the parent NNParams's
        default (`Activations.LEAKY_RELU`); a Builder-default
        mismatch here previously produced a different `state()` /
        `run.id` than the direct-kwarg ctor.

        Raises:
            ValueError: if `.vocab(size=...)`, `.layers(n=..., heads=...,
                d_model=...)`, or `.context(max_seq_len=...)` was not
                called before `.build()`. The message names the
                specific setter methods that are still missing so the
                user can complete the chain without consulting the
                dataclass schema.
        """
        missing: list[str] = []
        if "vocab_size" not in self._fields:
            missing.append(".vocab(size=...)")
        if "n_layers" not in self._fields:
            missing.append(".layers(n=..., heads=..., d_model=...)")
        if "max_seq_len" not in self._fields:
            missing.append(".context(max_seq_len=...)")
        if missing:
            raise ValueError(
                "NNTransformerParamsBuilder: call "
                + ", ".join(missing)
                + " before .build() — each setter fills the dataclass's "
                "required-no-default fields for the LM path."
            )
        return NNTransformerParams(
            hidden_dims=None,
            dropout_prob=0.0,
            activation=Activations.LEAKY_RELU,
            **self._fields,
        )

build() -> NNTransformerParams

Construct the dataclass.

Pre-empts the dataclass's missing-required-argument TypeError with an actionable Builder-level ValueError naming the setter methods that haven't been called yet — matches the [[builder-pattern-shape]] §11b convention that PR #52 established on NNTrainerParamsBuilder.

Fills in the dead parent-NNParams fields the TransformerNN net never reads but the parent dataclass requires at construction. activation mirrors the parent NNParams's default (Activations.LEAKY_RELU); a Builder-default mismatch here previously produced a different state() / run.id than the direct-kwarg ctor.

Raises:

Type	Description
ValueError	if .vocab(size=...), .layers(n=..., heads=..., d_model=...), or .context(max_seq_len=...) was not called before .build(). The message names the specific setter methods that are still missing so the user can complete the chain without consulting the dataclass schema.

Source code in nnx/nn/params/nn_transformer_params_builder.py

def build(self) -> NNTransformerParams:
    """Construct the dataclass.

    Pre-empts the dataclass's missing-required-argument TypeError
    with an actionable Builder-level ValueError naming the setter
    methods that haven't been called yet — matches the
    [[builder-pattern-shape]] §11b convention that PR #52
    established on NNTrainerParamsBuilder.

    Fills in the dead parent-NNParams fields the TransformerNN
    net never reads but the parent dataclass requires at
    construction. `activation` mirrors the parent NNParams's
    default (`Activations.LEAKY_RELU`); a Builder-default
    mismatch here previously produced a different `state()` /
    `run.id` than the direct-kwarg ctor.

    Raises:
        ValueError: if `.vocab(size=...)`, `.layers(n=..., heads=...,
            d_model=...)`, or `.context(max_seq_len=...)` was not
            called before `.build()`. The message names the
            specific setter methods that are still missing so the
            user can complete the chain without consulting the
            dataclass schema.
    """
    missing: list[str] = []
    if "vocab_size" not in self._fields:
        missing.append(".vocab(size=...)")
    if "n_layers" not in self._fields:
        missing.append(".layers(n=..., heads=..., d_model=...)")
    if "max_seq_len" not in self._fields:
        missing.append(".context(max_seq_len=...)")
    if missing:
        raise ValueError(
            "NNTransformerParamsBuilder: call "
            + ", ".join(missing)
            + " before .build() — each setter fills the dataclass's "
            "required-no-default fields for the LM path."
        )
    return NNTransformerParams(
        hidden_dims=None,
        dropout_prob=0.0,
        activation=Activations.LEAKY_RELU,
        **self._fields,
    )

context(*, max_seq_len: int, rope_base: Optional[float] = None) -> NNTransformerParamsBuilder

Context-length and RoPE base. max_seq_len is required; rope_base=None is the sentinel for "use the dataclass default (10000.0, the LLaMA / GPT convention)". The fluent contract is "last call wins": .context(rope_base=500000.0).context(max_seq_len=128) resets rope_base to the default — the second call's implicit rope_base=None drops the prior override.

Source code in nnx/nn/params/nn_transformer_params_builder.py

def context(
    self,
    *,
    max_seq_len: int,
    rope_base: Optional[float] = None,
) -> NNTransformerParamsBuilder:
    """Context-length and RoPE base. `max_seq_len` is required;
    `rope_base=None` is the sentinel for "use the dataclass default
    (10000.0, the LLaMA / GPT convention)". The fluent contract is
    "last call wins": `.context(rope_base=500000.0).context(max_seq_len=128)`
    resets `rope_base` to the default — the second call's implicit
    `rope_base=None` drops the prior override."""
    self._fields["max_seq_len"] = max_seq_len
    if rope_base is None:
        # Drop any prior override so the dataclass default governs
        # at build time; this is what makes last-call-wins work for
        # the None sentinel.
        self._fields.pop("rope_base", None)
    else:
        self._fields["rope_base"] = rope_base
    return self

dropout(*, attn: float = 0.0, resid: float = 0.0) -> NNTransformerParamsBuilder

Attention and residual dropout rates. Defaults are both 0.0 (modern LLM convention; regularization comes from data scale, not dropout).

Like .context(), a dropout() call specifies BOTH rates together — each call fully replaces the pair, and a rate left at its 0.0 default is reset, not carried over from a prior call. So .dropout(resid=0.3).dropout(attn=0.5) yields attn=0.5, resid=0.0 (the second call's implicit resid=0.0 drops the prior override); call .dropout(attn=0.5, resid=0.3) once to set both. Same-field last-call-wins still holds: .dropout(attn=0.5).dropout(attn=0.0) resets to attn=0.0. The dataclass's omit-when-default state() then handles run.id stability automatically.

Source code in nnx/nn/params/nn_transformer_params_builder.py

def dropout(
    self,
    *,
    attn: float = 0.0,
    resid: float = 0.0,
) -> NNTransformerParamsBuilder:
    """Attention and residual dropout rates. Defaults are both
    0.0 (modern LLM convention; regularization comes from data
    scale, not dropout).

    Like `.context()`, a `dropout()` call specifies BOTH rates
    together — each call fully replaces the pair, and a rate left
    at its 0.0 default is reset, not carried over from a prior
    call. So `.dropout(resid=0.3).dropout(attn=0.5)` yields
    `attn=0.5, resid=0.0` (the second call's implicit `resid=0.0`
    drops the prior override); call `.dropout(attn=0.5, resid=0.3)`
    once to set both. Same-field last-call-wins still holds:
    `.dropout(attn=0.5).dropout(attn=0.0)` resets to `attn=0.0`.
    The dataclass's omit-when-default `state()` then handles
    run.id stability automatically."""
    self._fields["attn_dropout"] = attn
    self._fields["resid_dropout"] = resid
    return self

ffn(*, mult: int) -> NNTransformerParamsBuilder

FFN expansion ratio. Default is 4 (the SwiGLU-friendly ratio); only call this method to override.

Source code in nnx/nn/params/nn_transformer_params_builder.py

def ffn(self, *, mult: int) -> NNTransformerParamsBuilder:
    """FFN expansion ratio. Default is 4 (the SwiGLU-friendly
    ratio); only call this method to override."""
    self._fields["ffn_mult"] = mult
    return self

layers(*, n: int, heads: int, d_model: int) -> NNTransformerParamsBuilder

Set depth (n_layers), attention head count (n_heads), and hidden dimension (d_model). Enforces d_model % heads == 0 immediately — this is the Builder's safety value-add over the direct-kwarg ctor, which only catches the mismatch at post_init time after all kwargs have already been typed.

Source code in nnx/nn/params/nn_transformer_params_builder.py

def layers(
    self,
    *,
    n: int,
    heads: int,
    d_model: int,
) -> NNTransformerParamsBuilder:
    """Set depth (`n_layers`), attention head count (`n_heads`),
    and hidden dimension (`d_model`). Enforces
    `d_model % heads == 0` immediately — this is the Builder's
    safety value-add over the direct-kwarg ctor, which only
    catches the mismatch at __post_init__ time after all kwargs
    have already been typed."""
    if d_model % heads != 0:
        raise ValueError(
            f"d_model={d_model} must be divisible by heads={heads} "
            "(transformer attention requires d_model / n_heads to be integral)"
        )
    self._fields["n_layers"] = n
    self._fields["n_heads"] = heads
    self._fields["d_model"] = d_model
    return self

tied_embeddings(value: bool) -> NNTransformerParamsBuilder

Toggle weight-tying between input embeddings and LM head. Default is True. The fluent contract is "last call wins" — a prior .tied_embeddings(False) followed by .tied_embeddings(True) leaves the dataclass at the default (which state() then omits).

Source code in nnx/nn/params/nn_transformer_params_builder.py

def tied_embeddings(self, value: bool) -> NNTransformerParamsBuilder:
    """Toggle weight-tying between input embeddings and LM head.
    Default is True. The fluent contract is "last call wins" — a
    prior `.tied_embeddings(False)` followed by `.tied_embeddings(True)`
    leaves the dataclass at the default (which `state()` then omits)."""
    self._fields["tie_embeddings"] = value
    return self

vocab(size: int) -> NNTransformerParamsBuilder

Set the vocabulary size. Mirrors into both input_dim and output_dim on the parent NNParams (the LM convention).

Source code in nnx/nn/params/nn_transformer_params_builder.py

def vocab(self, size: int) -> NNTransformerParamsBuilder:
    """Set the vocabulary size. Mirrors into both `input_dim` and
    `output_dim` on the parent NNParams (the LM convention)."""
    self._fields["vocab_size"] = size
    self._fields["input_dim"] = size
    self._fields["output_dim"] = size
    return self

nnx.nn.params.nn_tokenizer_params.NNTokenizerParams dataclass

Frozen dataclass holding a tokenizer + its on-disk pointer.

The dataclass is frozen so it can sit alongside NNTransformerParams / NNModelParams in an NNRun without inviting in-place mutation. The actual tokenizers.Tokenizer object is held in a repr=False field so it doesn't bloat the str() output.

Source code in nnx/nn/params/nn_tokenizer_params.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNTokenizerParams:
    """Frozen dataclass holding a tokenizer + its on-disk pointer.

    The dataclass is frozen so it can sit alongside NNTransformerParams /
    NNModelParams in an NNRun without inviting in-place mutation. The
    actual ``tokenizers.Tokenizer`` object is held in a repr=False field
    so it doesn't bloat the str() output.
    """

    path: str
    tokenizer: object = field(repr=False)

    # ---------- factories ----------

    @staticmethod
    def of(tokenizer: object, path: str) -> NNTokenizerParams:
        """Construct from a live Tokenizer instance and persist it to ``path``.

        This is the train-time entry point: train a tokenizer, then call
        ``NNTokenizerParams.of(tk, path="runs/tok.json")`` to wrap it
        with a paired on-disk artifact.
        """
        _require_tokenizers()
        # save() is the official HF Rust path — produces the same JSON
        # blob that ``Tokenizer.from_file`` consumes. Write to a temp
        # name and os.replace into place so an interrupt mid-save can't
        # leave a truncated tokenizer.json that from_state can never
        # load (matches the atomic-write convention of NNRun.save /
        # NNCheckpoint.to_file).
        # Parent dir first: the Rust-side save otherwise fails from a
        # fresh cwd with a cryptic "No such file or directory (os error
        # 2)" for paths like "artifacts/tok.json" (parity with
        # write_gguf / export_ollama_modelfile).
        parent = os.path.dirname(path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        tmp = path + ".tmp"
        tokenizer.save(tmp)
        os.replace(tmp, path)
        return NNTokenizerParams(path=path, tokenizer=tokenizer)

    @staticmethod
    def from_state(state: dict) -> NNTokenizerParams:
        """Load from a state dict produced by :meth:`state`. The single
        required key is ``path``; the tokenizer is reconstructed from
        the file the path points to.

        The path is stored exactly as the caller gave it to :meth:`of` —
        typically cwd-relative — so loading from a different working
        directory requires the same relative layout. That's deliberate:
        storing an absolute path would break run portability across
        machines, which is the more common need."""
        _require_tokenizers()
        path = state["path"]
        tk = Tokenizer.from_file(path)  # type: ignore[union-attr]
        return NNTokenizerParams(path=path, tokenizer=tk)

    # ---------- (de)serialization ----------

    def state(self) -> dict:
        """Return the serializable view — only the path goes into run.yaml."""
        return {"path": self.path}

    # ---------- ergonomic helpers ----------

    @property
    def vocab_size(self) -> int:
        return int(self.tokenizer.get_vocab_size())  # type: ignore[union-attr]

    def encode(self, text: str) -> list[int]:
        enc = self.tokenizer.encode(text)  # type: ignore[union-attr]
        return list(enc.ids)

    def decode(self, ids: list[int], skip_special_tokens: bool = True) -> str:
        return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)  # type: ignore[union-attr]

from_state(state: dict) -> NNTokenizerParams staticmethod

Load from a state dict produced by :meth:state. The single required key is path; the tokenizer is reconstructed from the file the path points to.

The path is stored exactly as the caller gave it to :meth:of — typically cwd-relative — so loading from a different working directory requires the same relative layout. That's deliberate: storing an absolute path would break run portability across machines, which is the more common need.

Source code in nnx/nn/params/nn_tokenizer_params.py

@staticmethod
def from_state(state: dict) -> NNTokenizerParams:
    """Load from a state dict produced by :meth:`state`. The single
    required key is ``path``; the tokenizer is reconstructed from
    the file the path points to.

    The path is stored exactly as the caller gave it to :meth:`of` —
    typically cwd-relative — so loading from a different working
    directory requires the same relative layout. That's deliberate:
    storing an absolute path would break run portability across
    machines, which is the more common need."""
    _require_tokenizers()
    path = state["path"]
    tk = Tokenizer.from_file(path)  # type: ignore[union-attr]
    return NNTokenizerParams(path=path, tokenizer=tk)

of(tokenizer: object, path: str) -> NNTokenizerParams staticmethod

Construct from a live Tokenizer instance and persist it to path.

This is the train-time entry point: train a tokenizer, then call NNTokenizerParams.of(tk, path="runs/tok.json") to wrap it with a paired on-disk artifact.

Source code in nnx/nn/params/nn_tokenizer_params.py

@staticmethod
def of(tokenizer: object, path: str) -> NNTokenizerParams:
    """Construct from a live Tokenizer instance and persist it to ``path``.

    This is the train-time entry point: train a tokenizer, then call
    ``NNTokenizerParams.of(tk, path="runs/tok.json")`` to wrap it
    with a paired on-disk artifact.
    """
    _require_tokenizers()
    # save() is the official HF Rust path — produces the same JSON
    # blob that ``Tokenizer.from_file`` consumes. Write to a temp
    # name and os.replace into place so an interrupt mid-save can't
    # leave a truncated tokenizer.json that from_state can never
    # load (matches the atomic-write convention of NNRun.save /
    # NNCheckpoint.to_file).
    # Parent dir first: the Rust-side save otherwise fails from a
    # fresh cwd with a cryptic "No such file or directory (os error
    # 2)" for paths like "artifacts/tok.json" (parity with
    # write_gguf / export_ollama_modelfile).
    parent = os.path.dirname(path)
    if parent:
        os.makedirs(parent, exist_ok=True)
    tmp = path + ".tmp"
    tokenizer.save(tmp)
    os.replace(tmp, path)
    return NNTokenizerParams(path=path, tokenizer=tokenizer)

state() -> dict

Return the serializable view — only the path goes into run.yaml.

Source code in nnx/nn/params/nn_tokenizer_params.py

def state(self) -> dict:
    """Return the serializable view — only the path goes into run.yaml."""
    return {"path": self.path}

nnx.nn.params.nn_tokenizer_params.train_bpe(files: Optional[list[str]] = None, *, vocab_size: int = 8192, texts: Optional[list[str]] = None, special_tokens: Optional[list[str]] = None, min_frequency: int = 2) -> Tokenizer

Train a BPE tokenizer on either a list of files or a list of texts.

Mirrors the HF "quick BPE" recipe — Whitespace pre-tokenizer + BPE model + BpeTrainer. Returns the trained Tokenizer instance; the caller is responsible for persisting via NNTokenizerParams.of(tk, path=...).

Parameters:

Name	Type	Description	Default
files	Optional[list[str]]	paths to plaintext files (one corpus line per file row). If None, texts is consulted instead.	None
vocab_size	int	target vocab. Actual size may be smaller for tiny corpora.	8192
texts	Optional[list[str]]	in-memory list of training strings — useful for unit tests and the examples without writing a temp file.	None
special_tokens	Optional[list[str]]	e.g. ["<pad>", "<bos>", "<eos>"]. Included in the vocab and not split during tokenization.	None
min_frequency	int	minimum pair frequency to merge — higher values give smaller, more conservative vocabs.	2

Returns:

Name	Type	Description
Tokenizer	Tokenizer	a trained tokenizers.Tokenizer ready for encode/decode + save.

Source code in nnx/nn/params/nn_tokenizer_params.py

def train_bpe(
    files: Optional[list[str]] = None,
    *,
    vocab_size: int = 8192,
    texts: Optional[list[str]] = None,
    special_tokens: Optional[list[str]] = None,
    min_frequency: int = 2,
) -> Tokenizer:  # bound to `None` when the optional `lm` extra isn't installed; annotation only evaluated lazily under `from __future__ import annotations`
    """Train a BPE tokenizer on either a list of files or a list of texts.

    Mirrors the HF "quick BPE" recipe — Whitespace pre-tokenizer + BPE
    model + BpeTrainer. Returns the trained Tokenizer instance; the
    caller is responsible for persisting via
    ``NNTokenizerParams.of(tk, path=...)``.

    Args:
        files: paths to plaintext files (one corpus line per file row).
            If None, ``texts`` is consulted instead.
        vocab_size: target vocab. Actual size may be smaller for tiny
            corpora.
        texts: in-memory list of training strings — useful for unit
            tests and the examples without writing a temp file.
        special_tokens: e.g. ``["<pad>", "<bos>", "<eos>"]``. Included
            in the vocab and not split during tokenization.
        min_frequency: minimum pair frequency to merge — higher values
            give smaller, more conservative vocabs.

    Returns:
        Tokenizer: a trained ``tokenizers.Tokenizer`` ready for encode/decode + save.
    """
    _require_tokenizers()
    if files is None and texts is None:
        raise ValueError("Must provide either `files` or `texts`.")

    tk = Tokenizer(BPE(unk_token="<unk>"))  # type: ignore[union-attr]
    tk.pre_tokenizer = Whitespace()  # type: ignore[union-attr]
    trainer = BpeTrainer(  # type: ignore[union-attr]
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=list(special_tokens or ["<unk>", "<pad>", "<bos>", "<eos>"]),
    )

    if files is not None:
        tk.train(files=list(files), trainer=trainer)
    else:
        tk.train_from_iterator(iter(texts), trainer=trainer, length=len(texts))

    return tk

nnx.nn.params.nn_run.NNRun dataclass

Source code in nnx/nn/params/nn_run.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNRun:
    net: NNParams
    train: NNTrainParams
    model: NNModelParams

    # Optional trainer-mode marker. Populated by Trainer.train(); None for
    # NNModel.train()-produced runs. state() omits it when None so existing
    # run.id hashes for NNModel runs are preserved exactly.
    trainer: Optional[NNTrainerParams] = field(default=None)

    _id: Optional[str] = field(repr=False, default=None)
    _state: Optional[dict] = field(repr=False, default=None)
    idps: Optional[list[NNIterationDataPoint]] = field(repr=False, default=None)

    def __str__(self):
        # Delegate to NNSchedulerParams.__str__ for the scheduler block —
        # it knows which fields apply to the configured `kind` (the
        # plateau-only `patience`/`cooldown`/`threshold` are misleading
        # for cosine/onecycle/linear-warmup schedulers).
        return (
            "{"
            f"loss={self.model.loss}"
            f", device={self.model.device}"
            f", net={self.model.net}"
            f", dims={self.net.dims}"
            f", dropout={self.net.dropout_prob}"
            f", activation={self.net.activation}"
            f", n_heads={self.net.n_heads}"
            f", n_epochs={self.train.n_epochs}"
            f", max_lr={self.train.optim.max_lr}"
            f", momentum={self.train.optim.momentum}"
            f", decay={self.train.optim.weight_decay}"
            f", scheduler={self.train.scheduler}"
            "}"
        )

    @property
    def id(self) -> str:
        return self._id

    def __post_init__(self):
        state = dict(model=self.model.state(), net=self.net.state(), train=self.train.state())
        # `trainer` is omitted when None so existing NNModel runs hash to
        # the same run.id as before this field existed. Same omit-when-
        # default pattern as NNTrainParams.seed / save_phase_checkpoints
        # and NNOptimParams.param_groups.
        if self.trainer is not None:
            state["trainer"] = self.trainer.state()

        id = hashlib.md5(str(state).encode("utf-8")).hexdigest()

        object.__setattr__(self, "_id", id)
        object.__setattr__(self, "_state", {"id": id, **state})

    def state(self) -> dict:
        return self._state

    def _repr_html_(self) -> str:
        """Jupyter rich-display: config table + per-epoch metric chart.

        Returns the HTML Jupyter will render when this :class:`NNRun`
        is the last expression in a cell. Outside Jupyter the same
        method can be called directly to grab the HTML string.

        The config table reflects the same state the run.id is hashed
        from. The metric chart plots train/val loss + error across
        epochs when ``self.idps`` is populated; otherwise the chart is
        omitted and only the config table appears.
        """
        config_html = self._render_config_table_html()
        # Empty list and None both collapse to "no chart" — only render
        # when self.idps actually has at least one idp to chart from.
        # The explicit `is not None and len()` form (vs truthy `if self.idps`)
        # narrows the type for static checkers and makes the empty-list
        # case visible to readers.
        chart_html = self._render_metric_chart_html() if self.idps is not None and len(self.idps) > 0 else ""
        return f'<div style="font-family: sans-serif;">{config_html}{chart_html}</div>'

    def _render_config_table_html(self) -> str:
        """HTML table of the canonical run config (subset of state())."""
        rows = [
            ("run.id", self.id),
            ("net", str(self.model.net)),
            ("device", str(self.model.device)),
            ("loss", str(self.model.loss)),
            ("input_dim → output_dim", f"{self.net.input_dim} → {self.net.output_dim}"),
            ("hidden_dims", str(self.net.hidden_dims)),
            ("dropout", str(self.net.dropout_prob)),
            ("activation", str(self.net.activation)),
            ("n_epochs", str(self.train.n_epochs)),
            ("optim", f"{self.train.optim.name} (max_lr={self.train.optim.max_lr})"),
        ]
        rows_html = "".join(
            f'<tr><td style="padding:2px 8px;font-weight:600;">{k}</td>'
            f'<td style="padding:2px 8px;font-family:monospace;">{v}</td></tr>'
            for k, v in rows
        )
        return f'<table style="border-collapse:collapse;border:1px solid #ddd;margin-bottom:8px;">{rows_html}</table>'

    def _render_metric_chart_html(self) -> str:
        """Plotly per-epoch metric chart embedded as HTML."""
        # Lazy-import plotly so test collection stays fast and
        # non-Jupyter callers who never trigger _repr_html_ don't pay
        # the import cost.
        # Group idps by epoch_idx. Each idp carries its epoch index
        # directly; the last idp per epoch may also carry val_edp.
        from collections import defaultdict

        import plotly.graph_objects as go

        epoch_buckets: dict[int, list] = defaultdict(list)
        for idp in self.idps:
            epoch_buckets[idp.epoch_idx].append(idp)

        epochs: list[int] = sorted(epoch_buckets.keys())
        if not epochs:
            return ""  # No data at all.

        train_losses: list[float] = []
        train_errs: list[float] = []
        val_losses: list[float] = []
        val_errs: list[float] = []

        for epoch_idx in epochs:
            idp_list = epoch_buckets[epoch_idx]
            losses = [idp.train_edp.loss for idp in idp_list if idp.train_edp.loss is not None]
            errs = [idp.train_edp.error for idp in idp_list if idp.train_edp.error is not None]
            train_losses.append(sum(losses) / len(losses) if losses else float("nan"))
            train_errs.append(sum(errs) / len(errs) if errs else float("nan"))
            # val_edp is set only on the last idp of each epoch (when a
            # val_loader was supplied). Use the last non-None val_edp found.
            val_idp = next((idp for idp in reversed(idp_list) if idp.val_edp is not None), None)
            val_losses.append(val_idp.val_edp.loss if val_idp and val_idp.val_edp.loss is not None else float("nan"))
            val_errs.append(val_idp.val_edp.error if val_idp and val_idp.val_edp.error is not None else float("nan"))

        fig = go.Figure()
        fig.add_trace(go.Scatter(x=epochs, y=train_losses, name="train_loss", mode="lines+markers"))
        # Only add val traces when at least one epoch had a val_edp.
        if any(v == v for v in val_losses):  # any non-NaN
            fig.add_trace(go.Scatter(x=epochs, y=val_losses, name="val_loss", mode="lines+markers"))
        if any(v == v for v in train_errs):
            fig.add_trace(go.Scatter(x=epochs, y=train_errs, name="train_err", mode="lines+markers", yaxis="y2"))
        if any(v == v for v in val_errs):
            fig.add_trace(go.Scatter(x=epochs, y=val_errs, name="val_err", mode="lines+markers", yaxis="y2"))
        fig.update_layout(
            title=f"NNRun {self.id[:8]}… — {len(epochs)} epoch{'s' if len(epochs) != 1 else ''}",
            xaxis_title="Epoch",
            yaxis=dict(title="Loss"),
            yaxis2=dict(title="Error", overlaying="y", side="right"),
            height=350,
            margin=dict(t=40, b=40, l=40, r=40),
        )
        return fig.to_html(full_html=False, include_plotlyjs="cdn")

    def with_idps(self, value: list[NNIterationDataPoint]) -> NNRun:
        return replace(self, idps=value)

    def checkpoints(self, root: Optional[str] = None) -> list[Optional[NNCheckpoint]]:
        """Load this run's five phase checkpoints, in cadence order
        (FIRST, Q1, Q2, Q3, LAST). Entries are None when the tag was
        never written — e.g. runs trained with
        ``save_phase_checkpoints=False`` write only LAST and BEST.

        BEST is deliberately excluded: it duplicates whichever phase
        checkpoint won, so including it would double-count. Load it
        directly via ``NNCheckpoint.load(run=run.id,
        type=Checkpoints.BEST)``.
        """
        return [
            NNCheckpoint.load(run=self.id, type=Checkpoints.FIRST, root=root),
            NNCheckpoint.load(run=self.id, type=Checkpoints.Q1, root=root),
            NNCheckpoint.load(run=self.id, type=Checkpoints.Q2, root=root),
            NNCheckpoint.load(run=self.id, type=Checkpoints.Q3, root=root),
            NNCheckpoint.load(run=self.id, type=Checkpoints.LAST, root=root),
        ]

    def save(self, root: Optional[str] = None) -> NNRun:
        runs_root = _runs_root(root)
        run_path = os.path.join(runs_root, self.id)
        best_run_path = os.path.join(runs_root, "best")

        csv_path = os.path.join(run_path, "idps.csv")
        yaml_path = os.path.join(run_path, "run.yaml")
        metadata_path = os.path.join(run_path, "metadata.yaml")

        if not os.path.exists(run_path):
            os.makedirs(run_path)

        # All three writes go through the atomic write helper so a
        # KeyboardInterrupt mid-save leaves either the old file or the
        # new file, never a half-written one. This is what makes the
        # per-epoch incremental save actually safe — without atomicity,
        # a partial write here corrupts the run.
        # safe_dump — never plain `yaml.dump`. NNRun.state() is a plain
        # dict of primitive types (the round-trip contract — see the
        # corresponding safe_load call in NNRun.load); using safe_dump
        # here makes that contract enforced at write-time too, so a
        # future state() change that smuggles in a non-primitive (e.g.
        # a torch.dtype or a numpy scalar) fails loudly here rather
        # than producing a run.yaml that fails to safe_load.
        # sort_keys=True is explicit so the on-disk YAML stays stable
        # across PyYAML versions (the default was False prior to 5.1
        # and True from 5.1 onward; we pin to the post-5.1 alphabetical
        # ordering so a future major-version bump can't silently shift
        # the file shape downstream tooling diffs against).
        _atomic_write_text(yaml_path, yaml.safe_dump(self.state(), sort_keys=True))

        # Env snapshot: written separately so it does NOT contribute to
        # run.id (which is md5(state())). Captures library/torch/python
        # versions + git commit so a run.yaml from six months ago is
        # debuggable even if the library has moved on.
        from ...seeding import env_snapshot

        _atomic_write_text(metadata_path, yaml.safe_dump(env_snapshot(), sort_keys=True))

        _atomic_write_text(
            csv_path,
            # `or []`: idps defaults to None on the dataclass; an empty
            # frame round-trips cleanly through read_csv on load.
            pd.json_normalize(data=[idp.state() for idp in (self.idps or [])]).to_csv(),
        )

        if not os.path.lexists(best_run_path) or not os.path.exists(best_run_path):
            # Either no symlink yet, or one dangling after a repo move — repoint.
            _point_best(best_run_path, run_path)
        else:
            # Resolve the current best target via the symlink OR pointer file —
            # `NNCheckpoint.load(run="best", ...)` only works under a symlink
            # layout, so on the Windows pointer-file fallback we go through
            # the run id explicitly to make a fair comparison.
            best_run_id = _read_best_pointer(best_run_path)
            best_ckpt = (
                NNCheckpoint.load(run=best_run_id, type=Checkpoints.BEST, root=root)
                if best_run_id is not None
                else None
            )
            best_err = _best_err(best_ckpt)
            curr_err = _best_err(NNCheckpoint.load(run=self.id, type=Checkpoints.BEST, root=root))

            if curr_err < best_err:
                _point_best(best_run_path, run_path)

        return self

    @staticmethod
    def load(id: str, root: Optional[str] = None) -> NNRun:
        # Reject path-traversal identifiers before joining; see
        # `_validate_run_id` for the threat model. Internal callers pass
        # md5 hex (always safe), so this is a defense-in-depth guard on
        # the public API surface.
        id = _validate_run_id(id)
        run_path = os.path.join(_runs_root(root), id)
        yaml_path = os.path.join(run_path, "run.yaml")
        csv_path = os.path.join(run_path, "idps.csv")

        with open(yaml_path, encoding="utf-8") as f:
            # safe_load — never FullLoader. NNRun.state() is a plain dict
            # of primitive types (strings / ints / floats / lists / nested
            # dicts), so safe_load round-trips it losslessly. Using
            # FullLoader would let a tampered or attacker-supplied
            # run.yaml instantiate arbitrary Python objects via the
            # `!!python/object/...` tag — defense-in-depth even though
            # the file is normally application-written. Matches the
            # safe_dump/safe_load pair the sibling metadata.yaml has
            # always used (see seeding.py's "yaml.safe_load-compatible"
            # comment for the metadata round-trip).
            rep = yaml.safe_load(f)
        if not isinstance(rep, dict):
            # An empty / truncated-to-zero file safe_loads to None and
            # would otherwise die on rep.get(...) with no file context.
            raise ValueError(f"malformed run.yaml at {yaml_path}: expected a mapping, got {type(rep).__name__}")

        try:
            raw_idps = pd.read_csv(csv_path).to_dict(orient="records")
        except pd.errors.EmptyDataError as e:
            # A zero-byte file (external truncation — our own atomic
            # writes never produce one; even an idps-less run writes the
            # frame header) raises a pandas error with no file context.
            raise ValueError(f"malformed idps.csv at {csv_path}: {e}") from e
        # Separate try-scopes per source file so a missing key names the
        # FILE that's actually corrupt (a dropped CSV column must not be
        # reported as a run.yaml problem, and vice versa).
        try:
            idps = [NNIterationDataPoint.from_state(idp) for idp in raw_idps]
        except KeyError as e:
            raise ValueError(f"malformed idps.csv at {csv_path}: missing column/key {e}") from e

        try:
            # Lazy import for trainer params — keeps `nnx.nn.params`
            # importable without dragging the trainer subpackage in, and
            # avoids a cycle if anything in `nnx.trainer` ever needs to
            # import NNRun.
            trainer_state = rep.get("trainer")
            if trainer_state is not None:
                from ...trainer.params import NNTrainerParams

                trainer = NNTrainerParams.from_state(trainer_state)
            else:
                trainer = None

            return NNRun(
                # resolve_from_state: a TRANSFORMER run's net params must
                # come back as NNTransformerParams, not be downgraded to
                # NNParams.
                net=NNParams.resolve_from_state(rep["net"]),
                train=NNTrainParams.from_state(rep["train"]),
                model=NNModelParams.from_state(rep["model"]),
                trainer=trainer,
                idps=idps,
            )
        except KeyError as e:
            # A hand-edited / truncated run.yaml otherwise surfaces as a
            # bare KeyError with no file context.
            raise ValueError(f"malformed run.yaml at {yaml_path}: missing key {e}") from e

    @staticmethod
    def all(root: Optional[str] = None) -> list[NNRun]:
        """List every saved NNRun under the runs root, skipping the `best`
        pointer. Returns [] when the runs/ directory doesn't exist yet.
        Non-directory entries (stray files, .DS_Store) are filtered out
        so they don't trigger spurious NNRun.load failures."""
        runs_root = _runs_root(root)
        if not os.path.isdir(runs_root):
            return []
        result: list[NNRun] = []
        for entry in os.listdir(runs_root):
            if entry == "best":
                continue
            if not os.path.isdir(os.path.join(runs_root, entry)):
                continue
            # Defensive: a directory in runs/ that lacks run.yaml isn't a
            # real run (could be a leftover from an aborted experiment).
            if not os.path.isfile(os.path.join(runs_root, entry, "run.yaml")):
                continue
            result.append(NNRun.load(id=entry, root=root))
        return result

all(root: Optional[str] = None) -> list[NNRun] staticmethod

List every saved NNRun under the runs root, skipping the best pointer. Returns [] when the runs/ directory doesn't exist yet. Non-directory entries (stray files, .DS_Store) are filtered out so they don't trigger spurious NNRun.load failures.

Source code in nnx/nn/params/nn_run.py

@staticmethod
def all(root: Optional[str] = None) -> list[NNRun]:
    """List every saved NNRun under the runs root, skipping the `best`
    pointer. Returns [] when the runs/ directory doesn't exist yet.
    Non-directory entries (stray files, .DS_Store) are filtered out
    so they don't trigger spurious NNRun.load failures."""
    runs_root = _runs_root(root)
    if not os.path.isdir(runs_root):
        return []
    result: list[NNRun] = []
    for entry in os.listdir(runs_root):
        if entry == "best":
            continue
        if not os.path.isdir(os.path.join(runs_root, entry)):
            continue
        # Defensive: a directory in runs/ that lacks run.yaml isn't a
        # real run (could be a leftover from an aborted experiment).
        if not os.path.isfile(os.path.join(runs_root, entry, "run.yaml")):
            continue
        result.append(NNRun.load(id=entry, root=root))
    return result

checkpoints(root: Optional[str] = None) -> list[Optional[NNCheckpoint]]

Load this run's five phase checkpoints, in cadence order (FIRST, Q1, Q2, Q3, LAST). Entries are None when the tag was never written — e.g. runs trained with save_phase_checkpoints=False write only LAST and BEST.

BEST is deliberately excluded: it duplicates whichever phase checkpoint won, so including it would double-count. Load it directly via NNCheckpoint.load(run=run.id, type=Checkpoints.BEST).

Source code in nnx/nn/params/nn_run.py

def checkpoints(self, root: Optional[str] = None) -> list[Optional[NNCheckpoint]]:
    """Load this run's five phase checkpoints, in cadence order
    (FIRST, Q1, Q2, Q3, LAST). Entries are None when the tag was
    never written — e.g. runs trained with
    ``save_phase_checkpoints=False`` write only LAST and BEST.

    BEST is deliberately excluded: it duplicates whichever phase
    checkpoint won, so including it would double-count. Load it
    directly via ``NNCheckpoint.load(run=run.id,
    type=Checkpoints.BEST)``.
    """
    return [
        NNCheckpoint.load(run=self.id, type=Checkpoints.FIRST, root=root),
        NNCheckpoint.load(run=self.id, type=Checkpoints.Q1, root=root),
        NNCheckpoint.load(run=self.id, type=Checkpoints.Q2, root=root),
        NNCheckpoint.load(run=self.id, type=Checkpoints.Q3, root=root),
        NNCheckpoint.load(run=self.id, type=Checkpoints.LAST, root=root),
    ]

nnx.nn.params.nn_checkpoint.NNCheckpoint dataclass

Source code in nnx/nn/params/nn_checkpoint.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNCheckpoint:
    net_params: NNParams
    net_state: OrderedDict
    model_params: NNModelParams
    idp: NNIterationDataPoint

    def to_file(self, path: str, format: Literal["pickle", "safetensors"] = "pickle") -> None:
        """Atomically write this NNCheckpoint to ``path``.

        Args:
            path: destination path. Parent directory is created if missing.
            format: one of:

                - ``"pickle"`` (default): a ``torch.save`` of the whole
                  NNCheckpoint dataclass. Bit-exact round-trip including
                  the OrderedDict state and the dataclass identity. The
                  on-disk format NNx has always written; back-compat
                  default for existing callers.
                - ``"safetensors"``: a ``.safetensors`` file with the
                  net's tensors as the data section and
                  NNParams + NNModelParams + NNIterationDataPoint
                  JSON-serialized into the metadata dict (str→str only,
                  per the safetensors spec). Safe to mmap, readable by
                  ComfyUI/vLLM/AutoGPTQ/HF tools, and proof against
                  arbitrary-code-execution on load. Requires the
                  ``thekaveh-nnx[hub]`` extra.

        Both formats write to ``<path>.tmp`` first and rename into place
        so a KeyboardInterrupt during the underlying save can never leave
        a half-written checkpoint at the destination — matching the
        atomicity guarantee NNRun.save offers for YAML/CSV.
        """
        dir_path = os.path.dirname(path)
        if dir_path and not os.path.exists(dir_path):
            os.makedirs(dir_path)

        if format == "pickle":
            _atomic_torch_save(self, path)
            return
        if format == "safetensors":
            self._to_safetensors_file(path)
            return
        raise ValueError(f"unknown checkpoint format: {format!r} (expected 'pickle' or 'safetensors')")

    def _to_safetensors_file(self, path: str) -> None:
        """Atomic safetensors write. Requires the ``thekaveh-nnx[hub]`` extra."""
        try:
            from safetensors.torch import save_file
        except ImportError as e:  # pragma: no cover — gated by optional dep
            raise ImportError(
                "safetensors checkpoint format requires the `hub` extra: `pip install thekaveh-nnx[hub]`."
            ) from e

        # safetensors metadata is str→str only — every value must be a
        # string. We JSON-encode each subsection so the reader can parse
        # them back without ambiguity.
        metadata = {
            "nnx_format_version": _SAFETENSORS_FORMAT_VERSION,
            "model_params": json.dumps(self.model_params.state()),
            "net_params": json.dumps(self.net_params.state()),
            "idp": json.dumps(self.idp.state()),
        }

        # safetensors save_file doesn't accept OrderedDict (only dict). The
        # iteration order is preserved either way in Python 3.7+, so coerce.
        # Detach drops any autograd graph attached to live params; clone
        # breaks storage sharing — safetensors rejects tied tensors
        # (TransformerNN's tok_embed/lm_head share storage by default,
        # and .contiguous() is a no-op on already-contiguous views).
        # On reload, load_state_dict assigns both identical copies back
        # into the tied parameter, so the tie survives the round-trip.
        tensors = {k: v.detach().contiguous().clone() for k, v in self.net_state.items()}

        tmp = path + ".tmp"
        save_file(tensors, tmp, metadata=metadata)
        os.replace(tmp, path)

    def save(
        self,
        run: str,
        type: Checkpoints,
        root: Optional[str] = None,
        optimizer_state: Optional[OrderedDict] = None,
    ) -> None:
        """Save the checkpoint to disk atomically.

        When `optimizer_state` is supplied, a sibling file is written at
        ``<id>/checkpoints/<type>.pt.opt.pt`` (the checkpoint path plus
        an ``.opt.pt`` suffix) holding the optimizer state dict.
        This sidecar is used by NNModel.train(resume_from=...) to warm-resume
        with the prior optimizer momentum / Adam state.

        Each of the two writes is individually atomic, but the PAIR is
        not: an interrupt between them leaves a fresh checkpoint beside
        the previous epoch's sidecar, and a warm-resume then restores
        one-epoch-stale optimizer state (weights stay correct; momentum
        is briefly mismatched). Accepted: a validation stamp would
        change the sidecar format for a one-epoch-soft failure mode.
        """
        ckpt_path = _checkpoint_path(run, type, root=root)
        self.to_file(path=ckpt_path)
        if optimizer_state is not None:
            _atomic_torch_save(optimizer_state, ckpt_path + ".opt.pt")

    @staticmethod
    def load_optimizer_state(
        run: str,
        type: Checkpoints,
        root: Optional[str] = None,
    ) -> Optional[OrderedDict]:
        """Load the optimizer state sidecar for a checkpoint. Returns None
        when no sidecar exists (e.g., checkpoints written before resume
        support was added).

        Loaded with ``weights_only=True`` — the optimizer state-dict
        contains only tensors and standard scalar/dict/list types, so the
        strict loader works AND it removes the arbitrary-code-execution
        risk that the main NNCheckpoint.from_file documents."""
        path = _checkpoint_path(run, type, root=root) + ".opt.pt"
        if not os.path.exists(path):
            return None
        return torch.load(path, weights_only=True)

    @staticmethod
    def from_file(path: str) -> Optional[NNCheckpoint]:
        """Load an NNCheckpoint from disk, auto-detecting pickle vs safetensors.

        Returns ``None`` if the path doesn't exist or the loaded pickle
        object isn't an NNCheckpoint instance.

        Dispatch is by magic bytes:

        - ``torch.save`` writes a ZIP archive in modern PyTorch
          (``_use_new_zipfile_serialization=True`` is the default since
          PyTorch 1.6), so the file starts with ``b"PK\\x03\\x04"``.
        - Legacy ``torch.save`` (with the zipfile serialization disabled)
          and bare pickle files begin with ``\\x80`` (the pickle PROTO
          opcode for protocol >= 2).
        - safetensors files begin with a little-endian u64 header length
          followed by a JSON object — byte 8 is always ``{``. The u64's
          LOW byte can legitimately be ``0x80`` (any header length
          ≡ 128 mod 256), which would collide with the pickle PROTO
          opcode — so safetensors is positively identified by byte 8
          BEFORE the ``\x80`` pickle check. The ZIP magic is checked
          first of all (a ZIP's byte 8 is the compression method, never
          ``{``; a torch-LEGACY pickle has the fixed magic byte ``0xf9``
          at offset 8, and a protocol ≥ 4 bare pickle has a frame-length
          byte there, ``0x00`` for any file under a terabyte. A
          protocol-2/3 *bare* pickle's byte 8 is content-dependent, but
          NNx never produces bare pickles and such a file failed under
          the old routing too).

        Anything matching none of the positive sniffs falls through to
        the safetensors loader, whose error on a genuinely corrupt file
        is clearer than a misleading unpickle attempt.

        SECURITY: the pickle branch calls ``torch.load(weights_only=False)``,
        which unpickles arbitrary Python objects. NEVER call this on a
        checkpoint file from an untrusted source — a malicious .pt file
        can execute arbitrary code at load time. The default
        ``./runs/<id>/checkpoints/`` layout assumes the files were
        produced locally by NNCheckpoint.save. For untrusted sources,
        use the safetensors path on save and load: safetensors has no
        arbitrary-code path.
        """
        if not os.path.exists(path):
            return None

        with open(path, "rb") as f:
            head = f.read(9)
        if not head:
            return None

        # Modern torch-save ZIP container.
        if head[:4] == b"PK\x03\x04":
            return NNCheckpoint._from_pickle_file(path)
        # safetensors: byte 8 is the opening brace of the JSON header.
        # Checked BEFORE the \x80 pickle sniff — see the docstring.
        if len(head) == 9 and head[8:9] == b"{":
            return NNCheckpoint._from_safetensors_file(path)
        # Legacy / bare-pickle protocol prefix.
        if head[:1] == b"\x80":
            return NNCheckpoint._from_pickle_file(path)
        return NNCheckpoint._from_safetensors_file(path)

    @staticmethod
    def _from_pickle_file(path: str) -> Optional[NNCheckpoint]:
        """Load the legacy pickle format (a torch.save'd NNCheckpoint)."""
        # NNCheckpoint files are pickled Python objects (not bare state dicts),
        # so the weights_only=True default introduced in torch>=2.6 would raise
        # UnpicklingError. See the security note in `from_file` before
        # widening this to externally-sourced files.
        ret = torch.load(path, weights_only=False)

        if not isinstance(ret, NNCheckpoint):
            return None

        return ret

    @staticmethod
    def _from_safetensors_file(path: str) -> NNCheckpoint:
        """Load a safetensors-format NNCheckpoint. Requires `thekaveh-nnx[hub]`."""
        try:
            from safetensors import safe_open
        except ImportError as e:  # pragma: no cover — gated by optional dep
            raise ImportError(
                "loading a safetensors checkpoint requires the `hub` extra: `pip install thekaveh-nnx[hub]`."
            ) from e

        net_state: OrderedDict[str, torch.Tensor] = OrderedDict()
        with safe_open(path, framework="pt") as f:
            meta = f.metadata() or {}
            # Preserve key order — safetensors guarantees insertion-order on
            # iteration in v0.5+; we rebuild the OrderedDict for parity with
            # the pickle path.
            for k in f.keys():
                net_state[k] = f.get_tensor(k)

        # Future-proof: today only version "1" exists. If a future writer
        # bumps the version, we'd add a dispatch here on `meta["nnx_format_version"]`.
        return NNCheckpoint(
            idp=_idp_from_nested_state(json.loads(meta["idp"])),
            model_params=NNModelParams.from_state(json.loads(meta["model_params"])),
            # resolve_from_state: transformer checkpoints round-trip as
            # NNTransformerParams instead of degrading to base NNParams.
            net_params=NNParams.resolve_from_state(json.loads(meta["net_params"])),
            net_state=net_state,
        )

    @staticmethod
    def load(run: str, type: Checkpoints, root: Optional[str] = None) -> Optional[NNCheckpoint]:
        return NNCheckpoint.from_file(path=_checkpoint_path(run, type, root=root))

from_file(path: str) -> Optional[NNCheckpoint] staticmethod

Load an NNCheckpoint from disk, auto-detecting pickle vs safetensors.

Returns None if the path doesn't exist or the loaded pickle object isn't an NNCheckpoint instance.

Dispatch is by magic bytes:

• torch.save writes a ZIP archive in modern PyTorch (_use_new_zipfile_serialization=True is the default since PyTorch 1.6), so the file starts with b"PK\x03\x04".
• Legacy torch.save (with the zipfile serialization disabled) and bare pickle files begin with \x80 (the pickle PROTO opcode for protocol >= 2).
• safetensors files begin with a little-endian u64 header length followed by a JSON object — byte 8 is always {. The u64's LOW byte can legitimately be 0x80 (any header length ≡ 128 mod 256), which would collide with the pickle PROTO opcode — so safetensors is positively identified by byte 8 BEFORE the € pickle check. The ZIP magic is checked first of all (a ZIP's byte 8 is the compression method, never {; a torch-LEGACY pickle has the fixed magic byte 0xf9 at offset 8, and a protocol ≥ 4 bare pickle has a frame-length byte there, 0x00 for any file under a terabyte. A protocol-2/3 bare pickle's byte 8 is content-dependent, but NNx never produces bare pickles and such a file failed under the old routing too).

Anything matching none of the positive sniffs falls through to the safetensors loader, whose error on a genuinely corrupt file is clearer than a misleading unpickle attempt.

SECURITY: the pickle branch calls torch.load(weights_only=False), which unpickles arbitrary Python objects. NEVER call this on a checkpoint file from an untrusted source — a malicious .pt file can execute arbitrary code at load time. The default ./runs/<id>/checkpoints/ layout assumes the files were produced locally by NNCheckpoint.save. For untrusted sources, use the safetensors path on save and load: safetensors has no arbitrary-code path.

Source code in nnx/nn/params/nn_checkpoint.py

@staticmethod
def from_file(path: str) -> Optional[NNCheckpoint]:
    """Load an NNCheckpoint from disk, auto-detecting pickle vs safetensors.

    Returns ``None`` if the path doesn't exist or the loaded pickle
    object isn't an NNCheckpoint instance.

    Dispatch is by magic bytes:

    - ``torch.save`` writes a ZIP archive in modern PyTorch
      (``_use_new_zipfile_serialization=True`` is the default since
      PyTorch 1.6), so the file starts with ``b"PK\\x03\\x04"``.
    - Legacy ``torch.save`` (with the zipfile serialization disabled)
      and bare pickle files begin with ``\\x80`` (the pickle PROTO
      opcode for protocol >= 2).
    - safetensors files begin with a little-endian u64 header length
      followed by a JSON object — byte 8 is always ``{``. The u64's
      LOW byte can legitimately be ``0x80`` (any header length
      ≡ 128 mod 256), which would collide with the pickle PROTO
      opcode — so safetensors is positively identified by byte 8
      BEFORE the ``\x80`` pickle check. The ZIP magic is checked
      first of all (a ZIP's byte 8 is the compression method, never
      ``{``; a torch-LEGACY pickle has the fixed magic byte ``0xf9``
      at offset 8, and a protocol ≥ 4 bare pickle has a frame-length
      byte there, ``0x00`` for any file under a terabyte. A
      protocol-2/3 *bare* pickle's byte 8 is content-dependent, but
      NNx never produces bare pickles and such a file failed under
      the old routing too).

    Anything matching none of the positive sniffs falls through to
    the safetensors loader, whose error on a genuinely corrupt file
    is clearer than a misleading unpickle attempt.

    SECURITY: the pickle branch calls ``torch.load(weights_only=False)``,
    which unpickles arbitrary Python objects. NEVER call this on a
    checkpoint file from an untrusted source — a malicious .pt file
    can execute arbitrary code at load time. The default
    ``./runs/<id>/checkpoints/`` layout assumes the files were
    produced locally by NNCheckpoint.save. For untrusted sources,
    use the safetensors path on save and load: safetensors has no
    arbitrary-code path.
    """
    if not os.path.exists(path):
        return None

    with open(path, "rb") as f:
        head = f.read(9)
    if not head:
        return None

    # Modern torch-save ZIP container.
    if head[:4] == b"PK\x03\x04":
        return NNCheckpoint._from_pickle_file(path)
    # safetensors: byte 8 is the opening brace of the JSON header.
    # Checked BEFORE the \x80 pickle sniff — see the docstring.
    if len(head) == 9 and head[8:9] == b"{":
        return NNCheckpoint._from_safetensors_file(path)
    # Legacy / bare-pickle protocol prefix.
    if head[:1] == b"\x80":
        return NNCheckpoint._from_pickle_file(path)
    return NNCheckpoint._from_safetensors_file(path)

load_optimizer_state(run: str, type: Checkpoints, root: Optional[str] = None) -> Optional[OrderedDict] staticmethod

Load the optimizer state sidecar for a checkpoint. Returns None when no sidecar exists (e.g., checkpoints written before resume support was added).

Loaded with weights_only=True — the optimizer state-dict contains only tensors and standard scalar/dict/list types, so the strict loader works AND it removes the arbitrary-code-execution risk that the main NNCheckpoint.from_file documents.

Source code in nnx/nn/params/nn_checkpoint.py

@staticmethod
def load_optimizer_state(
    run: str,
    type: Checkpoints,
    root: Optional[str] = None,
) -> Optional[OrderedDict]:
    """Load the optimizer state sidecar for a checkpoint. Returns None
    when no sidecar exists (e.g., checkpoints written before resume
    support was added).

    Loaded with ``weights_only=True`` — the optimizer state-dict
    contains only tensors and standard scalar/dict/list types, so the
    strict loader works AND it removes the arbitrary-code-execution
    risk that the main NNCheckpoint.from_file documents."""
    path = _checkpoint_path(run, type, root=root) + ".opt.pt"
    if not os.path.exists(path):
        return None
    return torch.load(path, weights_only=True)

save(run: str, type: Checkpoints, root: Optional[str] = None, optimizer_state: Optional[OrderedDict] = None) -> None

Save the checkpoint to disk atomically.

When optimizer_state is supplied, a sibling file is written at <id>/checkpoints/<type>.pt.opt.pt (the checkpoint path plus an .opt.pt suffix) holding the optimizer state dict. This sidecar is used by NNModel.train(resume_from=...) to warm-resume with the prior optimizer momentum / Adam state.

Each of the two writes is individually atomic, but the PAIR is not: an interrupt between them leaves a fresh checkpoint beside the previous epoch's sidecar, and a warm-resume then restores one-epoch-stale optimizer state (weights stay correct; momentum is briefly mismatched). Accepted: a validation stamp would change the sidecar format for a one-epoch-soft failure mode.

Source code in nnx/nn/params/nn_checkpoint.py

def save(
    self,
    run: str,
    type: Checkpoints,
    root: Optional[str] = None,
    optimizer_state: Optional[OrderedDict] = None,
) -> None:
    """Save the checkpoint to disk atomically.

    When `optimizer_state` is supplied, a sibling file is written at
    ``<id>/checkpoints/<type>.pt.opt.pt`` (the checkpoint path plus
    an ``.opt.pt`` suffix) holding the optimizer state dict.
    This sidecar is used by NNModel.train(resume_from=...) to warm-resume
    with the prior optimizer momentum / Adam state.

    Each of the two writes is individually atomic, but the PAIR is
    not: an interrupt between them leaves a fresh checkpoint beside
    the previous epoch's sidecar, and a warm-resume then restores
    one-epoch-stale optimizer state (weights stay correct; momentum
    is briefly mismatched). Accepted: a validation stamp would
    change the sidecar format for a one-epoch-soft failure mode.
    """
    ckpt_path = _checkpoint_path(run, type, root=root)
    self.to_file(path=ckpt_path)
    if optimizer_state is not None:
        _atomic_torch_save(optimizer_state, ckpt_path + ".opt.pt")

to_file(path: str, format: Literal['pickle', 'safetensors'] = 'pickle') -> None

Atomically write this NNCheckpoint to path.

Parameters:

Name	Type	Description	Default
path	str	destination path. Parent directory is created if missing.	required
format	Literal['pickle', 'safetensors']	one of: "pickle" (default): a torch.save of the whole NNCheckpoint dataclass. Bit-exact round-trip including the OrderedDict state and the dataclass identity. The on-disk format NNx has always written; back-compat default for existing callers. "safetensors": a .safetensors file with the net's tensors as the data section and NNParams + NNModelParams + NNIterationDataPoint JSON-serialized into the metadata dict (str→str only, per the safetensors spec). Safe to mmap, readable by ComfyUI/vLLM/AutoGPTQ/HF tools, and proof against arbitrary-code-execution on load. Requires the thekaveh-nnx[hub] extra.	'pickle'

Both formats write to <path>.tmp first and rename into place so a KeyboardInterrupt during the underlying save can never leave a half-written checkpoint at the destination — matching the atomicity guarantee NNRun.save offers for YAML/CSV.

Source code in nnx/nn/params/nn_checkpoint.py

def to_file(self, path: str, format: Literal["pickle", "safetensors"] = "pickle") -> None:
    """Atomically write this NNCheckpoint to ``path``.

    Args:
        path: destination path. Parent directory is created if missing.
        format: one of:

            - ``"pickle"`` (default): a ``torch.save`` of the whole
              NNCheckpoint dataclass. Bit-exact round-trip including
              the OrderedDict state and the dataclass identity. The
              on-disk format NNx has always written; back-compat
              default for existing callers.
            - ``"safetensors"``: a ``.safetensors`` file with the
              net's tensors as the data section and
              NNParams + NNModelParams + NNIterationDataPoint
              JSON-serialized into the metadata dict (str→str only,
              per the safetensors spec). Safe to mmap, readable by
              ComfyUI/vLLM/AutoGPTQ/HF tools, and proof against
              arbitrary-code-execution on load. Requires the
              ``thekaveh-nnx[hub]`` extra.

    Both formats write to ``<path>.tmp`` first and rename into place
    so a KeyboardInterrupt during the underlying save can never leave
    a half-written checkpoint at the destination — matching the
    atomicity guarantee NNRun.save offers for YAML/CSV.
    """
    dir_path = os.path.dirname(path)
    if dir_path and not os.path.exists(dir_path):
        os.makedirs(dir_path)

    if format == "pickle":
        _atomic_torch_save(self, path)
        return
    if format == "safetensors":
        self._to_safetensors_file(path)
        return
    raise ValueError(f"unknown checkpoint format: {format!r} (expected 'pickle' or 'safetensors')")

nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint dataclass

One row in the per-iteration training log.

train_edp is computed from the current batch only. val_edp is the per-epoch validation evaluation — populated only on the last idp of each epoch (the idp at which the validation loop ran). Other idps in the same epoch have val_edp=None. When reading idps.csv, group by epoch_idx and take the row with val_edp set for per-epoch validation metrics.

Source code in nnx/nn/params/nn_iteration_data_point.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNIterationDataPoint:
    """One row in the per-iteration training log.

    `train_edp` is computed from the current batch only. `val_edp` is the
    per-epoch validation evaluation — populated **only on the last idp of
    each epoch** (the idp at which the validation loop ran). Other idps in
    the same epoch have `val_edp=None`. When reading idps.csv, group by
    epoch_idx and take the row with val_edp set for per-epoch validation
    metrics.
    """

    lr: float
    iter_idx: int
    epoch_idx: int
    batch_idx: int
    train_edp: NNEvaluationDataPoint
    val_edp: Optional[NNEvaluationDataPoint] = None

    def with_val_edp(self, value: Optional[NNEvaluationDataPoint]) -> NNIterationDataPoint:
        return replace(self, val_edp=value)

    def state(self) -> dict:
        return dict(
            lr=self.lr,
            iter_idx=self.iter_idx,
            epoch_idx=self.epoch_idx,
            batch_idx=self.batch_idx,
            train_edp=self.train_edp.state(),
            val_edp=self.val_edp.state() if self.val_edp is not None else None,
        )

    @staticmethod
    def from_state(state: dict) -> NNIterationDataPoint:
        # Reassemble the `extra` dict from flattened CSV columns. After
        # NNRun.save, pd.json_normalize flattens nested {prefix: {name: v}}
        # into `<prefix>.extra.<name>` columns. We collect them back into
        # the inner state dict so NNEvaluationDataPoint.from_state can
        # populate the extra field correctly.
        def _collect_extra(prefix: str) -> dict:
            marker = f"{prefix}.extra."
            return {
                k[len(marker) :]: v
                for k, v in state.items()
                if k.startswith(marker)
                and v is not None
                # NaN values appear in CSV when other idps in the run had
                # the key set but this row didn't — filter via isna check.
                and not _is_nan(v)
            }

        def _field(key: str):
            # CSV round-trip (NNRun.load → pd.read_csv) yields NaN, not
            # None, for cells that were None at save time. Map both back
            # to None so loaded idps match what state() wrote.
            v = state.get(key)
            return None if v is None or _is_nan(v) else v

        val_edp = None
        if any(_field(f"val_edp.{k}") is not None for k in ("loss", "error", "accuracy", "f1", "recall", "precision")):
            val_edp = NNEvaluationDataPoint.from_state(
                dict(
                    loss=_field("val_edp.loss"),
                    error=_field("val_edp.error"),
                    accuracy=_field("val_edp.accuracy"),
                    f1=_field("val_edp.f1"),
                    recall=_field("val_edp.recall"),
                    precision=_field("val_edp.precision"),
                    extra=_collect_extra("val_edp"),
                )
            )
        return NNIterationDataPoint(
            lr=state["lr"],
            iter_idx=state["iter_idx"],
            epoch_idx=state["epoch_idx"],
            batch_idx=state["batch_idx"],
            train_edp=NNEvaluationDataPoint.from_state(
                dict(
                    loss=_field("train_edp.loss"),
                    error=_field("train_edp.error"),
                    accuracy=_field("train_edp.accuracy"),
                    f1=_field("train_edp.f1"),
                    recall=_field("train_edp.recall"),
                    precision=_field("train_edp.precision"),
                    extra=_collect_extra("train_edp"),
                )
            ),
            val_edp=val_edp,
        )

nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint dataclass

Per-batch / per-epoch evaluation metrics.

The four core fields (f1, recall, accuracy, precision) are computed by of() via sklearn. loss and error are typically attached after the fact by NNModel during training / evaluation.

extra is a free-form dict of user-supplied custom metric names to floats. Populated when NNTrainParams.extra_metrics or evaluate(metrics=) is set; empty by default (and omitted from state() when empty so that pre-extra runs hash to the same run.id and pre-extra YAML loads cleanly).

Source code in nnx/nn/params/nn_evaluation_data_point.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNEvaluationDataPoint:
    """Per-batch / per-epoch evaluation metrics.

    The four core fields (f1, recall, accuracy, precision) are computed by
    `of()` via sklearn. `loss` and `error` are typically attached after the
    fact by NNModel during training / evaluation.

    `extra` is a free-form dict of user-supplied custom metric names to
    floats. Populated when NNTrainParams.extra_metrics or evaluate(metrics=)
    is set; empty by default (and omitted from state() when empty so that
    pre-extra runs hash to the same run.id and pre-extra YAML loads cleanly).
    """

    f1: float
    recall: float
    accuracy: float
    precision: float
    loss: Optional[float] = None
    error: Optional[float] = None

    # Custom metrics injected by the caller. Keys are metric names; values
    # are floats. Default factory keeps the dataclass hashable-by-value via
    # the dict default.
    extra: dict = field(default_factory=dict)

    def with_loss(self, value: float):
        return replace(self, loss=value)

    def with_error(self, value: float):
        return replace(self, error=value)

    def with_extra(self, name: str, value: float) -> NNEvaluationDataPoint:
        merged = {**self.extra, name: float(value)}
        return replace(self, extra=merged)

    @staticmethod
    def of(
        Y: np.ndarray,
        Y_hat: np.ndarray,
        average: str = "macro",
        extra_metrics: Optional[Mapping[str, Callable]] = None,
    ):
        """Compute per-batch evaluation metrics.

        `average` controls how f1/precision/recall reduce across classes.
        Default "macro" treats all classes equally — the right choice for
        multi-class classification and the only one that makes f1/precision/
        recall mathematically distinct from accuracy. Pass "micro" to
        recover the legacy behavior (numerically identical to accuracy for
        single-label multi-class). Accuracy itself is not affected.

        `extra_metrics` is a {name -> callable(Y, Y_hat) -> float} map of
        user-supplied custom metrics. Each is invoked once on the aggregate
        predictions and stored in the returned object's `extra` dict.
        """
        extra: dict[str, float] = {}
        if extra_metrics:
            for name, fn in extra_metrics.items():
                extra[name] = float(fn(Y, Y_hat))

        return NNEvaluationDataPoint(
            accuracy=metrics.accuracy_score(y_true=Y, y_pred=Y_hat),
            f1=metrics.f1_score(y_true=Y, y_pred=Y_hat, average=average, zero_division=0),
            recall=metrics.recall_score(y_true=Y, y_pred=Y_hat, average=average, zero_division=0),
            precision=metrics.precision_score(y_true=Y, y_pred=Y_hat, average=average, zero_division=0),
            extra=extra,
        )

    @staticmethod
    def mean_of(edps: list[NNEvaluationDataPoint]) -> NNEvaluationDataPoint:
        """Unweighted-mean reduce a list of EDPs across every metric.

        .. warning::

            This is a **simple mean across edps**, NOT a sample-weighted
            mean. With unequal batch sizes (the common case), the result
            is statistically incorrect — a 1024-sample batch counts the
            same as an 8-sample tail batch. For correct sample-weighted
            metrics across batches, use :meth:`NNModel.evaluate`, which
            concatenates predictions across the loader and computes once
            on the full sample.

            ``mean_of`` is kept for back-compat with callers that already
            depend on the unweighted-mean semantics; new code should
            prefer :meth:`NNModel.evaluate` unless the unweighted form is
            specifically what's wanted (e.g., averaging across runs, not
            across batches within a run).

        An ``extra`` key present on some but not all edps is averaged over
        the edps where it IS present (skipped on the rest).
        """
        # Aggregate the standard fields with the existing logic.
        ret = NNEvaluationDataPoint(
            f1=np.mean([edp.f1 for edp in edps]),
            recall=np.mean([edp.recall for edp in edps]),
            accuracy=np.mean([edp.accuracy for edp in edps]),
            precision=np.mean([edp.precision for edp in edps]),
        )

        if len([edp.loss for edp in edps if edp.loss is not None]) > 0:
            ret = ret.with_loss(np.mean([edp.loss for edp in edps if edp.loss is not None]))

        if len([edp.error for edp in edps if edp.error is not None]) > 0:
            ret = ret.with_error(np.mean([edp.error for edp in edps if edp.error is not None]))

        # Propagate extras: union the key set across edps, mean per key
        # over the edps that have it. Keys missing from some edps are
        # skipped on those, not zero-filled.
        all_extra_keys: set[str] = set()
        for edp in edps:
            all_extra_keys.update(edp.extra.keys())
        if all_extra_keys:
            extra_mean: dict[str, float] = {}
            for k in all_extra_keys:
                values = [edp.extra[k] for edp in edps if k in edp.extra]
                if values:
                    extra_mean[k] = float(np.mean(values))
            ret = replace(ret, extra=extra_mean)

        return ret

    def state(self) -> dict:
        d = dict(
            f1=self.f1,
            recall=self.recall,
            accuracy=self.accuracy,
            precision=self.precision,
            loss=self.loss,
            error=self.error,
        )
        # Omit `extra` when empty so EDPs from before this field existed
        # remain bit-for-bit identical in state() form (preserves run.id
        # back-compat).
        if self.extra:
            d["extra"] = dict(self.extra)
        return d

    @staticmethod
    def from_state(state: dict) -> NNEvaluationDataPoint:
        return NNEvaluationDataPoint(
            f1=state["f1"],
            recall=state["recall"],
            accuracy=state["accuracy"],
            precision=state["precision"],
            loss=state["loss"],
            error=state["error"],
            extra=dict(state.get("extra") or {}),
        )

mean_of(edps: list[NNEvaluationDataPoint]) -> NNEvaluationDataPoint staticmethod

Unweighted-mean reduce a list of EDPs across every metric.

.. warning::

This is a **simple mean across edps**, NOT a sample-weighted
mean. With unequal batch sizes (the common case), the result
is statistically incorrect — a 1024-sample batch counts the
same as an 8-sample tail batch. For correct sample-weighted
metrics across batches, use :meth:`NNModel.evaluate`, which
concatenates predictions across the loader and computes once
on the full sample.

``mean_of`` is kept for back-compat with callers that already
depend on the unweighted-mean semantics; new code should
prefer :meth:`NNModel.evaluate` unless the unweighted form is
specifically what's wanted (e.g., averaging across runs, not
across batches within a run).

An extra key present on some but not all edps is averaged over the edps where it IS present (skipped on the rest).

Source code in nnx/nn/params/nn_evaluation_data_point.py

@staticmethod
def mean_of(edps: list[NNEvaluationDataPoint]) -> NNEvaluationDataPoint:
    """Unweighted-mean reduce a list of EDPs across every metric.

    .. warning::

        This is a **simple mean across edps**, NOT a sample-weighted
        mean. With unequal batch sizes (the common case), the result
        is statistically incorrect — a 1024-sample batch counts the
        same as an 8-sample tail batch. For correct sample-weighted
        metrics across batches, use :meth:`NNModel.evaluate`, which
        concatenates predictions across the loader and computes once
        on the full sample.

        ``mean_of`` is kept for back-compat with callers that already
        depend on the unweighted-mean semantics; new code should
        prefer :meth:`NNModel.evaluate` unless the unweighted form is
        specifically what's wanted (e.g., averaging across runs, not
        across batches within a run).

    An ``extra`` key present on some but not all edps is averaged over
    the edps where it IS present (skipped on the rest).
    """
    # Aggregate the standard fields with the existing logic.
    ret = NNEvaluationDataPoint(
        f1=np.mean([edp.f1 for edp in edps]),
        recall=np.mean([edp.recall for edp in edps]),
        accuracy=np.mean([edp.accuracy for edp in edps]),
        precision=np.mean([edp.precision for edp in edps]),
    )

    if len([edp.loss for edp in edps if edp.loss is not None]) > 0:
        ret = ret.with_loss(np.mean([edp.loss for edp in edps if edp.loss is not None]))

    if len([edp.error for edp in edps if edp.error is not None]) > 0:
        ret = ret.with_error(np.mean([edp.error for edp in edps if edp.error is not None]))

    # Propagate extras: union the key set across edps, mean per key
    # over the edps that have it. Keys missing from some edps are
    # skipped on those, not zero-filled.
    all_extra_keys: set[str] = set()
    for edp in edps:
        all_extra_keys.update(edp.extra.keys())
    if all_extra_keys:
        extra_mean: dict[str, float] = {}
        for k in all_extra_keys:
            values = [edp.extra[k] for edp in edps if k in edp.extra]
            if values:
                extra_mean[k] = float(np.mean(values))
        ret = replace(ret, extra=extra_mean)

    return ret

of(Y: np.ndarray, Y_hat: np.ndarray, average: str = 'macro', extra_metrics: Optional[Mapping[str, Callable]] = None) staticmethod

Compute per-batch evaluation metrics.

average controls how f1/precision/recall reduce across classes. Default "macro" treats all classes equally — the right choice for multi-class classification and the only one that makes f1/precision/ recall mathematically distinct from accuracy. Pass "micro" to recover the legacy behavior (numerically identical to accuracy for single-label multi-class). Accuracy itself is not affected.

extra_metrics is a {name -> callable(Y, Y_hat) -> float} map of user-supplied custom metrics. Each is invoked once on the aggregate predictions and stored in the returned object's extra dict.

Source code in nnx/nn/params/nn_evaluation_data_point.py

@staticmethod
def of(
    Y: np.ndarray,
    Y_hat: np.ndarray,
    average: str = "macro",
    extra_metrics: Optional[Mapping[str, Callable]] = None,
):
    """Compute per-batch evaluation metrics.

    `average` controls how f1/precision/recall reduce across classes.
    Default "macro" treats all classes equally — the right choice for
    multi-class classification and the only one that makes f1/precision/
    recall mathematically distinct from accuracy. Pass "micro" to
    recover the legacy behavior (numerically identical to accuracy for
    single-label multi-class). Accuracy itself is not affected.

    `extra_metrics` is a {name -> callable(Y, Y_hat) -> float} map of
    user-supplied custom metrics. Each is invoked once on the aggregate
    predictions and stored in the returned object's `extra` dict.
    """
    extra: dict[str, float] = {}
    if extra_metrics:
        for name, fn in extra_metrics.items():
            extra[name] = float(fn(Y, Y_hat))

    return NNEvaluationDataPoint(
        accuracy=metrics.accuracy_score(y_true=Y, y_pred=Y_hat),
        f1=metrics.f1_score(y_true=Y, y_pred=Y_hat, average=average, zero_division=0),
        recall=metrics.recall_score(y_true=Y, y_pred=Y_hat, average=average, zero_division=0),
        precision=metrics.precision_score(y_true=Y, y_pred=Y_hat, average=average, zero_division=0),
        extra=extra,
    )

4. Networks

nnx.nn.net.feed_fwd_nn.FeedFwdNN

Bases: Module

Source code in nnx/nn/net/feed_fwd_nn.py

class FeedFwdNN(nn.Module):
    def __init__(self, params: NNParams):
        super().__init__()

        self.params = params

        self.layers = nn.ModuleList(
            [
                nn.Linear(in_features=in_dim, out_features=out_dim)
                for in_dim, out_dim in zip(self.params.dims, self.params.dims[1:], strict=False)
            ]
        )

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        X = X.view(X.size(0), -1)

        for layer in self.layers[:-1]:
            X = layer(X)
            X = self.params.activation()(X)
            X = F.dropout(X, p=self.params.dropout_prob, training=self.training)

        X = self.layers[-1](X)

        return X

    def unpack_batch(self, batch):
        if isinstance(batch, (list, tuple)):
            X, Y = batch
            return (X,), Y
        # PyG batches go through the GNN subclasses (GraphConvNN /
        # GraphSageNN / GraphAttNN), but we keep the branch here for
        # back-compat with the original contract. Lazy import so module
        # load doesn't pay the torch_geometric import cost for users
        # who only ever pass standard (X, Y) tuples.
        from torch_geometric.data.data import Data as _PygData

        if isinstance(batch, _PygData):
            X, Y = batch.x, batch.y
            return (X,), Y
        raise TypeError(
            "FeedFwdNN.unpack_batch expects a (list/tuple) batch or a "
            "torch_geometric.data.Data instance; got "
            f"{type(batch).__name__}."
        )

    def __str__(self):
        return f"FeedFwdNN={self.params}"

    def to_file(self, path: str) -> None:
        torch.save(self.state_dict(), path)

    @staticmethod
    def from_file(path: str, params: NNParams) -> FeedFwdNN:
        # weights_only=True: a state-dict is plain tensors + standard
        # scalar/dict types — the strict loader works and removes the
        # arbitrary-code-execution risk on user-supplied paths. Matches
        # NNCheckpoint.load_optimizer_state and load_pretrained.
        net = FeedFwdNN(params)
        net.load_state_dict(torch.load(path, weights_only=True))

        return net

    @staticmethod
    def from_state(state_dict: dict, params: NNParams) -> FeedFwdNN:
        net = FeedFwdNN(params)
        net.load_state_dict(state_dict)

        return net

nnx.nn.net.graph_nn_base.GraphNNBase

Bases: Module

Abstract base for GNN architectures.

Subclasses must implement _build_layers() returning an nn.ModuleList of PyG message-passing layers. The forward loop applies all-but-last layers with the configured activation + dropout, then a bare final layer.

Source code in nnx/nn/net/graph_nn_base.py

class GraphNNBase(nn.Module):
    """Abstract base for GNN architectures.

    Subclasses must implement `_build_layers()` returning an `nn.ModuleList`
    of PyG message-passing layers. The forward loop applies all-but-last
    layers with the configured activation + dropout, then a bare final layer.
    """

    def __init__(self, params: NNParams):
        super().__init__()
        self.params = params
        self.layers = self._build_layers()

    def _build_layers(self) -> nn.ModuleList:
        raise NotImplementedError("subclass must implement _build_layers()")

    def forward(self, X: torch.Tensor, E: torch.Tensor) -> torch.Tensor:
        for layer in self.layers[:-1]:
            X = layer(X, E)
            X = self.params.activation()(X)
            X = F.dropout(X, p=self.params.dropout_prob, training=self.training)
        return self.layers[-1](X, E)

    def unpack_batch(self, batch) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
        return (batch.x, batch.edge_index), batch.y

    def seed_count(self, batch) -> Optional[int]:
        """Number of seed rows at the head of a NeighborLoader subgraph.

        NeighborLoader puts the ``batch_size`` seed nodes first and
        appends their sampled neighbors — which can belong to *other*
        splits. Loss and metrics must be computed on the seed rows only;
        scoring neighbor rows leaks val/test labels into the training
        loss and train labels into val metrics.

        Returns None (no slicing) for anything that isn't a
        NeighborLoader subgraph: plain full-graph ``Data`` has no
        ``batch_size``, and a multi-graph ``Batch.from_data_list``
        collation DOES carry ``batch_size`` (= ``num_graphs``) but no
        ``input_id`` — slicing there would truncate node-level output
        to the graph count. ``input_id`` is the NeighborLoader-specific
        marker (the seed indices), so it gates the slice.
        """
        if getattr(batch, "input_id", None) is None:
            return None
        n = getattr(batch, "batch_size", None)
        return int(n) if n is not None else None

seed_count(batch) -> Optional[int]

Number of seed rows at the head of a NeighborLoader subgraph.

NeighborLoader puts the batch_size seed nodes first and appends their sampled neighbors — which can belong to other splits. Loss and metrics must be computed on the seed rows only; scoring neighbor rows leaks val/test labels into the training loss and train labels into val metrics.

Returns None (no slicing) for anything that isn't a NeighborLoader subgraph: plain full-graph Data has no batch_size, and a multi-graph Batch.from_data_list collation DOES carry batch_size (= num_graphs) but no input_id — slicing there would truncate node-level output to the graph count. input_id is the NeighborLoader-specific marker (the seed indices), so it gates the slice.

Source code in nnx/nn/net/graph_nn_base.py

def seed_count(self, batch) -> Optional[int]:
    """Number of seed rows at the head of a NeighborLoader subgraph.

    NeighborLoader puts the ``batch_size`` seed nodes first and
    appends their sampled neighbors — which can belong to *other*
    splits. Loss and metrics must be computed on the seed rows only;
    scoring neighbor rows leaks val/test labels into the training
    loss and train labels into val metrics.

    Returns None (no slicing) for anything that isn't a
    NeighborLoader subgraph: plain full-graph ``Data`` has no
    ``batch_size``, and a multi-graph ``Batch.from_data_list``
    collation DOES carry ``batch_size`` (= ``num_graphs``) but no
    ``input_id`` — slicing there would truncate node-level output
    to the graph count. ``input_id`` is the NeighborLoader-specific
    marker (the seed indices), so it gates the slice.
    """
    if getattr(batch, "input_id", None) is None:
        return None
    n = getattr(batch, "batch_size", None)
    return int(n) if n is not None else None

nnx.nn.net.graph_conv_nn.GraphConvNN

Bases: GraphNNBase

Source code in nnx/nn/net/graph_conv_nn.py

class GraphConvNN(GraphNNBase):
    def _build_layers(self) -> nn.ModuleList:
        return nn.ModuleList(
            [
                pyg.nn.GCNConv(in_channels=in_dim, out_channels=out_dim)
                for in_dim, out_dim in zip(self.params.dims, self.params.dims[1:], strict=False)
            ]
        )

    def __str__(self) -> str:
        return f"GraphConvNN={self.params}"

nnx.nn.net.graph_sage_nn.GraphSageNN

Bases: GraphNNBase

Source code in nnx/nn/net/graph_sage_nn.py

class GraphSageNN(GraphNNBase):
    def _build_layers(self) -> nn.ModuleList:
        return nn.ModuleList(
            [
                pyg.nn.SAGEConv(in_channels=in_dim, out_channels=out_dim)
                for in_dim, out_dim in zip(self.params.dims, self.params.dims[1:], strict=False)
            ]
        )

    def __str__(self) -> str:
        return f"GraphSageNN={self.params}"

nnx.nn.net.graph_att_nn.GraphAttNN

Bases: GraphNNBase

Source code in nnx/nn/net/graph_att_nn.py

class GraphAttNN(GraphNNBase):
    def _build_layers(self) -> nn.ModuleList:
        if self.params.n_heads is None or self.params.n_heads <= 0:
            raise ValueError(f"GraphAttNN requires NNParams.n_heads > 0, got {self.params.n_heads!r}")
        n_heads = self.params.n_heads
        dim_pairs = list(zip(self.params.dims, self.params.dims[1:], strict=False))
        return nn.ModuleList(
            [
                pyg.nn.GATConv(
                    out_channels=out_dim,
                    heads=n_heads,
                    concat=(idx_dim != len(dim_pairs) - 1),
                    in_channels=in_dim if idx_dim == 0 else in_dim * n_heads,
                )
                for idx_dim, (in_dim, out_dim) in enumerate(dim_pairs)
            ]
        )

    def __str__(self) -> str:
        return f"GraphAttNN={self.params}"

nnx.nn.net.transformer_nn.TransformerNN

Bases: Module

Source code in nnx/nn/net/transformer_nn.py

class TransformerNN(nn.Module):
    def __init__(self, params: NNTransformerParams):
        super().__init__()
        self.params = params

        self.tok_embed = nn.Embedding(num_embeddings=params.vocab_size, embedding_dim=params.d_model)
        # GPT-2/LLaMA-style small-std init. nn.Embedding's default is
        # N(0, 1) — with tied embeddings, the input token's own logit
        # then includes e·e ≈ d_model, so an untrained model starts at
        # CE ≈ d_model instead of ln(vocab) and greedy/sampled decoding
        # degenerates into repeating the last prompt token. The shared
        # tensor also covers lm_head when tie_embeddings=True.
        nn.init.normal_(self.tok_embed.weight, mean=0.0, std=0.02)
        self.blocks = nn.ModuleList(
            [
                TransformerBlock(
                    d_model=params.d_model,
                    n_heads=params.n_heads,
                    ffn_mult=params.ffn_mult,
                    max_seq_len=params.max_seq_len,
                    rope_base=params.rope_base,
                    attn_dropout=params.attn_dropout,
                    resid_dropout=params.resid_dropout,
                )
                for _ in range(params.n_layers)
            ]
        )
        self.norm_out = RMSNorm(params.d_model)
        # LM head: a no-bias Linear from d_model → vocab_size. When
        # tie_embeddings=True, we share the weight tensor with the
        # token embedding (parameter-shared, not just initialized to
        # the same values) — the standard LLaMA trick that cuts
        # ~vocab*d_model params for free.
        self.lm_head = nn.Linear(params.d_model, params.vocab_size, bias=False)
        if params.tie_embeddings:
            # Identity-assign the parameter — not just copy the data —
            # so .parameters() doesn't double-count and gradients flow
            # through a single shared tensor.
            self.lm_head.weight = self.tok_embed.weight

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        """Args:
            tokens: (batch, seq) long tensor of token ids.

        Returns:
            (batch, seq, vocab_size) logits — pre-softmax.
        """
        b, t = tokens.shape
        if t > self.params.max_seq_len:
            raise ValueError(f"input sequence length ({t}) exceeds max_seq_len={self.params.max_seq_len}")
        x = self.tok_embed(tokens)  # (B, T, d_model)
        for block in self.blocks:
            x, _ = block(x, use_cache=False)
        x = self.norm_out(x)
        logits = self.lm_head(x)  # (B, T, vocab)
        return logits

    def forward_with_cache(
        self,
        tokens: torch.Tensor,
        past_kvs: Optional[list[LayerKV]] = None,
    ) -> tuple[torch.Tensor, list[LayerKV]]:
        """Cache-threading forward used by ``GenerativeNNModel.generate``.

        Behaves like ``forward`` but additionally accepts a per-layer
        list of (k, v) caches (or ``None`` entries on the first call)
        and returns the updated per-layer caches alongside the logits.

        The total attended-to length per layer is
        ``past_kv_len + tokens.shape[1]`` — the caller is responsible
        for ensuring that this stays within ``max_seq_len`` (the
        generate loop slides a window when it would otherwise overflow).

        Args:
            tokens: (batch, seq) long tensor of token ids. During
                incremental decode, ``seq == 1``; on the prefill step
                the prompt's full length is fed in one shot.
            past_kvs: list of length ``n_layers`` with each entry a
                ``(k, v)`` tuple or ``None``. ``None`` means "no
                history for this layer" (i.e., first call).

        Returns:
            A tuple ``(logits, new_kvs)`` where ``logits`` is
                ``(batch, seq, vocab)`` — the *new* tokens' logits (with
                ``past_kvs != None`` and ``seq=1`` the returned
                ``logits[:, -1, :]`` is the next-token distribution
                conditioned on the full cached prefix) — and ``new_kvs``
                is a list of length ``n_layers`` of updated ``(k, v)``
                tuples; pass this back in for the next step.
        """
        b, t = tokens.shape
        # Length already cached, taken from layer 0 (all layers share length).
        cached_len = 0
        if past_kvs is not None and past_kvs[0] is not None:
            cached_len = past_kvs[0][0].size(-2)
        total_len = cached_len + t
        if total_len > self.params.max_seq_len:
            raise ValueError(
                f"cached_len ({cached_len}) + new tokens ({t}) = {total_len} "
                f"exceeds max_seq_len={self.params.max_seq_len}"
            )

        if past_kvs is None:
            past_kvs = [None] * len(self.blocks)
        if len(past_kvs) != len(self.blocks):
            raise ValueError(f"past_kvs has {len(past_kvs)} entries but model has {len(self.blocks)} layers")

        x = self.tok_embed(tokens)  # (B, T, d_model)
        new_kvs: list[LayerKV] = []
        for block, layer_past in zip(self.blocks, past_kvs, strict=True):
            x, new_kv = block(x, past_kv=layer_past, use_cache=True)
            new_kvs.append(new_kv)
        x = self.norm_out(x)
        logits = self.lm_head(x)  # (B, T, vocab)
        return logits, new_kvs

    def unpack_batch(self, batch):
        """Make TransformerNN compatible with the standard supervised
        NNModel training loop.

        For an LM the canonical batch is ``(tokens, targets)`` where
        ``targets = tokens[:, 1:]`` shifted by one. We don't shift here —
        the caller assembles the tuple — but we accept either a 2-tuple
        ``(X, Y)`` or a plain tensor of tokens (next-token loss is then
        computed in the train step).
        """
        if isinstance(batch, (list, tuple)) and len(batch) == 2:
            X, Y = batch
            return (X,), Y
        if isinstance(batch, torch.Tensor):
            # Self-supervised next-token: targets are inputs shifted by 1.
            return (batch[:, :-1],), batch[:, 1:]
        raise TypeError(
            f"TransformerNN.unpack_batch expects a (tokens, targets) tuple "
            f"or a tensor of token ids; got {type(batch).__name__}"
        )

    def __str__(self) -> str:
        return f"TransformerNN={self.params}"

    def to_file(self, path: str) -> None:
        torch.save(self.state_dict(), path)

    @staticmethod
    def from_file(path: str, params: NNTransformerParams) -> TransformerNN:
        net = TransformerNN(params)
        net.load_state_dict(torch.load(path, weights_only=True))
        return net

    @staticmethod
    def from_state(state_dict: dict, params: NNTransformerParams) -> TransformerNN:
        net = TransformerNN(params)
        net.load_state_dict(state_dict)
        return net

forward(tokens: torch.Tensor) -> torch.Tensor

Parameters:

Name	Type	Description	Default
tokens	Tensor	(batch, seq) long tensor of token ids.	required

Returns:

Type	Description
Tensor	(batch, seq, vocab_size) logits — pre-softmax.

Source code in nnx/nn/net/transformer_nn.py

def forward(self, tokens: torch.Tensor) -> torch.Tensor:
    """Args:
        tokens: (batch, seq) long tensor of token ids.

    Returns:
        (batch, seq, vocab_size) logits — pre-softmax.
    """
    b, t = tokens.shape
    if t > self.params.max_seq_len:
        raise ValueError(f"input sequence length ({t}) exceeds max_seq_len={self.params.max_seq_len}")
    x = self.tok_embed(tokens)  # (B, T, d_model)
    for block in self.blocks:
        x, _ = block(x, use_cache=False)
    x = self.norm_out(x)
    logits = self.lm_head(x)  # (B, T, vocab)
    return logits

forward_with_cache(tokens: torch.Tensor, past_kvs: Optional[list[LayerKV]] = None) -> tuple[torch.Tensor, list[LayerKV]]

Cache-threading forward used by GenerativeNNModel.generate.

Behaves like forward but additionally accepts a per-layer list of (k, v) caches (or None entries on the first call) and returns the updated per-layer caches alongside the logits.

The total attended-to length per layer is past_kv_len + tokens.shape[1] — the caller is responsible for ensuring that this stays within max_seq_len (the generate loop slides a window when it would otherwise overflow).

Parameters:

Name	Type	Description	Default
tokens	Tensor	(batch, seq) long tensor of token ids. During incremental decode, seq == 1; on the prefill step the prompt's full length is fed in one shot.	required
past_kvs	Optional[list[LayerKV]]	list of length n_layers with each entry a (k, v) tuple or None. None means "no history for this layer" (i.e., first call).	None

Returns:

Type	Description
tuple[Tensor, list[LayerKV]]	A tuple (logits, new_kvs) where logits is (batch, seq, vocab) — the new tokens' logits (with past_kvs != None and seq=1 the returned logits[:, -1, :] is the next-token distribution conditioned on the full cached prefix) — and new_kvs is a list of length n_layers of updated (k, v) tuples; pass this back in for the next step.

Source code in nnx/nn/net/transformer_nn.py

def forward_with_cache(
    self,
    tokens: torch.Tensor,
    past_kvs: Optional[list[LayerKV]] = None,
) -> tuple[torch.Tensor, list[LayerKV]]:
    """Cache-threading forward used by ``GenerativeNNModel.generate``.

    Behaves like ``forward`` but additionally accepts a per-layer
    list of (k, v) caches (or ``None`` entries on the first call)
    and returns the updated per-layer caches alongside the logits.

    The total attended-to length per layer is
    ``past_kv_len + tokens.shape[1]`` — the caller is responsible
    for ensuring that this stays within ``max_seq_len`` (the
    generate loop slides a window when it would otherwise overflow).

    Args:
        tokens: (batch, seq) long tensor of token ids. During
            incremental decode, ``seq == 1``; on the prefill step
            the prompt's full length is fed in one shot.
        past_kvs: list of length ``n_layers`` with each entry a
            ``(k, v)`` tuple or ``None``. ``None`` means "no
            history for this layer" (i.e., first call).

    Returns:
        A tuple ``(logits, new_kvs)`` where ``logits`` is
            ``(batch, seq, vocab)`` — the *new* tokens' logits (with
            ``past_kvs != None`` and ``seq=1`` the returned
            ``logits[:, -1, :]`` is the next-token distribution
            conditioned on the full cached prefix) — and ``new_kvs``
            is a list of length ``n_layers`` of updated ``(k, v)``
            tuples; pass this back in for the next step.
    """
    b, t = tokens.shape
    # Length already cached, taken from layer 0 (all layers share length).
    cached_len = 0
    if past_kvs is not None and past_kvs[0] is not None:
        cached_len = past_kvs[0][0].size(-2)
    total_len = cached_len + t
    if total_len > self.params.max_seq_len:
        raise ValueError(
            f"cached_len ({cached_len}) + new tokens ({t}) = {total_len} "
            f"exceeds max_seq_len={self.params.max_seq_len}"
        )

    if past_kvs is None:
        past_kvs = [None] * len(self.blocks)
    if len(past_kvs) != len(self.blocks):
        raise ValueError(f"past_kvs has {len(past_kvs)} entries but model has {len(self.blocks)} layers")

    x = self.tok_embed(tokens)  # (B, T, d_model)
    new_kvs: list[LayerKV] = []
    for block, layer_past in zip(self.blocks, past_kvs, strict=True):
        x, new_kv = block(x, past_kv=layer_past, use_cache=True)
        new_kvs.append(new_kv)
    x = self.norm_out(x)
    logits = self.lm_head(x)  # (B, T, vocab)
    return logits, new_kvs

unpack_batch(batch)

Make TransformerNN compatible with the standard supervised NNModel training loop.

For an LM the canonical batch is (tokens, targets) where targets = tokens[:, 1:] shifted by one. We don't shift here — the caller assembles the tuple — but we accept either a 2-tuple (X, Y) or a plain tensor of tokens (next-token loss is then computed in the train step).

Source code in nnx/nn/net/transformer_nn.py

def unpack_batch(self, batch):
    """Make TransformerNN compatible with the standard supervised
    NNModel training loop.

    For an LM the canonical batch is ``(tokens, targets)`` where
    ``targets = tokens[:, 1:]`` shifted by one. We don't shift here —
    the caller assembles the tuple — but we accept either a 2-tuple
    ``(X, Y)`` or a plain tensor of tokens (next-token loss is then
    computed in the train step).
    """
    if isinstance(batch, (list, tuple)) and len(batch) == 2:
        X, Y = batch
        return (X,), Y
    if isinstance(batch, torch.Tensor):
        # Self-supervised next-token: targets are inputs shifted by 1.
        return (batch[:, :-1],), batch[:, 1:]
    raise TypeError(
        f"TransformerNN.unpack_batch expects a (tokens, targets) tuple "
        f"or a tensor of token ids; got {type(batch).__name__}"
    )

nnx.nn.net.vit_nn.ViTNN

Bases: Module

Small Vision Transformer encoder.

Forward contract:

forward(x: (B, C, H, W), mask: Optional[BoolTensor[B, n_patches]]=None) → (B, T_kept + 1, d_model) if mask provided (T_kept = mask.sum()) → (B, n_patches + 1, d_model) otherwise.

The leading token is the learned CLS. Patches are flattened in raster order (row-major over the patch grid). The optional mask is the I-JEPA "context" mask: True positions are kept, False ones are dropped before any attention runs, so gradients do not flow through masked patches.

__init__ requires image_size, patch_size, and in_channels for the patch-embedding convolution. image_size must be divisible by patch_size — validated at construction.

Source code in nnx/nn/net/vit_nn.py

class ViTNN(nn.Module):
    """Small Vision Transformer encoder.

    Forward contract:

      ``forward(x: (B, C, H, W), mask: Optional[BoolTensor[B, n_patches]]=None)``
      → ``(B, T_kept + 1, d_model)`` if ``mask`` provided (T_kept = mask.sum())
      → ``(B, n_patches + 1, d_model)`` otherwise.

    The leading token is the learned CLS. Patches are flattened in
    raster order (row-major over the patch grid). The optional ``mask``
    is the I-JEPA "context" mask: True positions are kept, False ones
    are dropped before any attention runs, so gradients do not flow
    through masked patches.

    ``__init__`` requires ``image_size``, ``patch_size``, and
    ``in_channels`` for the patch-embedding convolution. ``image_size``
    must be divisible by ``patch_size`` — validated at construction.
    """

    def __init__(
        self,
        *,
        image_size: int = 32,
        patch_size: int = 4,
        in_channels: int = 3,
        d_model: int = 64,
        n_layers: int = 4,
        n_heads: int = 4,
        ffn_mult: int = 4,
        attn_dropout: float = 0.0,
        resid_dropout: float = 0.0,
    ):
        super().__init__()
        # Positive-dimension guards run BEFORE the divisibility checks: a
        # zero/negative value can otherwise mask itself (e.g. d_model=0 passes
        # `0 % n_heads == 0` then zeroes head_dim; image_size=-32 passes
        # `-32 % 4 == 0` then yields a wrong n_patches) and build a silently
        # degenerate model with no error — the same [[params-boundary-validation]]
        # footgun guarded on NNTransformerParams. These do not affect any state().
        for _name, _value in (
            ("image_size", image_size),
            ("patch_size", patch_size),
            ("in_channels", in_channels),
            ("d_model", d_model),
            ("n_layers", n_layers),
            ("n_heads", n_heads),
            ("ffn_mult", ffn_mult),
        ):
            if _value <= 0:
                raise ValueError(f"ViTNN requires {_name} > 0, got {_value}")
        for _name, _value in (("attn_dropout", attn_dropout), ("resid_dropout", resid_dropout)):
            if not 0.0 <= _value <= 1.0:
                raise ValueError(f"ViTNN requires 0.0 <= {_name} <= 1.0, got {_value}")
        if image_size % patch_size != 0:
            raise ValueError(f"image_size={image_size} must be divisible by patch_size={patch_size}")
        if d_model % n_heads != 0:
            raise ValueError(f"d_model={d_model} must be divisible by n_heads={n_heads}")

        self.image_size = image_size
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.d_model = d_model
        self.n_patches = (image_size // patch_size) ** 2

        # Patch embedding: Conv2d with kernel=stride=patch_size produces
        # one (d_model)-vector per non-overlapping patch. (B, C, H, W) ->
        # (B, d_model, H/p, W/p) -> flatten -> (B, n_patches, d_model).
        self.patch_embed = nn.Conv2d(
            in_channels=in_channels,
            out_channels=d_model,
            kernel_size=patch_size,
            stride=patch_size,
        )
        # Learned CLS + position embeddings. One extra position for the
        # CLS token (prepended at the front of the sequence).
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.n_patches + 1, d_model))
        # Init pos_embed / cls_token with small noise so the symmetry
        # at step 0 doesn't lead to degenerate attention patterns.
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)

        self.blocks = nn.ModuleList(
            [
                ViTBlock(
                    d_model=d_model,
                    n_heads=n_heads,
                    ffn_mult=ffn_mult,
                    attn_dropout=attn_dropout,
                    resid_dropout=resid_dropout,
                )
                for _ in range(n_layers)
            ]
        )
        self.norm_out = RMSNorm(d_model)

    def patch_positions(self) -> torch.Tensor:
        """Return ``LongTensor[n_patches]`` of patch-token positions in
        the full sequence (i.e., ``arange(1, n_patches + 1)`` — CLS is
        position 0).

        Exposed so the I-JEPA step factory can derive its context /
        target position indices by boolean-masking this tensor instead
        of rebuilding the arange (see ``jepa_train_step_factory``).
        """
        return torch.arange(1, self.n_patches + 1, device=self.pos_embed.device)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Run the encoder.

        Args:
            x: (B, C, H, W) input image tensor.
            mask: optional BoolTensor of shape ``(B, n_patches)`` —
                True positions are kept, False ones are dropped *before*
                attention. Per-sample masks may have different
                ``True``-counts, but the resulting batch must have the
                same kept-count per row (asserted). I-JEPA's typical
                context mask is uniform across the batch (same set of
                patches kept on every sample in a step).

        Returns:
            ``(B, T_kept + 1, d_model)`` where T_kept is the number of
            kept patches (or ``n_patches`` when mask is None). The +1
            is the CLS token at position 0.
        """
        if x.dim() != 4:
            raise ValueError(f"ViTNN expects (B, C, H, W) input, got shape {tuple(x.shape)}")
        B = x.shape[0]
        # Patch embed -> (B, n_patches, d_model). Add positional embeds
        # to the patch tokens (pos 1..n_patches; index 0 is CLS-only).
        patches = self.patch_embed(x).flatten(2).transpose(1, 2)  # (B, n_patches, d_model)
        patches = patches + self.pos_embed[:, 1 : self.n_patches + 1, :]

        if mask is not None:
            if mask.shape != (B, self.n_patches):
                raise ValueError(f"mask shape {tuple(mask.shape)} does not match expected ({B}, {self.n_patches})")
            # Per-row kept-count must agree so we can stack into a
            # rectangular tensor. JEPA's typical use is a uniform
            # mask anyway, but we validate to fail loudly when callers
            # construct ragged masks.
            kept_counts = mask.sum(dim=1)
            if kept_counts.unique().numel() != 1:
                raise ValueError(
                    f"mask has ragged kept-counts per row {kept_counts.tolist()}; "
                    "ViTNN requires a uniform kept-count across the batch."
                )
            T_kept = int(kept_counts[0].item())
            # Gather kept tokens. ``mask`` is (B, n_patches) bool; result
            # reshapes to (B, T_kept, d_model).
            patches = patches[mask].view(B, T_kept, self.d_model)

        cls = self.cls_token.expand(B, -1, -1) + self.pos_embed[:, :1, :]
        x = torch.cat([cls, patches], dim=1)
        for block in self.blocks:
            x = block(x)
        return self.norm_out(x)

    def unpack_batch(self, batch):
        """Standard ``(X-tuple, Y)`` adapter. JEPA doesn't use Y but
        the supervised linear-probe path on top of a frozen ViTNN does.
        """
        if isinstance(batch, (list, tuple)):
            x = batch[0]
            y = batch[1] if len(batch) > 1 else None
            return (x,), y
        return (batch,), None

    def __str__(self) -> str:
        return (
            f"ViTNN[image={self.image_size}, patch={self.patch_size}, "
            f"d_model={self.d_model}, n_patches={self.n_patches}]"
        )

forward(x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor

Run the encoder.

Parameters:

Name	Type	Description	Default
x	Tensor	(B, C, H, W) input image tensor.	required
mask	Optional[Tensor]	optional BoolTensor of shape (B, n_patches) — True positions are kept, False ones are dropped before attention. Per-sample masks may have different True-counts, but the resulting batch must have the same kept-count per row (asserted). I-JEPA's typical context mask is uniform across the batch (same set of patches kept on every sample in a step).	None

Returns:

Type	Description
Tensor	(B, T_kept + 1, d_model) where T_kept is the number of
Tensor	kept patches (or n_patches when mask is None). The +1
Tensor	is the CLS token at position 0.

Source code in nnx/nn/net/vit_nn.py

def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Run the encoder.

    Args:
        x: (B, C, H, W) input image tensor.
        mask: optional BoolTensor of shape ``(B, n_patches)`` —
            True positions are kept, False ones are dropped *before*
            attention. Per-sample masks may have different
            ``True``-counts, but the resulting batch must have the
            same kept-count per row (asserted). I-JEPA's typical
            context mask is uniform across the batch (same set of
            patches kept on every sample in a step).

    Returns:
        ``(B, T_kept + 1, d_model)`` where T_kept is the number of
        kept patches (or ``n_patches`` when mask is None). The +1
        is the CLS token at position 0.
    """
    if x.dim() != 4:
        raise ValueError(f"ViTNN expects (B, C, H, W) input, got shape {tuple(x.shape)}")
    B = x.shape[0]
    # Patch embed -> (B, n_patches, d_model). Add positional embeds
    # to the patch tokens (pos 1..n_patches; index 0 is CLS-only).
    patches = self.patch_embed(x).flatten(2).transpose(1, 2)  # (B, n_patches, d_model)
    patches = patches + self.pos_embed[:, 1 : self.n_patches + 1, :]

    if mask is not None:
        if mask.shape != (B, self.n_patches):
            raise ValueError(f"mask shape {tuple(mask.shape)} does not match expected ({B}, {self.n_patches})")
        # Per-row kept-count must agree so we can stack into a
        # rectangular tensor. JEPA's typical use is a uniform
        # mask anyway, but we validate to fail loudly when callers
        # construct ragged masks.
        kept_counts = mask.sum(dim=1)
        if kept_counts.unique().numel() != 1:
            raise ValueError(
                f"mask has ragged kept-counts per row {kept_counts.tolist()}; "
                "ViTNN requires a uniform kept-count across the batch."
            )
        T_kept = int(kept_counts[0].item())
        # Gather kept tokens. ``mask`` is (B, n_patches) bool; result
        # reshapes to (B, T_kept, d_model).
        patches = patches[mask].view(B, T_kept, self.d_model)

    cls = self.cls_token.expand(B, -1, -1) + self.pos_embed[:, :1, :]
    x = torch.cat([cls, patches], dim=1)
    for block in self.blocks:
        x = block(x)
    return self.norm_out(x)

patch_positions() -> torch.Tensor

Return LongTensor[n_patches] of patch-token positions in the full sequence (i.e., arange(1, n_patches + 1) — CLS is position 0).

Exposed so the I-JEPA step factory can derive its context / target position indices by boolean-masking this tensor instead of rebuilding the arange (see jepa_train_step_factory).

Source code in nnx/nn/net/vit_nn.py

def patch_positions(self) -> torch.Tensor:
    """Return ``LongTensor[n_patches]`` of patch-token positions in
    the full sequence (i.e., ``arange(1, n_patches + 1)`` — CLS is
    position 0).

    Exposed so the I-JEPA step factory can derive its context /
    target position indices by boolean-masking this tensor instead
    of rebuilding the arange (see ``jepa_train_step_factory``).
    """
    return torch.arange(1, self.n_patches + 1, device=self.pos_embed.device)

unpack_batch(batch)

Standard (X-tuple, Y) adapter. JEPA doesn't use Y but the supervised linear-probe path on top of a frozen ViTNN does.

Source code in nnx/nn/net/vit_nn.py

def unpack_batch(self, batch):
    """Standard ``(X-tuple, Y)`` adapter. JEPA doesn't use Y but
    the supervised linear-probe path on top of a frozen ViTNN does.
    """
    if isinstance(batch, (list, tuple)):
        x = batch[0]
        y = batch[1] if len(batch) > 1 else None
        return (x,), y
    return (batch,), None

nnx.nn.net.vit_nn.ViTBlock

Bases: Module

Pre-norm ViT block: x = x + attn(RMSNorm(x)); x = x + ffn(RMSNorm(x)).

Same shape as :class:nnx.nn.net.transformer_layers.TransformerBlock but with bidirectional attention instead of causal. SwiGLU is reused unchanged.

Source code in nnx/nn/net/vit_nn.py

class ViTBlock(nn.Module):
    """Pre-norm ViT block: ``x = x + attn(RMSNorm(x)); x = x + ffn(RMSNorm(x))``.

    Same shape as :class:`nnx.nn.net.transformer_layers.TransformerBlock`
    but with bidirectional attention instead of causal. SwiGLU is
    reused unchanged.
    """

    def __init__(
        self,
        d_model: int,
        n_heads: int,
        ffn_mult: int = 4,
        attn_dropout: float = 0.0,
        resid_dropout: float = 0.0,
    ):
        super().__init__()
        # ViTBlock is a public top-level export (nnx.ViTBlock), so it owns its
        # own positive-dimension guards rather than relying on ViTNN's. ffn_mult=0
        # would otherwise build a zero-width SwiGLU FFN silently; d_model<=0 /
        # n_heads<=0 a zero attention. Matches the ViTNN guard above.
        for _name, _value in (("d_model", d_model), ("n_heads", n_heads), ("ffn_mult", ffn_mult)):
            if _value <= 0:
                raise ValueError(f"ViTBlock requires {_name} > 0, got {_value}")
        for _name, _value in (("attn_dropout", attn_dropout), ("resid_dropout", resid_dropout)):
            if not 0.0 <= _value <= 1.0:
                raise ValueError(f"ViTBlock requires 0.0 <= {_name} <= 1.0, got {_value}")
        self.norm1 = RMSNorm(d_model)
        self.attn = _MultiHeadSelfAttention(d_model, n_heads, attn_dropout=attn_dropout)
        self.norm2 = RMSNorm(d_model)
        self.ffn = SwiGLU(d_model=d_model, ffn_mult=ffn_mult)
        self.resid_dropout = resid_dropout

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + F.dropout(self.attn(self.norm1(x)), p=self.resid_dropout, training=self.training)
        x = x + F.dropout(self.ffn(self.norm2(x)), p=self.resid_dropout, training=self.training)
        return x

nnx.nn.moe.MoELinear

Bases: Module

Sparse top-k Mixture-of-Experts drop-in for :class:nn.Linear.

Forward pass:

1. Router (a bias-less :class:nn.Linear) projects input (B, in_features) → (B, num_experts) logits.
2. top_k largest logits per row are kept; a softmax over those k values produces the per-expert gating weight.
3. Each token is dispatched to its top-k experts; expert outputs are weighted by the gating weights and summed into the output tensor.
4. self.last_aux_loss is populated with the Switch-style load-balancing penalty num_experts · Σ_i f_i · P_i. This is a scalar tensor with gradients wired to the router so optimization of the main loss + this term pushes routing toward uniform expert usage.

Parameters:

Name	Type	Description	Default
in_features	int	input feature dimension (matches nn.Linear).	required
out_features	int	output feature dimension (matches nn.Linear).	required
num_experts	int	number of expert sub-networks. Must be ≥ 2. (num_experts=1 collapses to a plain linear with extra book-keeping; the layer rejects it to surface the misuse.)	required
top_k	int	number of experts each input is routed through. Must be ≥ 1 and ≤ num_experts. Defaults to 2 — the Switch-Transformer paper uses k=1, but k=2 is the broader MoE convention and tolerates a single misrouted expert without losing the entire token.	2

Attributes:

Name	Type	Description
router		bias-less :class:nn.Linear of shape (in_features, num_experts).
experts		:class:nn.ModuleList of num_experts :class:nn.Linear layers, each (in_features, out_features).
top_k		how many experts run per token.
num_experts		total expert count.
last_aux_loss	Tensor | None	scalar torch.Tensor set after each :meth:forward. None before the first forward.

Raises:

Type	Description
ValueError	if in_features <= 0, out_features <= 0, num_experts <= 1, top_k <= 0, or top_k > num_experts.

Source code in nnx/nn/moe.py

class MoELinear(nn.Module):
    """Sparse top-k Mixture-of-Experts drop-in for :class:`nn.Linear`.

    Forward pass:

      1. Router (a bias-less :class:`nn.Linear`) projects input
         ``(B, in_features) → (B, num_experts)`` logits.
      2. ``top_k`` largest logits per row are kept; a softmax over
         those ``k`` values produces the per-expert gating weight.
      3. Each token is dispatched to its top-``k`` experts; expert
         outputs are weighted by the gating weights and summed into
         the output tensor.
      4. ``self.last_aux_loss`` is populated with the Switch-style
         load-balancing penalty
         ``num_experts · Σ_i f_i · P_i``. This is a scalar tensor with
         gradients wired to the router so optimization of the main
         loss + this term pushes routing toward uniform expert usage.

    Args:
        in_features: input feature dimension (matches ``nn.Linear``).
        out_features: output feature dimension (matches ``nn.Linear``).
        num_experts: number of expert sub-networks. Must be ≥ 2.
            (``num_experts=1`` collapses to a plain linear with extra
            book-keeping; the layer rejects it to surface the misuse.)
        top_k: number of experts each input is routed through. Must
            be ≥ 1 and ≤ ``num_experts``. Defaults to 2 — the
            Switch-Transformer paper uses ``k=1``, but ``k=2`` is the
            broader MoE convention and tolerates a single misrouted
            expert without losing the entire token.

    Attributes:
        router: bias-less :class:`nn.Linear` of shape
            ``(in_features, num_experts)``.
        experts: :class:`nn.ModuleList` of ``num_experts``
            :class:`nn.Linear` layers, each ``(in_features, out_features)``.
        top_k: how many experts run per token.
        num_experts: total expert count.
        last_aux_loss: scalar ``torch.Tensor`` set after each
            :meth:`forward`. ``None`` before the first forward.

    Raises:
        ValueError: if ``in_features <= 0``, ``out_features <= 0``,
            ``num_experts <= 1``, ``top_k <= 0``, or ``top_k > num_experts``.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        *,
        num_experts: int,
        top_k: int = 2,
    ):
        super().__init__()
        if in_features <= 0:
            raise ValueError(f"in_features must be positive, got {in_features}")
        if out_features <= 0:
            raise ValueError(f"out_features must be positive, got {out_features}")
        if num_experts <= 1:
            raise ValueError(
                f"num_experts must be ≥ 2 (a single-expert MoELinear is just "
                f"an nn.Linear with extra book-keeping), got {num_experts}"
            )
        if top_k <= 0:
            raise ValueError(f"top_k must be positive, got {top_k}")
        if top_k > num_experts:
            raise ValueError(
                f"top_k ({top_k}) cannot exceed num_experts ({num_experts}) — "
                "a token cannot be dispatched to more experts than exist."
            )

        self.in_features = in_features
        self.out_features = out_features
        self.num_experts = num_experts
        self.top_k = top_k

        # Router has no bias by convention — the additive offset would
        # bias the softmax toward fixed experts regardless of input,
        # the opposite of what we want for load balancing.
        self.router = nn.Linear(in_features, num_experts, bias=False)
        self.experts = nn.ModuleList([nn.Linear(in_features, out_features) for _ in range(num_experts)])

        # Populated by the first forward. Typed as Optional so static
        # checkers see the pre-forward state correctly.
        self.last_aux_loss: torch.Tensor | None = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # nn.Linear contract: accept (..., in_features) — including the
        # unbatched 1-D case. The routing / dispatch logic below is
        # written for 2-D input, so flatten/promote the leading dims
        # here and restore them on return — a (B, T, C) sequence batch
        # routes each token independently, which is the standard
        # MoE-in-transformer semantics.
        lead_shape = x.shape[:-1]
        if x.dim() != 2:
            x = x.reshape(-1, self.in_features)

        # Router logits: (B, num_experts). The full softmax across
        # all experts is used in the load-balancing penalty (Switch's
        # P_i term needs ALL experts' probability mass, not just the
        # top-k subset).
        logits = self.router(x)
        probs_full = logits.softmax(dim=-1)  # (B, num_experts)

        # Top-k routing: select the k largest logits per row, then
        # renormalize across the k via softmax so the gating weights
        # sum to 1 per token. (Re-softmaxing the k logits — rather
        # than re-normalizing the full probs — is the original Switch
        # formulation; it gives sharper weights when the routed
        # experts have similar pre-top-k logits.)
        topk_vals, topk_idx = logits.topk(self.top_k, dim=-1)  # both (B, top_k)
        gate_weights = topk_vals.softmax(dim=-1)  # (B, top_k)

        B = x.size(0)
        out = torch.zeros(B, self.out_features, device=x.device, dtype=x.dtype)

        # Dispatch loop: for each of the top_k slots and each expert,
        # pull the rows whose slot routes to that expert, run them
        # through the expert, and scatter back into `out` weighted by
        # the gate. This is O(top_k · num_experts) Python iterations
        # — fine at tutorial scale, replaced by block-sparse kernels
        # at production scale.
        for k in range(self.top_k):
            for e in range(self.num_experts):
                mask = topk_idx[:, k] == e
                if mask.any():
                    expert_out = self.experts[e](x[mask])
                    out[mask] = out[mask] + gate_weights[mask, k : k + 1] * expert_out

        # Switch-style load-balancing aux loss
        # ``α · N · Σ_i f_i · P_i`` where:
        #   f_i = fraction of dispatched (token, slot) pairs routed
        #         to expert i, averaged over batch and top_k slots.
        #   P_i = mean router probability for expert i over the batch.
        # Both terms are in [0, 1] and sum to 1 across i. The product
        # is minimized at uniform f and P; at uniform routing each
        # equals 1/N, the sum equals N · 1/N² = 1/N, and the whole
        # penalty equals 1. So the minimum value of this loss term
        # is 1, not 0 — that's the property tested by
        # ``test_moe_linear_aux_loss_zero_at_uniform``.
        #
        # Broadcasting: topk_idx is (B, top_k); the [..., None]
        # expansion yields (B, top_k, 1), comparing against
        # arange[None, None, :] of shape (1, 1, num_experts) gives a
        # boolean tensor of shape (B, top_k, num_experts). Mean over
        # the first two dims collapses to (num_experts,).
        expert_idx = torch.arange(self.num_experts, device=x.device)
        dispatch_mask = topk_idx[..., None] == expert_idx[None, None, :]
        dispatch_frac = dispatch_mask.float().mean(dim=(0, 1))  # (num_experts,)
        mean_prob = probs_full.mean(dim=0)  # (num_experts,)
        self.last_aux_loss = self.num_experts * (dispatch_frac * mean_prob).sum()

        if len(lead_shape) != 1:
            out = out.reshape(*lead_shape, self.out_features)
        return out

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"num_experts={self.num_experts}, top_k={self.top_k}"
        )

5. Datasets

nnx.nn.dataset.nn_dataset_base.NNDatasetBase dataclass

Bases: ABC

Source code in nnx/nn/dataset/nn_dataset_base.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNDatasetBase(ABC):
    name: str = field(init=False)

    input_dim: int = field(init=False)
    output_dim: int = field(init=False)

    train_loader: DatasetProtocol = field(init=False)
    # val/test may be None: the tabular and preference subclasses set
    # them to None for empty splits (their documented contract).
    val_loader: Optional[DatasetProtocol] = field(init=False)
    test_loader: Optional[DatasetProtocol] = field(init=False)

    _state: dict = field(repr=False, init=False)

    def state(self) -> dict:
        return self._state

nnx.nn.dataset.nn_dataset.NNDataset dataclass

Bases: NNDatasetBase

Vision dataset wrapper. val_proportion carves a validation slice out of the source train=True split (NOT out of the test split, which stays untouched for final evaluation).

Source code in nnx/nn/dataset/nn_dataset.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNDataset(NNDatasetBase):
    """Vision dataset wrapper. `val_proportion` carves a validation slice
    out of the source `train=True` split (NOT out of the test split, which
    stays untouched for final evaluation)."""

    ds_class: type[VisionDataset]
    root_dir: str = "./data"
    download: bool = True
    transform: Optional[Callable] = None
    # Per-split batch size. None for any entry means "use the full split as
    # one batch" (resolved in __post_init__ once the split sizes are known).
    batch_sizes: tuple[Optional[int], Optional[int], Optional[int]] = (None, None, None)
    val_proportion: float = 0.1
    # Deterministic split when set — same `seed` + same `val_proportion`
    # round-trips to the same train/val ids across runs. Default None falls
    # back to the global torch RNG (the pre-fix behavior). Mirrors the
    # NNPreferenceDataset contract that the seeded-split family already used.
    seed: Optional[int] = None

    def __post_init__(self):
        if not 0.0 <= self.val_proportion < 1.0:
            raise ValueError(f"val_proportion must be in [0, 1), got {self.val_proportion}")
        full_train_dataset, test_dataset = (
            self.ds_class(root=self.root_dir, train=True, download=self.download, transform=self.transform),
            self.ds_class(root=self.root_dir, train=False, download=self.download, transform=self.transform),
        )

        # Carve val out of train so the test set stays held-out for final eval.
        # Compute val_size first, derive train_size as the remainder so the two
        # sum exactly to len(full_train_dataset) (int truncation safe).
        full_train_len = len(full_train_dataset)
        val_size = int(full_train_len * self.val_proportion)
        train_size = full_train_len - val_size
        # seed=None must genuinely fall back to the global torch RNG (the
        # documented contract): a fresh torch.Generator() is NOT that —
        # it always carries the same fixed default seed, which would make
        # every unseeded split bit-identical and deaf to torch.manual_seed.
        gen = torch.Generator().manual_seed(int(self.seed)) if self.seed is not None else torch.default_generator
        train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size], generator=gen)

        object.__setattr__(self, "name", self.ds_class.__name__)

        # Fail fast on the no-transform PIL case: torchvision datasets
        # yield PIL Images without `transform`, and everything downstream
        # (input_dim inference, batching) needs tensors.
        sample = full_train_dataset[0][0]
        if not hasattr(sample, "shape"):
            raise ValueError(
                f"{self.ds_class.__name__} samples have no .shape (got {type(sample).__name__}) — "
                "pass transform=torchvision.transforms.ToTensor() (or a pipeline ending in it)."
            )

        train_batch_size = self.batch_sizes[0] or len(train_dataset)
        # max(1, ...): an empty split would otherwise resolve to
        # DataLoader(batch_size=0), which raises.
        val_batch_size = self.batch_sizes[1] or max(1, len(val_dataset))
        test_batch_size = self.batch_sizes[2] or max(1, len(test_dataset))
        resolved_batch_sizes = (train_batch_size, val_batch_size, test_batch_size)

        object.__setattr__(self, "batch_sizes", resolved_batch_sizes)

        object.__setattr__(
            self, "train_loader", DataLoader(shuffle=True, dataset=train_dataset, batch_size=resolved_batch_sizes[0])
        )

        # val_proportion=0.0 → val_loader=None, matching the tabular /
        # preference siblings' documented empty-split contract (an empty
        # DataLoader would instead make train() run a zero-sample
        # validate pass and crash at the end of the first epoch).
        object.__setattr__(
            self,
            "val_loader",
            DataLoader(shuffle=False, dataset=val_dataset, batch_size=resolved_batch_sizes[1])
            if len(val_dataset) > 0
            else None,
        )

        object.__setattr__(
            self, "test_loader", DataLoader(shuffle=False, dataset=test_dataset, batch_size=resolved_batch_sizes[2])
        )

        # train_loader.dataset is now a Subset (from random_split); shape/classes
        # come from the underlying full_train_dataset instead.
        object.__setattr__(self, "input_dim", reduce(lambda x, y: x * y, sample.shape))

        object.__setattr__(self, "output_dim", len(full_train_dataset.classes))

        state = dict(
            name=self.name,
            input_dim=self.input_dim,
            output_dim=self.output_dim,
            train_batch_size=f"{self.batch_sizes[0]:,}",
            # 0 when the val split is empty (val_loader is None) — the
            # resolved batch_sizes carry a placeholder 1 there.
            val_batch_size=f"{self.batch_sizes[1] if len(val_dataset) > 0 else 0:,}",
            test_batch_size=f"{self.batch_sizes[2]:,}",
        )

        object.__setattr__(self, "_state", state)

nnx.nn.dataset.nn_graph_dataset.NNGraphDataset dataclass

Bases: NNDatasetBase

Source code in nnx/nn/dataset/nn_graph_dataset.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNGraphDataset(NNDatasetBase):
    ds_class: type[Dataset]
    # n_neighbors is required for sampler="neighbor"; unused for "full".
    # Kept Optional so that full-batch callers need not supply a meaningless list.
    n_neighbors: Optional[list[int]] = None
    root_dir: str = "./data"
    transform: Optional[Callable] = None
    n_workers: int = 4
    # Per-split batch size. None for any entry means "use every node in the
    # split mask" (resolved in __post_init__ from the train/val/test masks).
    batch_sizes: tuple[Optional[int], Optional[int], Optional[int]] = (None, None, None)
    # Deterministic neighbor sampling when set — the train loader shuffles
    # (RandomSampler reads `generator`) and the default `n_workers=4` spawns
    # worker processes whose numpy/python RNG must be pinned via
    # `dataloader_worker_init_fn`. Default None falls back to the global torch
    # RNG (the pre-fix behavior). Mirrors the seed contract the other datasets
    # (NNDataset / NNTabularDataset / NNPreferenceDataset) already expose.
    # Ignored for sampler="full" (no sampling randomness to control).
    seed: Optional[int] = None
    # "neighbor": NeighborLoader-based mini-batch sampling (default, today's
    #   behavior).  Requires n_neighbors.
    # "full": one-batch-per-epoch loader that carries the entire graph
    #   permuted so the split's seed nodes lead.  Does NOT require pyg-lib
    #   or torch-sparse — works on Apple Silicon without pre-built wheels.
    sampler: Literal["neighbor", "full"] = "neighbor"

    def __post_init__(self):
        if self.sampler == "neighbor" and self.n_neighbors is None:
            raise ValueError(
                "n_neighbors is required when sampler='neighbor'. "
                "Pass n_neighbors=[k1, k2, ...] or switch to sampler='full'."
            )

        dataset = self.ds_class(root=self.root_dir, transform=self.transform)
        # Single-graph datasets expose the underlying Data via dataset[0].
        # This replaces the historical private `dataset._data` access, which
        # was renamed/removed across PyG versions.
        data = dataset[0]

        object.__setattr__(self, "name", self.ds_class.__name__)

        train_batch_size = self.batch_sizes[0] or int(data.train_mask.sum())
        val_batch_size = self.batch_sizes[1] or int(data.val_mask.sum())
        test_batch_size = self.batch_sizes[2] or int(data.test_mask.sum())
        resolved_batch_sizes = (train_batch_size, val_batch_size, test_batch_size)

        object.__setattr__(self, "batch_sizes", resolved_batch_sizes)

        if self.sampler == "full":
            object.__setattr__(self, "train_loader", _full_batch_loader(data, data.train_mask))
            object.__setattr__(self, "val_loader", _full_batch_loader(data, data.val_mask))
            object.__setattr__(self, "test_loader", _full_batch_loader(data, data.test_mask))
        else:
            # seed=None must genuinely fall back to the global torch RNG (the
            # documented contract): a fresh torch.Generator() always carries the
            # same fixed default seed, which would make every unseeded run
            # bit-identical and deaf to torch.manual_seed.
            gen = torch.Generator().manual_seed(int(self.seed)) if self.seed is not None else torch.default_generator

            object.__setattr__(
                self,
                "train_loader",
                NeighborLoader(
                    shuffle=True,
                    data=data,
                    num_workers=self.n_workers,
                    num_neighbors=self.n_neighbors,
                    batch_size=resolved_batch_sizes[0],
                    input_nodes=data.train_mask,
                    generator=gen,
                    worker_init_fn=dataloader_worker_init_fn,
                ),
            )

            object.__setattr__(
                self,
                "val_loader",
                NeighborLoader(
                    shuffle=False,
                    data=data,
                    num_workers=self.n_workers,
                    num_neighbors=self.n_neighbors,
                    batch_size=resolved_batch_sizes[1],
                    input_nodes=data.val_mask,
                    generator=gen,
                    worker_init_fn=dataloader_worker_init_fn,
                ),
            )

            object.__setattr__(
                self,
                "test_loader",
                NeighborLoader(
                    shuffle=False,
                    data=data,
                    num_workers=self.n_workers,
                    num_neighbors=self.n_neighbors,
                    batch_size=resolved_batch_sizes[2],
                    input_nodes=data.test_mask,
                    generator=gen,
                    worker_init_fn=dataloader_worker_init_fn,
                ),
            )

        object.__setattr__(self, "input_dim", dataset.num_features)

        object.__setattr__(self, "output_dim", dataset.num_classes)

        state: dict = dict(
            name=self.name,
            sampler=self.sampler,
            input_dim=self.input_dim,
            output_dim=self.output_dim,
            train_batch_size=f"{self.batch_sizes[0]:,}",
            val_batch_size=f"{self.batch_sizes[1]:,}",
            test_batch_size=f"{self.batch_sizes[2]:,}",
        )
        if self.sampler == "neighbor":
            state["n_workers"] = self.n_workers
            state["n_neighbors"] = self.n_neighbors

        object.__setattr__(self, "_state", state)

nnx.nn.dataset.nn_tabular_dataset.NNTabularDataset dataclass

Bases: NNDatasetBase

Wrap a pandas DataFrame as train/val/test DataLoaders.

feature_cols columns are stacked into the input tensor; target_col is the integer label column. Targets are coerced to int64 (long), so use this for classification. For regression, prefer to construct the DataLoaders yourself and pass them through NNTrainParams.

Source code in nnx/nn/dataset/nn_tabular_dataset.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNTabularDataset(NNDatasetBase):
    """Wrap a pandas DataFrame as train/val/test DataLoaders.

    `feature_cols` columns are stacked into the input tensor; `target_col`
    is the integer label column. Targets are coerced to int64 (long), so
    use this for classification. For regression, prefer to construct the
    DataLoaders yourself and pass them through NNTrainParams.
    """

    df: pd.DataFrame
    feature_cols: list[str]
    target_col: str

    # Per-split batch size. None for any entry means "use the full split as
    # one batch" (resolved in __post_init__ once the split sizes are known).
    batch_sizes: tuple[Optional[int], Optional[int], Optional[int]] = (None, None, None)
    val_proportion: float = 0.15
    test_proportion: float = 0.15
    name_override: Optional[str] = None
    feature_dtype: torch.dtype = field(default=torch.float32)
    # Deterministic split when set — same `seed` + same `val_proportion` /
    # `test_proportion` round-trips to the same train/val/test ids across
    # runs. Default None falls back to the global torch RNG (the pre-fix
    # behavior). Mirrors NNPreferenceDataset's seeded-split contract.
    seed: Optional[int] = None

    def __post_init__(self):
        if not 0.0 <= self.val_proportion < 1.0:
            raise ValueError(f"val_proportion must be in [0, 1), got {self.val_proportion}")
        if not 0.0 <= self.test_proportion < 1.0:
            raise ValueError(f"test_proportion must be in [0, 1), got {self.test_proportion}")
        if self.val_proportion + self.test_proportion >= 1.0:
            raise ValueError(
                f"val_proportion + test_proportion must be < 1, got {self.val_proportion + self.test_proportion}"
            )

        # Validate columns up-front so missing-column errors point at user
        # input rather than failing deep inside torch.tensor with a KeyError.
        missing_features = [c for c in self.feature_cols if c not in self.df.columns]
        if missing_features:
            raise KeyError(f"NNTabularDataset feature_cols not in DataFrame: {missing_features}")
        if self.target_col not in self.df.columns:
            raise KeyError(f"NNTabularDataset target_col {self.target_col!r} not in DataFrame")
        if self.target_col in self.feature_cols:
            # Silent label leakage: the model would train on its own
            # target as an input feature and report near-perfect val
            # accuracy (classic feature_cols=list(df.columns) mistake).
            raise ValueError(
                f"target_col {self.target_col!r} must not appear in feature_cols — that trains on the label."
            )

        # NaN anywhere in the modeled columns is silent poison: NaN
        # features flow into NaN losses, and a NaN target's float→int64
        # cast is UNDEFINED (class 0 on ARM, INT64_MIN on x86 → CUDA
        # device assert) — and the contiguity check below can't see it
        # because pandas min/max/nunique skip NaN.
        # dict.fromkeys dedupes (order-preserving) duplicates WITHIN
        # feature_cols — target/feature overlap is rejected above.
        modeled = self.df[list(dict.fromkeys([*self.feature_cols, self.target_col]))]
        if modeled.isna().any().any():
            bad_cols = [c for c in modeled.columns if modeled[c].isna().any()]
            raise ValueError(
                f"NaN values in columns {bad_cols} — drop or impute rows before constructing NNTabularDataset."
            )

        # Coerce features + target → tensors. Trust `dtype=` on torch.tensor
        # rather than going through an extra np.float32 intermediate copy.
        X = torch.tensor(
            self.df[self.feature_cols].to_numpy(),
            dtype=self.feature_dtype,
        )
        y = torch.tensor(
            self.df[self.target_col].to_numpy(),
            dtype=torch.long,
        )
        n_total = len(X)
        if n_total == 0:
            raise ValueError("NNTabularDataset requires a non-empty DataFrame")

        full_dataset = TensorDataset(X, y)

        # Sizes computed as (n_total - val - test, val, test) so the three
        # sum exactly even with int truncation.
        n_val = int(n_total * self.val_proportion)
        n_test = int(n_total * self.test_proportion)
        n_train = n_total - n_val - n_test
        # seed=None must genuinely fall back to the global torch RNG (the
        # documented contract): a fresh torch.Generator() is NOT that —
        # it always carries the same fixed default seed, which would make
        # every unseeded split bit-identical and deaf to torch.manual_seed.
        gen = torch.Generator().manual_seed(int(self.seed)) if self.seed is not None else torch.default_generator
        train_ds, val_ds, test_ds = random_split(full_dataset, [n_train, n_val, n_test], generator=gen)

        object.__setattr__(self, "name", self.name_override or "NNTabularDataset")

        train_batch_size = self.batch_sizes[0] or n_train
        val_batch_size = self.batch_sizes[1] or max(1, n_val)
        test_batch_size = self.batch_sizes[2] or max(1, n_test)
        resolved_batch_sizes = (train_batch_size, val_batch_size, test_batch_size)
        object.__setattr__(self, "batch_sizes", resolved_batch_sizes)

        object.__setattr__(
            self,
            "train_loader",
            DataLoader(train_ds, batch_size=resolved_batch_sizes[0], shuffle=True),
        )
        # Val/test loaders default to size 0 when the proportion is zero;
        # skip constructing them in that case and use None so callers can
        # check `ds.val_loader is None`.
        object.__setattr__(
            self,
            "val_loader",
            DataLoader(val_ds, batch_size=resolved_batch_sizes[1], shuffle=False) if n_val > 0 else None,
        )
        object.__setattr__(
            self,
            "test_loader",
            DataLoader(test_ds, batch_size=resolved_batch_sizes[2], shuffle=False) if n_test > 0 else None,
        )

        object.__setattr__(self, "input_dim", len(self.feature_cols))
        # Classification target dim = number of unique classes in the DF.
        # Labels must be contiguous 0..K-1: nunique() on e.g. {0, 5}
        # would size the model at 2 outputs and the mismatch only
        # surfaces much later inside cross-entropy as an opaque index /
        # device-side assert error. Fail fast with a fixable message.
        n_classes = int(self.df[self.target_col].nunique())
        target_min = int(self.df[self.target_col].min())
        target_max = int(self.df[self.target_col].max())
        if target_min != 0 or target_max != n_classes - 1:
            raise ValueError(
                f"target_col {self.target_col!r} labels must be contiguous integers 0..K-1; "
                f"got min={target_min}, max={target_max}, n_unique={n_classes}. "
                "Remap labels (e.g. pd.factorize) before constructing the dataset."
            )
        object.__setattr__(self, "output_dim", n_classes)

        object.__setattr__(
            self,
            "_state",
            dict(
                name=self.name,
                input_dim=self.input_dim,
                output_dim=self.output_dim,
                n_train=n_train,
                n_val=n_val,
                n_test=n_test,
                feature_cols=list(self.feature_cols),
                target_col=self.target_col,
            ),
        )

nnx.nn.dataset.nn_preference_dataset.NNPreferenceDataset dataclass

Bases: NNDatasetBase

Wrap parallel lists of (prompt, chosen, rejected) strings as DPO loaders.

Tokenizes every triple through tokenizer.encode once at construction, pads/truncates to fixed lengths, then splits into train / val / test DataLoader\ s with the same shape as the rest of :class:NNDatasetBase (so callbacks and the standard training loop work unchanged).

Each batch yielded is (prompt_ids, chosen_ids, rejected_ids) where each entry is (B, T_*) torch.LongTensor.

Source code in nnx/nn/dataset/nn_preference_dataset.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNPreferenceDataset(NNDatasetBase):
    """Wrap parallel lists of (prompt, chosen, rejected) strings as DPO loaders.

    Tokenizes every triple through ``tokenizer.encode`` once at
    construction, pads/truncates to fixed lengths, then splits into
    train / val / test ``DataLoader``\\ s with the same shape as the
    rest of :class:`NNDatasetBase` (so callbacks and the standard
    training loop work unchanged).

    Each batch yielded is ``(prompt_ids, chosen_ids, rejected_ids)``
    where each entry is ``(B, T_*)`` ``torch.LongTensor``.
    """

    prompts: list[str]
    chosen: list[str]
    rejected: list[str]
    tokenizer: object  # NNTokenizerParams — typed `object` to keep the lm extra optional

    max_prompt_len: int = 64
    max_response_len: int = 64
    pad_token_id: int = 0

    batch_sizes: tuple[Optional[int], Optional[int], Optional[int]] = (None, None, None)
    val_proportion: float = 0.1
    test_proportion: float = 0.1
    name_override: Optional[str] = None
    seed: Optional[int] = field(default=None)

    def __post_init__(self):
        if not 0.0 <= self.val_proportion < 1.0:
            raise ValueError(f"val_proportion must be in [0, 1), got {self.val_proportion}")
        if not 0.0 <= self.test_proportion < 1.0:
            raise ValueError(f"test_proportion must be in [0, 1), got {self.test_proportion}")
        if self.val_proportion + self.test_proportion >= 1.0:
            raise ValueError(
                f"val_proportion + test_proportion must be < 1, got {self.val_proportion + self.test_proportion}"
            )
        if not (len(self.prompts) == len(self.chosen) == len(self.rejected)):
            raise ValueError(
                "NNPreferenceDataset: prompts, chosen, rejected must align — got "
                f"{len(self.prompts)} / {len(self.chosen)} / {len(self.rejected)}"
            )
        if len(self.prompts) == 0:
            raise ValueError("NNPreferenceDataset requires non-empty input lists")

        # Encode everything via the tokenizer up-front.
        encode = self.tokenizer.encode  # type: ignore[attr-defined]
        p_ids = [encode(p) for p in self.prompts]
        c_ids = [encode(c) for c in self.chosen]
        r_ids = [encode(r) for r in self.rejected]

        full = _PreferenceCorpus(
            prompt_ids=p_ids,
            chosen_ids=c_ids,
            rejected_ids=r_ids,
            max_prompt_len=self.max_prompt_len,
            max_response_len=self.max_response_len,
            pad_token_id=self.pad_token_id,
        )

        n_total = len(full)
        n_val = int(n_total * self.val_proportion)
        n_test = int(n_total * self.test_proportion)
        n_train = n_total - n_val - n_test
        if n_train <= 0:
            raise ValueError(
                f"NNPreferenceDataset: not enough rows ({n_total}) for the requested "
                f"val/test split ({self.val_proportion}, {self.test_proportion})"
            )

        # Deterministic split when seed is given — matches NNTabularDataset's
        # implicit contract that the same inputs round-trip to the same
        # train/val/test ids when seeded externally.
        # seed=None must genuinely fall back to the global torch RNG: a
        # fresh torch.Generator() always carries the same fixed default
        # seed, which would make unseeded splits bit-identical across
        # runs and deaf to torch.manual_seed.
        gen = torch.Generator().manual_seed(int(self.seed)) if self.seed is not None else torch.default_generator
        train_ds, val_ds, test_ds = random_split(full, [n_train, n_val, n_test], generator=gen)

        object.__setattr__(self, "name", self.name_override or "NNPreferenceDataset")

        train_batch_size = self.batch_sizes[0] or n_train
        val_batch_size = self.batch_sizes[1] or max(1, n_val)
        test_batch_size = self.batch_sizes[2] or max(1, n_test)
        object.__setattr__(self, "batch_sizes", (train_batch_size, val_batch_size, test_batch_size))

        object.__setattr__(
            self,
            "train_loader",
            DataLoader(train_ds, batch_size=train_batch_size, shuffle=True),
        )
        object.__setattr__(
            self,
            "val_loader",
            DataLoader(val_ds, batch_size=val_batch_size, shuffle=False) if n_val > 0 else None,
        )
        object.__setattr__(
            self,
            "test_loader",
            DataLoader(test_ds, batch_size=test_batch_size, shuffle=False) if n_test > 0 else None,
        )

        # input_dim / output_dim: there's no "feature dimension" for a
        # raw text-pair dataset — report the tokenizer's vocab size so
        # downstream sizing checks have a number to look at.
        vocab_size = int(self.tokenizer.vocab_size)  # type: ignore[attr-defined]
        object.__setattr__(self, "input_dim", vocab_size)
        object.__setattr__(self, "output_dim", vocab_size)

        object.__setattr__(
            self,
            "_state",
            dict(
                name=self.name,
                vocab_size=vocab_size,
                n_train=n_train,
                n_val=n_val,
                n_test=n_test,
                max_prompt_len=self.max_prompt_len,
                max_response_len=self.max_response_len,
                pad_token_id=self.pad_token_id,
            ),
        )

6. Enums

nnx.nn.enum.activations.Activations

Bases: Enum

Source code in nnx/nn/enum/activations.py

class Activations(Enum):
    ELU = "elu"
    SELU = "selu"
    TANH = "tanh"
    RELU = "relu"
    SOFTMAX = "softmax"
    SIGMOID = "sigmoid"
    SOFTPLUS = "softplus"
    LEAKY_RELU = "leaky_relu"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(self) -> Callable:
        match self:
            case Activations.ELU:
                return F.elu
            case Activations.SELU:
                return F.selu
            case Activations.TANH:
                return F.tanh
            case Activations.RELU:
                return F.relu
            # F.softmax needs an explicit dim; -1 matches the last-axis
            # convention used everywhere in nnx (batch-major). functools.partial
            # (rather than a lambda) is picklable, which matters under DDP /
            # multiprocessing dataloader workers.
            case Activations.SOFTMAX:
                return partial(F.softmax, dim=-1)
            case Activations.SIGMOID:
                return F.sigmoid
            case Activations.SOFTPLUS:
                return F.softplus
            case Activations.LEAKY_RELU:
                return F.leaky_relu

nnx.nn.enum.checkpoints.Checkpoints

Bases: Enum

Source code in nnx/nn/enum/checkpoints.py

class Checkpoints(Enum):
    Q1 = "q1"
    Q2 = "q2"
    Q3 = "q3"
    BEST = "best"
    LAST = "last"
    FIRST = "first"

    def __str__(self):
        return self.value

    def __repr__(self):
        return str(self)

nnx.nn.enum.devices.Devices

Bases: Enum

Source code in nnx/nn/enum/devices.py

class Devices(Enum):
    CPU = "cpu"
    MPS = "mps"
    CUDA = "cuda"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(self) -> torch.device:
        return torch.device(self.value)

    def torch_device(self) -> torch.device:
        """Explicit alias for ``self()`` — more readable in code that mixes
        the enum and torch.device usage."""
        return torch.device(self.value)

    @staticmethod
    def get() -> Devices:
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return Devices.MPS
        elif torch.cuda.is_available():
            return Devices.CUDA
        else:
            return Devices.CPU

    @staticmethod
    def get_torch_device() -> torch.device:
        """Convenience: auto-detect and return the corresponding torch.device
        directly. Equivalent to ``Devices.get().torch_device()``."""
        return Devices.get().torch_device()

get_torch_device() -> torch.device staticmethod

Convenience: auto-detect and return the corresponding torch.device directly. Equivalent to Devices.get().torch_device().

Source code in nnx/nn/enum/devices.py

@staticmethod
def get_torch_device() -> torch.device:
    """Convenience: auto-detect and return the corresponding torch.device
    directly. Equivalent to ``Devices.get().torch_device()``."""
    return Devices.get().torch_device()

torch_device() -> torch.device

Explicit alias for self() — more readable in code that mixes the enum and torch.device usage.

Source code in nnx/nn/enum/devices.py

def torch_device(self) -> torch.device:
    """Explicit alias for ``self()`` — more readable in code that mixes
    the enum and torch.device usage."""
    return torch.device(self.value)

nnx.nn.enum.losses.Losses

Bases: Enum

Source code in nnx/nn/enum/losses.py

class Losses(Enum):
    CROSS_ENTROPY = "cross_entropy"
    MEAN_SQUARED_ERROR = "mean_squared_error"
    BINARY_CROSS_ENTROPY = "binary_cross_entropy"
    NEGATIVE_LOG_LIKELIHOOD = "negative_log_likelihood"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(self) -> nn.Module:
        match self:
            case Losses.CROSS_ENTROPY:
                return nn.CrossEntropyLoss()
            case Losses.MEAN_SQUARED_ERROR:
                return nn.MSELoss()
            case Losses.BINARY_CROSS_ENTROPY:
                return nn.BCELoss()
            case Losses.NEGATIVE_LOG_LIKELIHOOD:
                return nn.NLLLoss()

nnx.nn.enum.nets.Nets

Bases: Enum

Source code in nnx/nn/enum/nets.py

class Nets(Enum):
    FEED_FWD = "feed_fwd"
    GRAPH_ATT = "graph_att"
    GRAPH_CONV = "graph_conv"
    GRAPH_SAGE = "graph_sage"
    # Decoder-only Transformer; consumed by NNTransformerParams. Adding
    # this enum variant is back-compat-safe: existing run.yaml files
    # that don't reference it deserialize unchanged through Nets(<str>).
    TRANSFORMER = "transformer"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(self, params: NNParams) -> nn.Module:
        if params is None:
            raise ValueError("params must not be None")

        match self:
            case Nets.FEED_FWD:
                return FeedFwdNN(params=params)
            case Nets.GRAPH_ATT:
                return GraphAttNN(params=params)
            case Nets.GRAPH_CONV:
                return GraphConvNN(params=params)
            case Nets.GRAPH_SAGE:
                return GraphSageNN(params=params)
            case Nets.TRANSFORMER:
                return TransformerNN(params=params)

nnx.nn.enum.optims.Optims

Bases: Enum

Source code in nnx/nn/enum/optims.py

class Optims(Enum):
    SGD = "sgd"
    ADAM = "adam"
    ADAM_AMSGRAD = "adam_amsgrad"
    SGD_NESTEROV = "sgd_nesterov"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(
        self,
        net: nn.Module,
        lr_start: float,
        weight_decay: float,
        momentum: Union[float, tuple[float, float]],
        param_groups: Optional[list] = None,
        strict_param_groups: bool = False,
    ) -> optim.Optimizer:
        """Build the underlying torch optimizer.

        When ``param_groups`` is None (back-compat default), constructs
        the optimizer with a single group: every trainable parameter of
        ``net`` at ``lr=lr_start``, ``weight_decay=weight_decay``.

        When ``param_groups`` is set to a list of
        :class:`nnx.finetune.NNParamGroupSpec`, dispatches to
        :func:`nnx.finetune.param_groups.build_param_groups` to bucket
        parameters by fnmatch pattern and apply per-group LR /
        weight_decay overrides. Frozen parameters (``requires_grad=False``)
        are dropped — the optimizer doesn't need to know about them.

        ``strict_param_groups`` toggles between fine-tuning and
        multi-optimizer-Trainer semantics: when False (default),
        unmatched parameters go into a default group at ``lr_start``;
        when True, unmatched parameters are dropped from the optimizer
        entirely. The Trainer passes True so disjoint optimizers don't
        end up co-owning the same params via implicit default buckets.
        """
        if net is None:
            raise ValueError("net must not be None")

        if param_groups is not None:
            # Lazy import — defers the finetune subpackage until a
            # param-grouped optimizer is actually built (no cycle today;
            # matches NNOptimParams' deferral style).
            from ...finetune.param_groups import build_param_groups

            params_or_groups = build_param_groups(
                net,
                param_groups,
                default_lr=lr_start,
                default_weight_decay=weight_decay,
                strict=strict_param_groups,
            )
        else:
            params_or_groups = net.parameters()

        match self:
            case Optims.SGD:
                return optim.SGD(
                    params_or_groups,
                    lr=lr_start,
                    momentum=momentum,
                    weight_decay=weight_decay,
                )
            case Optims.ADAM:
                return optim.Adam(
                    params_or_groups,
                    lr=lr_start,
                    betas=momentum,
                    weight_decay=weight_decay,
                )
            case Optims.ADAM_AMSGRAD:
                return optim.Adam(
                    params_or_groups,
                    amsgrad=True,
                    lr=lr_start,
                    betas=momentum,
                    weight_decay=weight_decay,
                )
            case Optims.SGD_NESTEROV:
                return optim.SGD(
                    params_or_groups,
                    nesterov=True,
                    lr=lr_start,
                    momentum=momentum,
                    weight_decay=weight_decay,
                )

__call__(net: nn.Module, lr_start: float, weight_decay: float, momentum: Union[float, tuple[float, float]], param_groups: Optional[list] = None, strict_param_groups: bool = False) -> optim.Optimizer

Build the underlying torch optimizer.

When param_groups is None (back-compat default), constructs the optimizer with a single group: every trainable parameter of net at lr=lr_start, weight_decay=weight_decay.

When param_groups is set to a list of :class:nnx.finetune.NNParamGroupSpec, dispatches to :func:nnx.finetune.param_groups.build_param_groups to bucket parameters by fnmatch pattern and apply per-group LR / weight_decay overrides. Frozen parameters (requires_grad=False) are dropped — the optimizer doesn't need to know about them.

strict_param_groups toggles between fine-tuning and multi-optimizer-Trainer semantics: when False (default), unmatched parameters go into a default group at lr_start; when True, unmatched parameters are dropped from the optimizer entirely. The Trainer passes True so disjoint optimizers don't end up co-owning the same params via implicit default buckets.

Source code in nnx/nn/enum/optims.py

def __call__(
    self,
    net: nn.Module,
    lr_start: float,
    weight_decay: float,
    momentum: Union[float, tuple[float, float]],
    param_groups: Optional[list] = None,
    strict_param_groups: bool = False,
) -> optim.Optimizer:
    """Build the underlying torch optimizer.

    When ``param_groups`` is None (back-compat default), constructs
    the optimizer with a single group: every trainable parameter of
    ``net`` at ``lr=lr_start``, ``weight_decay=weight_decay``.

    When ``param_groups`` is set to a list of
    :class:`nnx.finetune.NNParamGroupSpec`, dispatches to
    :func:`nnx.finetune.param_groups.build_param_groups` to bucket
    parameters by fnmatch pattern and apply per-group LR /
    weight_decay overrides. Frozen parameters (``requires_grad=False``)
    are dropped — the optimizer doesn't need to know about them.

    ``strict_param_groups`` toggles between fine-tuning and
    multi-optimizer-Trainer semantics: when False (default),
    unmatched parameters go into a default group at ``lr_start``;
    when True, unmatched parameters are dropped from the optimizer
    entirely. The Trainer passes True so disjoint optimizers don't
    end up co-owning the same params via implicit default buckets.
    """
    if net is None:
        raise ValueError("net must not be None")

    if param_groups is not None:
        # Lazy import — defers the finetune subpackage until a
        # param-grouped optimizer is actually built (no cycle today;
        # matches NNOptimParams' deferral style).
        from ...finetune.param_groups import build_param_groups

        params_or_groups = build_param_groups(
            net,
            param_groups,
            default_lr=lr_start,
            default_weight_decay=weight_decay,
            strict=strict_param_groups,
        )
    else:
        params_or_groups = net.parameters()

    match self:
        case Optims.SGD:
            return optim.SGD(
                params_or_groups,
                lr=lr_start,
                momentum=momentum,
                weight_decay=weight_decay,
            )
        case Optims.ADAM:
            return optim.Adam(
                params_or_groups,
                lr=lr_start,
                betas=momentum,
                weight_decay=weight_decay,
            )
        case Optims.ADAM_AMSGRAD:
            return optim.Adam(
                params_or_groups,
                amsgrad=True,
                lr=lr_start,
                betas=momentum,
                weight_decay=weight_decay,
            )
        case Optims.SGD_NESTEROV:
            return optim.SGD(
                params_or_groups,
                nesterov=True,
                lr=lr_start,
                momentum=momentum,
                weight_decay=weight_decay,
            )

nnx.nn.enum.schedulers.Schedulers

Bases: Enum

Source code in nnx/nn/enum/schedulers.py

class Schedulers(Enum):
    REDUCE_LR_ON_PLATEAU = "reduce_lr_on_plateau"
    STEP = "step"
    COSINE_ANNEALING = "cosine_annealing"
    ONE_CYCLE = "one_cycle"
    LINEAR_WARMUP_DECAY = "linear_warmup_decay"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(
        self,
        optimizer: Optimizer,
        params: NNSchedulerParams,
        n_epochs: int,
    ):
        match self:
            case Schedulers.REDUCE_LR_ON_PLATEAU:
                return lr_scheduler.ReduceLROnPlateau(
                    optimizer,
                    mode="min",
                    min_lr=params.min_lr,
                    factor=params.factor,
                    cooldown=params.cooldown,
                    patience=params.patience,
                    threshold=params.threshold,
                )
            case Schedulers.STEP:
                step_size = params.step_size if params.step_size is not None else max(1, n_epochs // 3)
                gamma = params.factor
                return lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
            case Schedulers.COSINE_ANNEALING:
                T_max = params.T_max if params.T_max is not None else n_epochs
                return lr_scheduler.CosineAnnealingLR(
                    optimizer,
                    T_max=T_max,
                    eta_min=params.min_lr,
                )
            case Schedulers.ONE_CYCLE:
                max_lr = params.max_lr if params.max_lr is not None else optimizer.param_groups[0]["lr"]
                total_steps = params.total_steps if params.total_steps is not None else n_epochs
                _reject_short_total_steps(total_steps, n_epochs)
                return lr_scheduler.OneCycleLR(
                    optimizer,
                    max_lr=max_lr,
                    total_steps=total_steps,
                )
            case Schedulers.LINEAR_WARMUP_DECAY:
                warmup_steps = params.warmup_steps if params.warmup_steps is not None else max(1, n_epochs // 10)
                total_steps = params.total_steps if params.total_steps is not None else n_epochs
                _reject_short_total_steps(total_steps, n_epochs)

                def _lr_lambda(step: int) -> float:
                    if step < warmup_steps:
                        # step+1: the scheduler steps once per EPOCH, so
                        # a 0.0 factor at step 0 would train the entire
                        # first epoch at LR=0 (HF's per-batch stepping
                        # only wastes one batch with the 0-based form).
                        return float(step + 1) / float(max(1, warmup_steps))
                    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
                    return max(0.0, 1.0 - progress)

                return lr_scheduler.LambdaLR(optimizer, lr_lambda=_lr_lambda)

7. Callbacks

nnx.nn.callbacks.Callback

Base class for training callbacks. Override any subset of the hooks.

Source code in nnx/nn/callbacks.py

class Callback:
    """Base class for training callbacks. Override any subset of the hooks."""

    def on_train_begin(self, ctx: _CallbackContext) -> None:
        pass

    def on_epoch_begin(self, ctx: _CallbackContext) -> None:
        pass

    def on_epoch_end(self, ctx: _CallbackContext) -> None:
        pass

    def on_train_end(self, ctx: _CallbackContext) -> None:
        pass

nnx.nn.callbacks.EarlyStopping

Bases: Callback

Stop training when the monitored metric stops improving.

Parameters:

Name	Type	Description	Default
monitor	str	which IDP field to track. "val_edp.error" (default), "val_edp.loss", "train_edp.error", or "train_edp.loss".	'val_edp.error'
patience	int	epochs with no improvement before stopping.	10
min_delta	float	minimum change to qualify as improvement.	0.0
mode	str	"min" (default) for loss/error; "max" for accuracy/f1.	'min'

Source code in nnx/nn/callbacks.py

class EarlyStopping(Callback):
    """Stop training when the monitored metric stops improving.

    Args:
        monitor: which IDP field to track. "val_edp.error" (default), "val_edp.loss",
                 "train_edp.error", or "train_edp.loss".
        patience: epochs with no improvement before stopping.
        min_delta: minimum change to qualify as improvement.
        mode: "min" (default) for loss/error; "max" for accuracy/f1.
    """

    def __init__(
        self,
        monitor: str = "val_edp.error",
        patience: int = 10,
        min_delta: float = 0.0,
        mode: str = "min",
    ):
        if mode not in ("min", "max"):
            raise ValueError(f"mode must be 'min' or 'max', got {mode!r}")
        self.monitor = monitor
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self._best: Optional[float] = None
        self._wait: int = 0

    def _lookup_monitored(self, idp: NNIterationDataPoint) -> Optional[float]:
        # Named distinctly from nnx._metrics._resolve_metric (the
        # val→train / error→loss fallback resolver) — this one just
        # dereferences the user's `monitor` string, e.g. "val_edp.loss".
        edp_name, _, field = self.monitor.partition(".")
        edp = getattr(idp, edp_name, None)
        if edp is None:
            return None
        return getattr(edp, field, None)

    def on_train_begin(self, ctx: _CallbackContext) -> None:
        # Fresh run, fresh patience: without this reset, reusing one
        # EarlyStopping instance across train() calls compares against
        # the previous run's best and can stop the new run immediately.
        self._best = None
        self._wait = 0

    def _is_improvement(self, current: float, best: float) -> bool:
        if self.mode == "min":
            return current < best - self.min_delta
        return current > best + self.min_delta

    def on_epoch_end(self, ctx: _CallbackContext) -> None:
        if ctx.idp is None:
            return
        current = self._lookup_monitored(ctx.idp)
        if current is None:
            return
        if self._best is None or self._is_improvement(current, self._best):
            self._best = current
            self._wait = 0
        else:
            self._wait += 1
            if self._wait >= self.patience:
                ctx.should_stop = True

nnx.nn.callbacks.LRMonitor

Bases: Callback

Logs the current LR each epoch. History exposed at .history.

Source code in nnx/nn/callbacks.py

class LRMonitor(Callback):
    """Logs the current LR each epoch. History exposed at `.history`."""

    def __init__(self):
        self.history: list[float] = []

    def on_epoch_end(self, ctx: _CallbackContext) -> None:
        lr = ctx.optimizer.param_groups[0]["lr"]
        self.history.append(lr)

nnx.nn.callbacks.ModelCheckpoint

Bases: Callback

Save a custom-tagged checkpoint at user-specified epochs.

The standard train() loop already saves FIRST / Q1 / Q2 / Q3 / LAST / BEST via the Checkpoints enum. This callback adds ad-hoc save points outside that cycle — useful for sampling at fixed milestones (e.g., epoch 10, 20, 50) for downstream inspection.

Each match writes <cwd>/runs/<run.id>/checkpoints/<tag>_e<epoch>.pt — cwd-relative, matching what :meth:NNRun.save and :class:NNCheckpoint use when called from inside :meth:NNModel.train (the train() entry point doesn't accept a root= parameter). The epoch suffix prevents successive matches from overwriting each other when epochs has multiple entries.

Parameters:

Name	Type	Description	Default
epochs	Optional[list[int]]	list of 0-indexed epoch numbers at which to save. Empty / None means the callback never fires (and never saves anything).	None
tag	str	prefix in the filename, defaults to "custom".	'custom'

Source code in nnx/nn/callbacks.py

class ModelCheckpoint(Callback):
    """Save a custom-tagged checkpoint at user-specified epochs.

    The standard train() loop already saves FIRST / Q1 / Q2 / Q3 / LAST / BEST
    via the Checkpoints enum. This callback adds ad-hoc save points outside
    that cycle — useful for sampling at fixed milestones (e.g., epoch 10,
    20, 50) for downstream inspection.

    Each match writes ``<cwd>/runs/<run.id>/checkpoints/<tag>_e<epoch>.pt``
    — cwd-relative, matching what :meth:`NNRun.save` and :class:`NNCheckpoint`
    use when called from inside :meth:`NNModel.train` (the train() entry
    point doesn't accept a ``root=`` parameter). The epoch suffix
    prevents successive matches from overwriting each other when
    ``epochs`` has multiple entries.

    Args:
        epochs: list of 0-indexed epoch numbers at which to save. Empty /
            None means the callback never fires (and never saves anything).
        tag: prefix in the filename, defaults to ``"custom"``.
    """

    def __init__(self, epochs: Optional[list[int]] = None, tag: str = "custom"):
        self.epochs = set(epochs or [])
        self.tag = tag

    def on_epoch_end(self, ctx: _CallbackContext) -> None:
        if ctx.epoch not in self.epochs:
            return
        # Build the NNCheckpoint inline — same shape as NNModel._save_checkpoints
        # but with a custom path so it doesn't collide with the Checkpoints enum
        # tags. Goes through NNCheckpoint.to_file for the atomic-write guarantee.
        ckpt = NNCheckpoint(
            idp=ctx.idp,
            model_params=ctx.model.params,
            net_params=ctx.model.net_params,
            net_state=ctx.model.net.state_dict(),
        )
        # Same cwd-relative `runs/<id>/checkpoints/` layout NNCheckpoint.save
        # uses through _checkpoint_path; we hand-build the path here because
        # the user-supplied tag isn't part of the Checkpoints enum.
        path = os.path.join(
            "runs",
            ctx.run.id,
            "checkpoints",
            f"{self.tag}_e{ctx.epoch}.pt",
        )
        ckpt.to_file(path)

nnx.nn.callbacks.TensorBoardCallback

Bases: Callback

Stream train/val metrics + LR to a TensorBoard SummaryWriter.

Requires tensorboard to be installed — imported lazily so users who don't use this callback don't pay the dependency cost.

Parameters:

Name	Type	Description	Default
log_dir	Optional[str]	directory passed to SummaryWriter. None lets TensorBoard pick its default (runs/).	None
flush_each_epoch	bool	when True (default), calls writer.flush() so partial training is visible in TB even if the process crashes.	True

Source code in nnx/nn/callbacks.py

class TensorBoardCallback(Callback):
    """Stream train/val metrics + LR to a TensorBoard SummaryWriter.

    Requires `tensorboard` to be installed — imported lazily so users who
    don't use this callback don't pay the dependency cost.

    Args:
        log_dir: directory passed to SummaryWriter. None lets TensorBoard
            pick its default (runs/<datetime>).
        flush_each_epoch: when True (default), calls writer.flush() so
            partial training is visible in TB even if the process crashes.
    """

    def __init__(self, log_dir: Optional[str] = None, flush_each_epoch: bool = True):
        try:
            from torch.utils.tensorboard import SummaryWriter
        except ImportError as e:
            raise ImportError(
                "TensorBoardCallback requires `tensorboard`. "
                "Install with `pip install thekaveh-nnx[tensorboard]` or `pip install tensorboard`."
            ) from e
        self._writer = SummaryWriter(log_dir=log_dir)
        self._flush_each_epoch = flush_each_epoch

    def on_epoch_end(self, ctx: _CallbackContext) -> None:
        idp = ctx.idp
        if idp is None:
            return
        step = idp.epoch_idx

        for name, v in _edp_metric_iter(idp.train_edp):
            self._writer.add_scalar(f"train/{name}", v, step)
        for name, v in _edp_metric_iter(idp.val_edp):
            self._writer.add_scalar(f"val/{name}", v, step)
        self._writer.add_scalar("lr", ctx.optimizer.param_groups[0]["lr"], step)

        if self._flush_each_epoch:
            self._writer.flush()

    def on_train_end(self, ctx: _CallbackContext) -> None:
        self._writer.close()

nnx.nn.callbacks.WandbCallback

Bases: Callback

Stream train/val metrics + LR to Weights & Biases.

Requires wandb — lazily imported. Pass project= to start a new run, or wandb_run= to attach to an externally-managed run.

Source code in nnx/nn/callbacks.py

class WandbCallback(Callback):
    """Stream train/val metrics + LR to Weights & Biases.

    Requires `wandb` — lazily imported. Pass `project=` to start a new run,
    or `wandb_run=` to attach to an externally-managed run.
    """

    def __init__(
        self,
        project: Optional[str] = None,
        wandb_run=None,
        **init_kwargs,
    ):
        if wandb_run is None:
            try:
                import wandb
            except ImportError as e:
                raise ImportError(
                    "WandbCallback requires `wandb`. Install with `pip install thekaveh-nnx[wandb]` or `pip install wandb`."
                ) from e
            self._run = wandb.init(project=project, **init_kwargs)
            self._owns_run = True
        else:
            self._run = wandb_run
            self._owns_run = False

    def on_epoch_end(self, ctx: _CallbackContext) -> None:
        idp = ctx.idp
        if idp is None:
            return

        log: dict = {"epoch": idp.epoch_idx, "lr": ctx.optimizer.param_groups[0]["lr"]}
        for name, v in _edp_metric_iter(idp.train_edp):
            log[f"train/{name}"] = v
        for name, v in _edp_metric_iter(idp.val_edp):
            log[f"val/{name}"] = v
        self._run.log(log, step=idp.epoch_idx)

    def on_train_end(self, ctx: _CallbackContext) -> None:
        if self._owns_run:
            self._run.finish()

8. Fine-tuning (nnx.finetune)

nnx.finetune.freezing.freeze(module: nn.Module, *patterns: str) -> int

Set requires_grad=False on every parameter under module whose dotted name matches any of patterns.

Patterns use fnmatch shell-glob semantics: * matches any sequence of characters including dots (not just one path segment), ? matches a single character, [seq] matches one character from the set. Match is against the parameter's full dotted name, e.g., encoder.layer.5.weight. So "encoder.*" matches every parameter under the encoder subtree, including deeply nested ones like encoder.layer.5.weight.

Parameters:

Name	Type	Description	Default
module	Module	any nn.Module.	required
*patterns	str	one or more fnmatch globs. If no patterns are given, raises ValueError (freeze-all-by-default is too dangerous to be the no-arg behavior).	()

Returns:

Type	Description
int	The number of parameters newly frozen (i.e., previously had
int	requires_grad=True). Useful for assertion / logging.

Source code in nnx/finetune/freezing.py

def freeze(module: nn.Module, *patterns: str) -> int:
    """Set ``requires_grad=False`` on every parameter under ``module``
    whose dotted name matches any of ``patterns``.

    Patterns use ``fnmatch`` shell-glob semantics: ``*`` matches any
    sequence of characters **including dots** (not just one path segment),
    ``?`` matches a single character, ``[seq]`` matches one character
    from the set. Match is against the parameter's full dotted name,
    e.g., ``encoder.layer.5.weight``. So ``"encoder.*"`` matches every
    parameter under the encoder subtree, including deeply nested ones
    like ``encoder.layer.5.weight``.

    Args:
        module: any ``nn.Module``.
        *patterns: one or more fnmatch globs. If no patterns are
            given, raises ``ValueError`` (freeze-all-by-default is too
            dangerous to be the no-arg behavior).

    Returns:
        The number of parameters newly frozen (i.e., previously had
        ``requires_grad=True``). Useful for assertion / logging.
    """
    if not patterns:
        raise ValueError("freeze() requires at least one pattern; pass '*' to freeze every parameter")
    n = 0
    for name, param in module.named_parameters():
        if any(fnmatch.fnmatchcase(name, p) for p in patterns):
            if param.requires_grad:
                n += 1
            param.requires_grad = False
    return n

nnx.finetune.freezing.unfreeze(module: nn.Module, *patterns: str) -> int

Mirror of :func:freeze — set requires_grad=True on matching parameters. Returns the count newly unfrozen.

Source code in nnx/finetune/freezing.py

def unfreeze(module: nn.Module, *patterns: str) -> int:
    """Mirror of :func:`freeze` — set ``requires_grad=True`` on matching
    parameters. Returns the count newly unfrozen."""
    if not patterns:
        raise ValueError("unfreeze() requires at least one pattern; pass '*' to unfreeze every parameter")
    n = 0
    for name, param in module.named_parameters():
        if any(fnmatch.fnmatchcase(name, p) for p in patterns):
            if not param.requires_grad:
                n += 1
            param.requires_grad = True
    return n

nnx.finetune.freezing.frozen(module: nn.Module) -> list[str]

List the dotted parameter names currently frozen under module.

Returned list is sorted by name for stable test assertions. Useful for logging at train() entry so users can see exactly which parameters are excluded from training.

Source code in nnx/finetune/freezing.py

def frozen(module: nn.Module) -> list[str]:
    """List the dotted parameter names currently frozen under ``module``.

    Returned list is sorted by name for stable test assertions. Useful
    for logging at ``train()`` entry so users can see exactly which
    parameters are excluded from training.
    """
    return sorted(name for name, param in module.named_parameters() if not param.requires_grad)

nnx.finetune.loading.load_pretrained(module: nn.Module, source: Union[str, Path, dict, nn.Module], *, key_map: Optional[dict[str, str]] = None, strict: bool = False, prefix: Optional[str] = None) -> LoadPretrainedResult

Load weights into module from an external source.

The source can be

• a path (str or Path) to a .pt / .pth file holding a state-dict (loaded with weights_only=True for safety);
• a state-dict (dict) already in memory;
• another nn.Module, in which case its state-dict is used.

Parameters:

Name	Type	Description	Default
module	Module	target module to load into. Mutated in place.	required
source	Union[str, Path, dict, Module]	see above.	required
key_map	Optional[dict[str, str]]	optional remapping from source keys to target keys, applied AFTER prefix stripping and BEFORE matching. Each entry is a prefix substitution: for the first key in key_map whose prefix matches the source key, that prefix is replaced with the mapped value. E.g., {"backbone.": "net."} rewrites backbone.conv1.weight to net.conv1.weight; later occurrences of backbone. mid-string are NOT touched. First-match-wins; subsequent entries don't fire once a key has been remapped.	None
strict	bool	when True, raise if any source key has no target match OR any target key has no source. Default False (fine-tuning commonly partial-loads).	False
prefix	Optional[str]	optional prefix to strip from source keys before matching. E.g., prefix="model." turns model.layer.0 into layer.0. Applied BEFORE key_map.	None

Returns:

Type	Description
LoadPretrainedResult	class:LoadPretrainedResult with the loaded / missing /
LoadPretrainedResult	unexpected key sets.

Source code in nnx/finetune/loading.py

def load_pretrained(
    module: nn.Module,
    source: Union[str, Path, dict, nn.Module],
    *,
    key_map: Optional[dict[str, str]] = None,
    strict: bool = False,
    prefix: Optional[str] = None,
) -> LoadPretrainedResult:
    """Load weights into ``module`` from an external source.

    The source can be:
      - a path (str or Path) to a ``.pt`` / ``.pth`` file holding a
        state-dict (loaded with ``weights_only=True`` for safety);
      - a state-dict (``dict``) already in memory;
      - another ``nn.Module``, in which case its state-dict is used.

    Args:
        module: target module to load into. Mutated in place.
        source: see above.
        key_map: optional remapping from source keys to target keys,
            applied AFTER ``prefix`` stripping and BEFORE matching.
            Each entry is a **prefix** substitution: for the first
            key in ``key_map`` whose prefix matches the source key,
            that prefix is replaced with the mapped value. E.g.,
            ``{"backbone.": "net."}`` rewrites ``backbone.conv1.weight``
            to ``net.conv1.weight``; later occurrences of ``backbone.``
            mid-string are NOT touched. First-match-wins; subsequent
            entries don't fire once a key has been remapped.
        strict: when True, raise if any source key has no target match
            OR any target key has no source. Default False (fine-tuning
            commonly partial-loads).
        prefix: optional prefix to strip from source keys before
            matching. E.g., ``prefix="model."`` turns ``model.layer.0``
            into ``layer.0``. Applied BEFORE ``key_map``.

    Returns:
        :class:`LoadPretrainedResult` with the loaded / missing /
        unexpected key sets.
    """
    # 1. Resolve `source` to a state-dict.
    if isinstance(source, nn.Module):
        src: dict = dict(source.state_dict())
    elif isinstance(source, (str, Path)):
        # weights_only=True is safe — state-dicts are tensors + scalars
        # only. Avoids the ACE risk that NNCheckpoint.from_file warns
        # about. Caller passing a malicious file gets a clean
        # UnpicklingError, not arbitrary code execution.
        src = torch.load(str(source), weights_only=True, map_location="cpu")
    elif isinstance(source, dict):
        src = dict(source)
    else:
        raise TypeError(f"load_pretrained: source must be a path, dict, or nn.Module; got {type(source).__name__}")

    # 2. Apply prefix-stripping (if requested) then key_map remapping.
    remapped: dict = {}
    for key, val in src.items():
        new_key = key
        if prefix is not None and new_key.startswith(prefix):
            new_key = new_key[len(prefix) :]
        if key_map:
            for foreign, local in key_map.items():
                if new_key.startswith(foreign):
                    new_key = local + new_key[len(foreign) :]
                    break
        if new_key in remapped:
            # Two source keys collapsed onto one target after prefix /
            # key_map rewriting. Identical tensors (tied weights stored
            # under two names) are harmless; differing ones would load
            # whichever came last — unpredictable, so raise.
            if not torch.equal(remapped[new_key], val):
                raise ValueError(
                    f"load_pretrained: remapping collapses multiple DIFFERING source keys onto {new_key!r} "
                    "— adjust `prefix` / `key_map` so targets stay unique."
                )
        remapped[new_key] = val

    # 3. Match against the target module's keys and pick the overlap.
    target_keys = set(module.state_dict().keys())
    source_keys = set(remapped.keys())

    loaded_keys = sorted(target_keys & source_keys)
    missing_keys = sorted(target_keys - source_keys)
    unexpected_keys = sorted(source_keys - target_keys)

    if strict and (missing_keys or unexpected_keys):
        raise RuntimeError(
            "load_pretrained(strict=True) found mismatches:\n"
            f"  missing in source:    {missing_keys}\n"
            f"  unexpected in source: {unexpected_keys}"
        )

    # 4. Apply only the overlapping keys. PyTorch's load_state_dict
    # with strict=False does this, but we've already computed the
    # intersection so we can pass just those entries.
    overlap = {k: remapped[k] for k in loaded_keys}
    module.load_state_dict(overlap, strict=False)

    return LoadPretrainedResult(
        loaded_keys=loaded_keys,
        missing_keys=missing_keys,
        unexpected_keys=unexpected_keys,
    )

nnx.finetune.loading.LoadPretrainedResult dataclass

Outcome of a :func:load_pretrained call.

Compared with :meth:torch.nn.Module.load_state_dict, this gives you back not just the missing/unexpected keys but also the list of keys actually applied (after any remapping) — useful for confirming the load did what you intended.

Source code in nnx/finetune/loading.py

@dataclass(frozen=True, slots=True)
class LoadPretrainedResult:
    """Outcome of a :func:`load_pretrained` call.

    Compared with :meth:`torch.nn.Module.load_state_dict`, this gives
    you back not just the missing/unexpected keys but also the list
    of keys actually applied (after any remapping) — useful for
    confirming the load did what you intended.
    """

    loaded_keys: list[str]  # keys present in source AND module (after remap)
    missing_keys: list[str]  # in module's state_dict, not in source
    unexpected_keys: list[str]  # in source, no match in module

nnx.finetune.param_groups.NNParamGroupSpec dataclass

One row in :attr:NNOptimParams.param_groups.

Matches parameters whose dotted name matches name_pattern (fnmatch glob) and applies the specified lr (absolute) or lr_multiplier (multiplied by NNOptimParams.max_lr) and optional weight_decay override.

Exactly one of lr and lr_multiplier may be set. If both are None the matched parameters use the optimizer's default LR — handy when you only want to override weight_decay for a group.

Example

Freeze nothing, but train the backbone at 1/100th the head's LR

and disable weight_decay on every bias term.

NNOptimParams( name=Optims.ADAM, max_lr=1e-3, momentum=(0.9, 0.999), weight_decay=5e-4, param_groups=[ NNParamGroupSpec(name_pattern="encoder.", lr_multiplier=0.01), NNParamGroupSpec(name_pattern=".bias", weight_decay=0.0), ], )

Source code in nnx/finetune/param_groups.py

@dataclass(frozen=True, kw_only=True, slots=True)
class NNParamGroupSpec:
    """One row in :attr:`NNOptimParams.param_groups`.

    Matches parameters whose dotted name matches ``name_pattern``
    (fnmatch glob) and applies the specified ``lr`` (absolute) or
    ``lr_multiplier`` (multiplied by ``NNOptimParams.max_lr``) and
    optional ``weight_decay`` override.

    Exactly one of ``lr`` and ``lr_multiplier`` may be set. If both are
    None the matched parameters use the optimizer's default LR — handy
    when you only want to override ``weight_decay`` for a group.

    Example:
        # Freeze nothing, but train the backbone at 1/100th the head's LR
        # and disable weight_decay on every bias term.
        NNOptimParams(
            name=Optims.ADAM,
            max_lr=1e-3,
            momentum=(0.9, 0.999),
            weight_decay=5e-4,
            param_groups=[
                NNParamGroupSpec(name_pattern="encoder.*", lr_multiplier=0.01),
                NNParamGroupSpec(name_pattern="*.bias", weight_decay=0.0),
            ],
        )
    """

    name_pattern: str
    lr: Optional[float] = None
    lr_multiplier: Optional[float] = None
    weight_decay: Optional[float] = None

    def __post_init__(self) -> None:
        if self.lr is not None and self.lr_multiplier is not None:
            raise ValueError(
                f"NNParamGroupSpec(name_pattern={self.name_pattern!r}): "
                "specify at most one of `lr` and `lr_multiplier`, not both"
            )
        # Fail-fast on lr/scale footguns (same convention as NNOptimParams /
        # NNSchedulerParams): lr or lr_multiplier <= 0 silently zeros the group's
        # effective learning rate (params never update) or, when negative, runs
        # gradient ascent on that group. weight_decay == 0 is valid (disables it),
        # so only reject a negative.
        if self.lr is not None and self.lr <= 0:
            raise ValueError(
                f"NNParamGroupSpec(name_pattern={self.name_pattern!r}): lr must be positive, "
                f"got {self.lr} (to freeze a group use nnx.finetune.freeze, not lr=0)"
            )
        if self.lr_multiplier is not None and self.lr_multiplier <= 0:
            raise ValueError(
                f"NNParamGroupSpec(name_pattern={self.name_pattern!r}): lr_multiplier must be "
                f"positive, got {self.lr_multiplier}"
            )
        if self.weight_decay is not None and self.weight_decay < 0:
            raise ValueError(
                f"NNParamGroupSpec(name_pattern={self.name_pattern!r}): weight_decay must be "
                f"non-negative, got {self.weight_decay}"
            )

    def state(self) -> dict:
        d: dict = dict(name_pattern=self.name_pattern)
        if self.lr is not None:
            d["lr"] = self.lr
        if self.lr_multiplier is not None:
            d["lr_multiplier"] = self.lr_multiplier
        if self.weight_decay is not None:
            d["weight_decay"] = self.weight_decay
        return d

    @staticmethod
    def from_state(state: dict) -> NNParamGroupSpec:
        return NNParamGroupSpec(
            name_pattern=state["name_pattern"],
            lr=state.get("lr"),
            lr_multiplier=state.get("lr_multiplier"),
            weight_decay=state.get("weight_decay"),
        )

nnx.finetune.param_groups.build_param_groups(module: nn.Module, specs: list[NNParamGroupSpec], *, default_lr: float, default_weight_decay: float, strict: bool = False) -> list[dict]

Walk module's parameters, bucket them by the first matching spec (or into a fallback default group), and return the list of param-group dicts the optimizer expects.

Parameters with requires_grad=False are dropped — they're frozen, the optimizer doesn't need to see them. (Without this, the optimizer would still hold them in its state but they'd never update; harmless but wasteful and confusing in optimizer.param_groups.)

Parameters:

Name	Type	Description	Default
module	Module	source of parameters to bucket.	required
specs	list[NNParamGroupSpec]	list of :class:NNParamGroupSpec in priority order. The first spec whose name_pattern matches a parameter's dotted name wins.	required
default_lr	float	LR for parameters that don't match any spec, or for specs that omit both lr and lr_multiplier.	required
default_weight_decay	float	WD for parameters that don't match any spec's weight_decay override.	required
strict	bool	when False (default, fine-tuning semantics), parameters that match no spec go into a default group at default_lr so every trainable parameter ends up in the optimizer. When True (multi-optimizer Trainer semantics), unmatched parameters are DROPPED from the optimizer entirely — the contract is "this optimizer owns only what the specs explicitly select", which is what allows disjoint optimizers in :class:nnx.trainer.Trainer.	False

Returns:

Type	Description
list[dict]	A list of dicts suitable for ``torch.optim.Optimizer(
list[dict]	params, ...)— each entry has"params"`` plus any overrides.

Source code in nnx/finetune/param_groups.py

def build_param_groups(
    module: nn.Module,
    specs: list[NNParamGroupSpec],
    *,
    default_lr: float,
    default_weight_decay: float,
    strict: bool = False,
) -> list[dict]:
    """Walk ``module``'s parameters, bucket them by the first matching
    spec (or into a fallback default group), and return the list of
    param-group dicts the optimizer expects.

    Parameters with ``requires_grad=False`` are dropped — they're
    frozen, the optimizer doesn't need to see them. (Without this, the
    optimizer would still hold them in its state but they'd never
    update; harmless but wasteful and confusing in `optimizer.param_groups`.)

    Args:
        module: source of parameters to bucket.
        specs: list of :class:`NNParamGroupSpec` in priority order.
            The first spec whose ``name_pattern`` matches a parameter's
            dotted name wins.
        default_lr: LR for parameters that don't match any spec, or
            for specs that omit both ``lr`` and ``lr_multiplier``.
        default_weight_decay: WD for parameters that don't match any
            spec's ``weight_decay`` override.
        strict: when False (default, fine-tuning semantics), parameters
            that match no spec go into a default group at ``default_lr``
            so every trainable parameter ends up in the optimizer. When
            True (multi-optimizer Trainer semantics), unmatched parameters
            are DROPPED from the optimizer entirely — the contract is
            "this optimizer owns only what the specs explicitly select",
            which is what allows disjoint optimizers in
            :class:`nnx.trainer.Trainer`.

    Returns:
        A list of dicts suitable for ``torch.optim.Optimizer(
        params, ...)`` — each entry has ``"params"`` plus any overrides.
    """
    # One bucket per spec, in priority order. Each parameter goes into
    # the FIRST matching spec's bucket (break on hit, below); anything
    # unmatched falls through to default_bucket.
    buckets: list[tuple[NNParamGroupSpec, list[nn.Parameter]]] = [(spec, []) for spec in specs]
    default_bucket: list[nn.Parameter] = []

    for name, param in module.named_parameters():
        if not param.requires_grad:
            continue
        for spec, bucket in buckets:
            if fnmatch.fnmatchcase(name, spec.name_pattern):
                bucket.append(param)
                break
        else:
            default_bucket.append(param)

    out: list[dict] = []
    for spec, bucket in buckets:
        if not bucket:
            continue
        group: dict = {"params": bucket}
        if spec.lr is not None:
            group["lr"] = spec.lr
        elif spec.lr_multiplier is not None:
            group["lr"] = default_lr * spec.lr_multiplier
        else:
            group["lr"] = default_lr
        group["weight_decay"] = spec.weight_decay if spec.weight_decay is not None else default_weight_decay
        out.append(group)

    # Default bucket: included under fine-tuning semantics (every trainable
    # param ends up in the optimizer), suppressed under strict mode (the
    # caller — typically nnx.trainer.Trainer — wants this optimizer to own
    # only what the specs explicitly select).
    if default_bucket and not strict:
        out.append(
            {
                "params": default_bucket,
                "lr": default_lr,
                "weight_decay": default_weight_decay,
            }
        )

    if not out:
        raise ValueError(
            "build_param_groups produced no parameter groups — every parameter "
            "is either frozen, unmatched (under strict mode), or the module "
            "has none. Check the specs' name_pattern against module.named_parameters()."
        )

    return out

9. Parameter-efficient fine-tuning (nnx.peft)

LoRA + DoRA + IA3 + Prefix-Tuning + Prompt-Tuning + Adapters. All methods share the same in-place wrap + save/load idiom (per-method save_*_weights / load_*_weights persist only the trainable delta).

9.1. LoRA

nnx.peft.lora.LoRALinear

Bases: Module

Linear layer wrapped with a LoRA low-rank residual.

The original :class:nn.Linear lives at self.base with its parameters frozen (requires_grad=False) on construction. lora_A and lora_B are trainable; lora_A uses Kaiming-uniform init and lora_B is zero-initialized so the layer's output at step 0 equals the base layer's output exactly — fine-tuning starts from the pretrained behavior and diverges only as B picks up gradient.

The wrapper preserves the base layer's in_features / out_features, so consumers that read base.weight.shape or pass tensors through the layer don't change.

Source code in nnx/peft/lora.py

class LoRALinear(nn.Module):
    """Linear layer wrapped with a LoRA low-rank residual.

    The original :class:`nn.Linear` lives at ``self.base`` with its
    parameters frozen (``requires_grad=False``) on construction.
    ``lora_A`` and ``lora_B`` are trainable; ``lora_A`` uses
    Kaiming-uniform init and ``lora_B`` is zero-initialized so the
    layer's output at step 0 equals the base layer's output exactly
    — fine-tuning starts from the pretrained behavior and diverges
    only as B picks up gradient.

    The wrapper preserves the base layer's ``in_features`` /
    ``out_features``, so consumers that read ``base.weight.shape`` or
    pass tensors through the layer don't change.
    """

    def __init__(
        self,
        base: nn.Linear,
        *,
        r: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        if not isinstance(base, nn.Linear):
            raise TypeError(f"LoRALinear requires an nn.Linear base, got {type(base).__name__}")
        if r <= 0:
            raise ValueError(f"LoRA rank r must be positive, got {r}")
        if alpha <= 0:
            raise ValueError(f"LoRA alpha must be positive, got {alpha}")
        if not (0.0 <= dropout < 1.0):
            raise ValueError(f"LoRA dropout must be in [0, 1), got {dropout}")

        self.base = base
        # Freeze every parameter of the wrapped layer — the LoRA residual
        # is the only trainable bit going forward.
        for p in self.base.parameters():
            p.requires_grad = False

        self.r = r
        self.alpha = alpha
        # `scaling = alpha / r` keeps the effective LR independent of r
        # — common LoRA convention.
        self.scaling = alpha / r

        in_features = base.in_features
        out_features = base.out_features

        # A: (r, in). Init with Kaiming-uniform (sqrt(5) gain), matching
        # the original LoRA implementation. B: (out, r), zero-init so
        # the residual contributes 0 at step 0.
        self.lora_A = nn.Parameter(torch.empty(r, in_features))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))

        self.lora_dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()

    @property
    def in_features(self) -> int:
        return self.base.in_features

    @property
    def out_features(self) -> int:
        return self.base.out_features

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Base forward stays unchanged — frozen W·x + b. The LoRA path
        # adds (α/r) · B(A(dropout(x))).
        base_out = self.base(x)
        lora_out = self.lora_dropout(x) @ self.lora_A.t() @ self.lora_B.t()
        return base_out + lora_out * self.scaling

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"r={self.r}, alpha={self.alpha}, scaling={self.scaling}"
        )

nnx.peft.lora.apply_lora_to(module: nn.Module, *name_patterns: str, r: int = 8, alpha: float = 16.0, dropout: float = 0.0) -> int

Wrap every :class:nn.Linear submodule whose dotted name matches any of name_patterns with a :class:LoRALinear. Returns the number of layers wrapped.

Patterns use shell-style globs (fnmatch) against the dotted submodule name as it appears in module.named_modules() — e.g., "layers.0", "encoder.*", "*" for every Linear.

The wrap is in-place: each matched layer is removed from its parent and replaced with a :class:LoRALinear wrapping it. The base layer's parameters end up frozen as a side effect of LoRALinear's construction; the LoRA parameters (lora_A / lora_B) are trainable by default.

Parameters:

Name	Type	Description	Default
module	Module	root module to walk. The function mutates module in place.	required
name_patterns	str	at least one fnmatch glob. Empty raises.	()
r	int	LoRA rank — passed through to :class:LoRALinear.	8
alpha	float	LoRA scaling numerator — passed through.	16.0
dropout	float	dropout on the LoRA path — passed through.	0.0

Returns:

Type	Description
int	The count of layers wrapped (may be 0 if no patterns match).

Raises:

Type	Description
ValueError	if name_patterns is empty.

Idempotency note: if a layer is already a :class:LoRALinear, its inner .base is skipped — re-applying apply_lora_to against the same patterns is a no-op for layers that already carry a LoRA wrapper. The function returns the count of NEW wraps.

Source code in nnx/peft/lora.py

def apply_lora_to(
    module: nn.Module,
    *name_patterns: str,
    r: int = 8,
    alpha: float = 16.0,
    dropout: float = 0.0,
) -> int:
    """Wrap every :class:`nn.Linear` submodule whose dotted name matches
    any of ``name_patterns`` with a :class:`LoRALinear`. Returns the
    number of layers wrapped.

    Patterns use shell-style globs (``fnmatch``) against the dotted
    submodule name as it appears in ``module.named_modules()`` — e.g.,
    ``"layers.0"``, ``"encoder.*"``, ``"*"`` for every Linear.

    The wrap is in-place: each matched layer is removed from its parent
    and replaced with a :class:`LoRALinear` wrapping it. The base
    layer's parameters end up frozen as a side effect of LoRALinear's
    construction; the LoRA parameters (``lora_A`` / ``lora_B``) are
    trainable by default.

    Args:
        module: root module to walk. The function mutates ``module``
            in place.
        name_patterns: at least one fnmatch glob. Empty raises.
        r: LoRA rank — passed through to :class:`LoRALinear`.
        alpha: LoRA scaling numerator — passed through.
        dropout: dropout on the LoRA path — passed through.

    Returns:
        The count of layers wrapped (may be 0 if no patterns match).

    Raises:
        ValueError: if ``name_patterns`` is empty.

    **Idempotency note:** if a layer is already a :class:`LoRALinear`,
    its inner ``.base`` is skipped — re-applying ``apply_lora_to``
    against the same patterns is a no-op for layers that already
    carry a LoRA wrapper. The function returns the count of NEW wraps.
    """
    if not name_patterns:
        raise ValueError("apply_lora_to requires at least one name pattern")

    # Two-phase: collect targets first, then mutate. Iterating
    # named_modules() while reassigning child attributes would
    # invalidate the traversal in ways that depend on dict iteration
    # order. Collecting first keeps the loop predictable.
    targets: list[str] = []
    for name, child in module.named_modules():
        if not name:
            # named_modules() yields the root itself under "" — it has
            # no parent attribute to reassign, so an in-place wrap is
            # impossible. Skip it (wrap the root yourself if needed).
            continue
        if not isinstance(child, nn.Linear):
            continue
        # Skip the inner .base of an existing LoRALinear — its parent
        # is a LoRALinear, which we surface via get_submodule.
        parent_path, _, _ = name.rpartition(".")
        parent = module if not parent_path else module.get_submodule(parent_path)
        if isinstance(parent, LoRALinear):
            continue
        if any(fnmatch.fnmatchcase(name, p) for p in name_patterns):
            targets.append(name)

    for name in targets:
        parent_path, _, attr = name.rpartition(".")
        parent = module if not parent_path else module.get_submodule(parent_path)
        old = getattr(parent, attr)
        setattr(parent, attr, LoRALinear(old, r=r, alpha=alpha, dropout=dropout))

    return len(targets)

nnx.peft.lora.save_lora_weights(module: nn.Module, path: Union[str, Path]) -> str

Save ONLY the LoRA parameters of module to path.

The output is a plain torch.save of a dict-subset of the full state_dict, containing only keys with lora_A or lora_B in them. Loadable via :func:load_lora_weights.

Parameters:

Name	Type	Description	Default
module	Module	any module that has been processed by :func:apply_lora_to. If no LoRA params exist, an empty dict is saved (the caller decides whether that's an error).	required
path	Union[str, Path]	destination file path.	required

Returns:

Type	Description
str	The path written (so calls can be chained).

Source code in nnx/peft/lora.py

def save_lora_weights(module: nn.Module, path: Union[str, Path]) -> str:
    """Save ONLY the LoRA parameters of ``module`` to ``path``.

    The output is a plain ``torch.save`` of a dict-subset of the full
    state_dict, containing only keys with ``lora_A`` or ``lora_B`` in
    them. Loadable via :func:`load_lora_weights`.

    Args:
        module: any module that has been processed by
            :func:`apply_lora_to`. If no LoRA params exist, an empty
            dict is saved (the caller decides whether that's an error).
        path: destination file path.

    Returns:
        The path written (so calls can be chained).
    """
    sd = _lora_keys_only(module.state_dict())
    torch.save(sd, str(path))
    return str(path)

nnx.peft.lora.load_lora_weights(module: nn.Module, source: Union[str, Path, dict]) -> int

Load LoRA parameters into module from source.

Parameters:

Name	Type	Description	Default
module	Module	must already have :class:LoRALinear wrappers in the same positions as the source — apply_lora_to FIRST, then call this. Otherwise the keys won't match and 0 params load.	required
source	Union[str, Path, dict]	either a path to a file produced by :func:save_lora_weights, or a state-dict dict directly.	required

Returns:

Type	Description
int	The number of parameter tensors loaded.

Loads via module.load_state_dict(..., strict=False) so the base layer's frozen weights — which are NOT in the LoRA-only checkpoint — don't trigger a missing-keys error.

Source code in nnx/peft/lora.py

def load_lora_weights(module: nn.Module, source: Union[str, Path, dict]) -> int:
    """Load LoRA parameters into ``module`` from ``source``.

    Args:
        module: must already have :class:`LoRALinear` wrappers in the
            same positions as the source — apply_lora_to FIRST, then
            call this. Otherwise the keys won't match and 0 params load.
        source: either a path to a file produced by
            :func:`save_lora_weights`, or a state-dict dict directly.

    Returns:
        The number of parameter tensors loaded.

    Loads via ``module.load_state_dict(..., strict=False)`` so the
    base layer's frozen weights — which are NOT in the LoRA-only
    checkpoint — don't trigger a missing-keys error.
    """
    sd = _resolve_source_to_state_dict(source, "load_lora_weights")
    sd = _lora_keys_only(sd)
    result = module.load_state_dict(sd, strict=False)
    # strict=False silently drops keys that don't exist on the module
    # (e.g. loading into an un-adapted model) — subtract them so the
    # return value is the number of tensors that actually landed.
    return len(sd) - len(result.unexpected_keys)

9.2. DoRA

nnx.peft.dora.DoRALinear

Bases: LoRALinear

Linear layer wrapped with a DoRA weight decomposition.

Subclasses :class:LoRALinear to inherit the frozen-base + trainable low-rank residual machinery (lora_A, lora_B, alpha/r scaling, optional dropout, base-freeze-on-construction). Adds a trainable per-output-row magnitude parameter (shape: out_features) initialized from the column-wise L2 norm of the base weight.

The forward composes the LoRA residual into a combined weight V = W_0 + (α/r) · BA, normalizes V row-wise, then re-scales by the trainable magnitude:

``W = magnitude.unsqueeze(1) * V / ||V||_c``
``y = W · x + b``

At step 0, B is zero-initialized (inherited from LoRALinear) so V = W_0 and ||V||_c == magnitude, giving W == W_0 exactly — fine-tuning starts from the pretrained behavior.

Parameters:

Name	Type	Description	Default
base	Linear	the :class:nn.Linear to wrap. Its parameters are frozen on construction (inherited from LoRALinear).	required
r	int	low-rank dim for the LoRA residual. Must be positive.	8
alpha	float	scaling numerator. Effective LoRA scale is alpha / r.	16.0
dropout	float	dropout on the LoRA update path. Range [0, 1).	0.0

Source code in nnx/peft/dora.py

class DoRALinear(LoRALinear):
    """Linear layer wrapped with a DoRA weight decomposition.

    Subclasses :class:`LoRALinear` to inherit the frozen-base + trainable
    low-rank residual machinery (``lora_A``, ``lora_B``, alpha/r scaling,
    optional dropout, base-freeze-on-construction). Adds a trainable
    per-output-row ``magnitude`` parameter (shape: ``out_features``)
    initialized from the column-wise L2 norm of the base weight.

    The forward composes the LoRA residual into a combined weight
    ``V = W_0 + (α/r) · BA``, normalizes ``V`` row-wise, then re-scales
    by the trainable magnitude:

        ``W = magnitude.unsqueeze(1) * V / ||V||_c``
        ``y = W · x + b``

    At step 0, ``B`` is zero-initialized (inherited from LoRALinear)
    so ``V = W_0`` and ``||V||_c == magnitude``, giving ``W == W_0``
    exactly — fine-tuning starts from the pretrained behavior.

    Args:
        base: the :class:`nn.Linear` to wrap. Its parameters are frozen
            on construction (inherited from LoRALinear).
        r: low-rank dim for the LoRA residual. Must be positive.
        alpha: scaling numerator. Effective LoRA scale is ``alpha / r``.
        dropout: dropout on the LoRA update path. Range ``[0, 1)``.
    """

    def __init__(
        self,
        base: nn.Linear,
        *,
        r: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.0,
    ):
        super().__init__(base, r=r, alpha=alpha, dropout=dropout)
        # Magnitude initialized from the per-output-row L2 norm of the
        # base weight. With ``B = 0`` (LoRA zero-init), the combined
        # weight ``V = W_0``, so ``magnitude * V / ||V||_c`` reduces to
        # ``||W_0||_c * (W_0 / ||W_0||_c) == W_0`` — output equals
        # base(x) exactly at step 0.
        with torch.no_grad():
            init_mag = self.base.weight.norm(p=2, dim=1)  # (out_features,)
        self.magnitude = nn.Parameter(init_mag.clone())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Compose V = W_0 + (α/r) · BA. Note: LoRALinear stores
        # ``self.scaling = alpha / r``, used for the inherited LoRA
        # forward; we reuse it here for the DoRA recomposition.
        lora_update = (self.lora_B @ self.lora_A) * self.scaling
        # Optional dropout on the LoRA matrix (not the input). NOTE:
        # this is DropConnect-style — it zeroes/rescales entries of the
        # composed BA *weight matrix* batch-wide, unlike LoRALinear,
        # which drops entries of the *input activations* per sample.
        # Deliberate: dropping x here would also perturb the row norms
        # the magnitude renormalization below depends on.
        # ``self.lora_dropout`` is Identity when dropout=0, so this is
        # a no-op in the common case.
        lora_update = self.lora_dropout(lora_update)
        V = self.base.weight + lora_update
        # Per-output-row L2 norm. clamp_min avoids divide-by-zero in
        # the pathological case where a row of V is all zeros (which
        # shouldn't happen for a sensibly-init'd base, but be safe).
        norm = V.norm(p=2, dim=1, keepdim=True).clamp_min(1e-8)
        V_normalized = V / norm
        W = self.magnitude.unsqueeze(1) * V_normalized
        return torch.nn.functional.linear(x, W, self.base.bias)

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"r={self.r}, alpha={self.alpha}, scaling={self.scaling} (DoRA)"
        )

nnx.peft.dora.apply_dora_to(module: nn.Module, *name_patterns: str, r: int = 8, alpha: float = 16.0, dropout: float = 0.0) -> int

Wrap every :class:nn.Linear submodule whose dotted name matches any of name_patterns with a :class:DoRALinear. Returns the number of layers wrapped.

Mirrors :func:nnx.peft.apply_lora_to — same fnmatch glob conventions, same two-phase (collect-then-mutate) traversal, same idempotency contract (existing DoRA/LoRA wrappers are not re-wrapped — the parent-is-LoRALinear check covers DoRALinear by inheritance).

Parameters:

Name	Type	Description	Default
module	Module	root module to walk. Mutated in place.	required
name_patterns	str	at least one fnmatch glob.	()
r	int	LoRA rank — passed through.	8
alpha	float	LoRA scaling numerator — passed through.	16.0
dropout	float	dropout on the LoRA update path — passed through.	0.0

Returns:

Type	Description
int	The count of layers wrapped (may be 0 if no patterns match
int	or every match is already wrapped).

Raises:

Type	Description
ValueError	if name_patterns is empty.

Source code in nnx/peft/dora.py

def apply_dora_to(
    module: nn.Module,
    *name_patterns: str,
    r: int = 8,
    alpha: float = 16.0,
    dropout: float = 0.0,
) -> int:
    """Wrap every :class:`nn.Linear` submodule whose dotted name matches
    any of ``name_patterns`` with a :class:`DoRALinear`. Returns the
    number of layers wrapped.

    Mirrors :func:`nnx.peft.apply_lora_to` — same fnmatch glob conventions,
    same two-phase (collect-then-mutate) traversal, same idempotency
    contract (existing DoRA/LoRA wrappers are not re-wrapped — the
    parent-is-LoRALinear check covers DoRALinear by inheritance).

    Args:
        module: root module to walk. Mutated in place.
        name_patterns: at least one fnmatch glob.
        r: LoRA rank — passed through.
        alpha: LoRA scaling numerator — passed through.
        dropout: dropout on the LoRA update path — passed through.

    Returns:
        The count of layers wrapped (may be 0 if no patterns match
        or every match is already wrapped).

    Raises:
        ValueError: if ``name_patterns`` is empty.
    """
    if not name_patterns:
        raise ValueError("apply_dora_to requires at least one name pattern")

    # Two-phase traversal — same rationale as apply_lora_to: avoid
    # invalidating the iterator while reassigning child attributes.
    targets: list[str] = []
    for name, child in module.named_modules():
        if not name:
            # named_modules() yields the root itself under "" — it has
            # no parent attribute to reassign, so an in-place wrap is
            # impossible. Skip it (wrap the root yourself if needed).
            continue
        if not isinstance(child, nn.Linear):
            continue
        # Skip the inner .base of an existing LoRALinear (which covers
        # DoRALinear via inheritance) — re-applying must be idempotent.
        parent_path, _, _ = name.rpartition(".")
        parent = module if not parent_path else module.get_submodule(parent_path)
        if isinstance(parent, LoRALinear):
            continue
        if any(fnmatch.fnmatchcase(name, p) for p in name_patterns):
            targets.append(name)

    for name in targets:
        parent_path, _, attr = name.rpartition(".")
        parent = module if not parent_path else module.get_submodule(parent_path)
        old = getattr(parent, attr)
        setattr(parent, attr, DoRALinear(old, r=r, alpha=alpha, dropout=dropout))

    return len(targets)

9.3. IA3

nnx.peft.ia3.IA3Linear

Bases: Module

Linear layer wrapped with an IA3 per-output-dim scaling vector.

The original :class:nn.Linear lives at self.base with its parameters frozen (requires_grad=False) on construction. scaling is the only trainable parameter: a length-out_features vector initialized to all-ones so the layer's output at step 0 equals the base layer's output exactly.

Forward: y = base(x) * scaling (broadcast over the trailing dim).

Parameters:

Name	Type	Description	Default
base	Linear	the :class:nn.Linear to wrap.	required

Source code in nnx/peft/ia3.py

class IA3Linear(nn.Module):
    """Linear layer wrapped with an IA3 per-output-dim scaling vector.

    The original :class:`nn.Linear` lives at ``self.base`` with its
    parameters frozen (``requires_grad=False``) on construction.
    ``scaling`` is the only trainable parameter: a length-``out_features``
    vector initialized to all-ones so the layer's output at step 0
    equals the base layer's output exactly.

    Forward: ``y = base(x) * scaling`` (broadcast over the trailing dim).

    Args:
        base: the :class:`nn.Linear` to wrap.
    """

    def __init__(self, base: nn.Linear):
        super().__init__()
        if not isinstance(base, nn.Linear):
            raise TypeError(f"IA3Linear requires an nn.Linear base, got {type(base).__name__}")

        self.base = base
        # Freeze every parameter of the wrapped layer — the scaling
        # vector is the only trainable bit going forward.
        for p in self.base.parameters():
            p.requires_grad = False

        # Per-output-dim scaling vector, initialized to all-ones so
        # the layer's output at step 0 equals base(x) exactly.
        self.scaling = nn.Parameter(torch.ones(base.out_features))

    @property
    def in_features(self) -> int:
        return self.base.in_features

    @property
    def out_features(self) -> int:
        return self.base.out_features

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # `self.scaling` broadcasts over leading batch / sequence dims:
        # base(x) is (..., out_features) and scaling is (out_features,).
        return self.base(x) * self.scaling

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features} (IA3)"

nnx.peft.ia3.apply_ia3_to(module: nn.Module, *name_patterns: str) -> int

Wrap every :class:nn.Linear submodule whose dotted name matches any of name_patterns with an :class:IA3Linear. Returns the number of layers wrapped.

Mirrors :func:nnx.peft.apply_lora_to — same fnmatch glob conventions, same two-phase (collect-then-mutate) traversal, same idempotency contract (existing IA3 wrappers are skipped via the parent-is-IA3Linear check).

Parameters:

Name	Type	Description	Default
module	Module	root module to walk. Mutated in place.	required
name_patterns	str	at least one fnmatch glob.	()

Returns:

Type	Description
int	The count of layers wrapped (may be 0 if no patterns match
int	or every match is already wrapped).

Raises:

Type	Description
ValueError	if name_patterns is empty.

Source code in nnx/peft/ia3.py

def apply_ia3_to(module: nn.Module, *name_patterns: str) -> int:
    """Wrap every :class:`nn.Linear` submodule whose dotted name matches
    any of ``name_patterns`` with an :class:`IA3Linear`. Returns the
    number of layers wrapped.

    Mirrors :func:`nnx.peft.apply_lora_to` — same fnmatch glob conventions,
    same two-phase (collect-then-mutate) traversal, same idempotency
    contract (existing IA3 wrappers are skipped via the parent-is-IA3Linear
    check).

    Args:
        module: root module to walk. Mutated in place.
        name_patterns: at least one fnmatch glob.

    Returns:
        The count of layers wrapped (may be 0 if no patterns match
        or every match is already wrapped).

    Raises:
        ValueError: if ``name_patterns`` is empty.
    """
    if not name_patterns:
        raise ValueError("apply_ia3_to requires at least one name pattern")

    targets: list[str] = []
    for name, child in module.named_modules():
        if not name:
            # named_modules() yields the root itself under "" — it has
            # no parent attribute to reassign, so an in-place wrap is
            # impossible. Skip it (wrap the root yourself if needed).
            continue
        if not isinstance(child, nn.Linear):
            continue
        # Skip the inner .base of an existing IA3Linear — re-applying
        # must be idempotent.
        parent_path, _, _ = name.rpartition(".")
        parent = module if not parent_path else module.get_submodule(parent_path)
        if isinstance(parent, IA3Linear):
            continue
        if any(fnmatch.fnmatchcase(name, p) for p in name_patterns):
            targets.append(name)

    for name in targets:
        parent_path, _, attr = name.rpartition(".")
        parent = module if not parent_path else module.get_submodule(parent_path)
        old = getattr(parent, attr)
        setattr(parent, attr, IA3Linear(old))

    return len(targets)

nnx.peft.ia3.save_ia3_weights(module: nn.Module, path: Union[str, Path]) -> str

Save ONLY the IA3 scaling parameters of module to path.

The output is a torch.save of a dict-subset of the full state_dict, containing only keys whose name includes scaling. Loadable via :func:load_ia3_weights.

Parameters:

Name	Type	Description	Default
module	Module	any module that has been processed by :func:apply_ia3_to. If no IA3 params exist, an empty dict is saved.	required
path	Union[str, Path]	destination file path.	required

Returns:

Type	Description
str	The path written (so calls can be chained).

Source code in nnx/peft/ia3.py

def save_ia3_weights(module: nn.Module, path: Union[str, Path]) -> str:
    """Save ONLY the IA3 ``scaling`` parameters of ``module`` to ``path``.

    The output is a ``torch.save`` of a dict-subset of the full
    state_dict, containing only keys whose name includes ``scaling``.
    Loadable via :func:`load_ia3_weights`.

    Args:
        module: any module that has been processed by
            :func:`apply_ia3_to`. If no IA3 params exist, an empty
            dict is saved.
        path: destination file path.

    Returns:
        The path written (so calls can be chained).
    """
    sd = _ia3_keys_only(module.state_dict())
    torch.save(sd, str(path))
    return str(path)

nnx.peft.ia3.load_ia3_weights(module: nn.Module, source: Union[str, Path, dict]) -> int

Load IA3 scaling parameters into module from source.

Parameters:

Name	Type	Description	Default
module	Module	must already have :class:IA3Linear wrappers in the same positions as the source — apply_ia3_to FIRST, then call this. Otherwise the keys won't match and 0 params load.	required
source	Union[str, Path, dict]	either a path to a file produced by :func:save_ia3_weights, or a state-dict dict directly.	required

Returns:

Type	Description
int	The number of parameter tensors loaded.

Loads via module.load_state_dict(..., strict=False) so the base layer's frozen weights — which are NOT in the IA3-only checkpoint — don't trigger a missing-keys error.

Source code in nnx/peft/ia3.py

def load_ia3_weights(module: nn.Module, source: Union[str, Path, dict]) -> int:
    """Load IA3 ``scaling`` parameters into ``module`` from ``source``.

    Args:
        module: must already have :class:`IA3Linear` wrappers in the
            same positions as the source — apply_ia3_to FIRST, then
            call this. Otherwise the keys won't match and 0 params load.
        source: either a path to a file produced by
            :func:`save_ia3_weights`, or a state-dict dict directly.

    Returns:
        The number of parameter tensors loaded.

    Loads via ``module.load_state_dict(..., strict=False)`` so the
    base layer's frozen weights — which are NOT in the IA3-only
    checkpoint — don't trigger a missing-keys error.
    """
    sd = _resolve_source_to_state_dict(source, "load_ia3_weights")
    sd = _ia3_keys_only(sd)
    result = module.load_state_dict(sd, strict=False)
    # strict=False silently drops keys that don't exist on the module
    # (e.g. loading into an un-adapted model) — subtract them so the
    # return value is the number of tensors that actually landed.
    return len(sd) - len(result.unexpected_keys)

9.4. Prefix Tuning

nnx.peft.prefix.PrefixTuner

Bases: Module

Wrap a :class:TransformerNN with learnable per-layer K/V prefixes.

Freezes every parameter of the wrapped model on construction and registers n_layers pairs of (n_prefix, n_heads, head_dim) K / V tensors as the only trainable parameters.

Parameters:

Name	Type	Description	Default
model	TransformerNN	a :class:TransformerNN instance. Its parameters are mutated in place (set to requires_grad=False); the attention forward of each targeted block is monkey-patched.	required
n_prefix	int	number of virtual prefix tokens per layer. Must be > 0.	10
n_layers	Optional[int]	number of leading transformer blocks to attach a prefix to. None (default) targets every block in model.blocks. When set, the first n_layers blocks are targeted; later blocks run un-prefixed.	None

Note on shape: the prefix uses n_heads and head_dim taken from the model's params — there's no per-block override, since every block in a TransformerNN shares the same attention shape.

Raises:

Type	Description
TypeError	if model is not a :class:TransformerNN.
ValueError	if n_prefix or n_layers is out of range, or if model is already prefix-tuned — a second tuner would silently hijack the patched forwards (they read the MHA's _nnx_prefix_tuner ref, which the second tuner overwrites, so the first tuner's parameters stop receiving gradients).

Source code in nnx/peft/prefix.py

class PrefixTuner(nn.Module):
    """Wrap a :class:`TransformerNN` with learnable per-layer K/V prefixes.

    Freezes every parameter of the wrapped model on construction and
    registers ``n_layers`` pairs of ``(n_prefix, n_heads, head_dim)``
    K / V tensors as the only trainable parameters.

    Args:
        model: a :class:`TransformerNN` instance. Its parameters are
            mutated in place (set to ``requires_grad=False``); the
            attention forward of each targeted block is monkey-patched.
        n_prefix: number of virtual prefix tokens per layer. Must be > 0.
        n_layers: number of leading transformer blocks to attach a prefix
            to. ``None`` (default) targets every block in
            ``model.blocks``. When set, the first ``n_layers`` blocks
            are targeted; later blocks run un-prefixed.

    Note on shape: the prefix uses ``n_heads`` and ``head_dim`` taken
    from the model's ``params`` — there's no per-block override, since
    every block in a TransformerNN shares the same attention shape.

    Raises:
        TypeError: if ``model`` is not a :class:`TransformerNN`.
        ValueError: if ``n_prefix`` or ``n_layers`` is out of range, or
            if ``model`` is already prefix-tuned — a second tuner would
            silently hijack the patched forwards (they read the MHA's
            ``_nnx_prefix_tuner`` ref, which the second tuner overwrites,
            so the first tuner's parameters stop receiving gradients).
    """

    def __init__(
        self,
        model: TransformerNN,
        *,
        n_prefix: int = 10,
        n_layers: Optional[int] = None,
    ):
        super().__init__()
        if not isinstance(model, TransformerNN):
            raise TypeError(f"PrefixTuner requires a TransformerNN, got {type(model).__name__}")
        if n_prefix <= 0:
            raise ValueError(f"n_prefix must be positive, got {n_prefix}")

        total_blocks = len(model.blocks)
        if n_layers is None:
            n_layers = total_blocks
        if n_layers <= 0:
            raise ValueError(f"n_layers must be positive, got {n_layers}")
        if n_layers > total_blocks:
            raise ValueError(f"n_layers={n_layers} exceeds model depth ({total_blocks} blocks)")
        if any(hasattr(block.attn, "_nnx_prefix_tuner") for block in model.blocks):
            raise ValueError(
                "PrefixTuner: this TransformerNN is already prefix-tuned. A second "
                "tuner would silently hijack the patched attention forwards (the "
                "first tuner's parameters would stop receiving gradients while its "
                "forward injects the new tuner's prefixes). Wrap a fresh model, or "
                "copy.deepcopy the tuned one first."
            )

        self.model = model
        self.n_prefix = n_prefix

        # Freeze every parameter of the wrapped TransformerNN — the
        # prefix tensors are the only trainable bit going forward.
        for p in self.model.parameters():
            p.requires_grad = False

        # Allocate (n_prefix, n_heads, head_dim) K and V per targeted
        # block. We store them as a ParameterList so they live in
        # state_dict and are visible to `self.parameters()`.
        n_heads = model.params.n_heads
        head_dim = model.params.d_model // n_heads

        self.prefix_keys = nn.ParameterList(
            [nn.Parameter(torch.empty(n_prefix, n_heads, head_dim)) for _ in range(n_layers)]
        )
        self.prefix_values = nn.ParameterList(
            [nn.Parameter(torch.empty(n_prefix, n_heads, head_dim)) for _ in range(n_layers)]
        )
        # Init with small Gaussian noise — matches the original
        # prefix-tuning paper's "random init" baseline.
        for p in self.prefix_keys:
            nn.init.normal_(p, std=0.02)
        for p in self.prefix_values:
            nn.init.normal_(p, std=0.02)

        # Monkey-patch each targeted block's MHA forward to inject the
        # learned prefix. We keep a reference to each MHA's original
        # forward in a sidecar list so the patch can be undone in
        # principle (not exposed as public API, but useful for clean
        # teardown in tests).
        self._original_attn_forwards: list = []
        for i in range(n_layers):
            block = model.blocks[i]
            mha = block.attn
            self._original_attn_forwards.append(mha.forward)
            # Deepcopy-safety: the patch is a MODULE-LEVEL function bound
            # via MethodType, with the tuner/layer refs stored on the MHA
            # itself (raw __dict__ entries — object.__setattr__ bypasses
            # Module registration, avoiding a tuner<->model submodule
            # cycle). copy.deepcopy rebinds the method to the COPIED mha
            # and deep-copies the refs through the memo, so a deepcopy of
            # a prefix-tuned net is fully independent. A closure here
            # would be copied atomically, leaving the copy's attention
            # silently reading the ORIGINAL weights/prefix (bit ptq,
            # born_again teachers, and surgery copies).
            object.__setattr__(mha, "_nnx_prefix_tuner", self)
            mha._nnx_prefix_layer_idx = i
            mha.forward = types.MethodType(_prefix_patched_forward, mha)

    # ------------------------------------------------------------------
    # Forward + trainable params
    # ------------------------------------------------------------------

    def forward(self, *args, **kwargs):
        """Delegate to the wrapped model. The prefix injection happens
        inside each block's monkey-patched MHA forward."""
        return self.model(*args, **kwargs)

    def trainable_parameters(self) -> Iterator[nn.Parameter]:
        """Yield only the learned prefix tensors.

        The wrapped model's parameters are frozen on construction; this
        is the iterator you hand to an optimizer.
        """
        yield from self.prefix_keys
        yield from self.prefix_values

    # ------------------------------------------------------------------
    # Save / load
    # ------------------------------------------------------------------

    def prefix_state_dict(self) -> dict:
        """Return a state-dict containing only the prefix tensors,
        keyed for round-trip via :meth:`load_prefix_weights`.

        The keys are the same as ``self.state_dict()`` filtered to the
        prefix entries — i.e., ``prefix_keys.0``, ``prefix_values.0``,
        ``prefix_keys.1``, …
        """
        return {k: v for k, v in self.state_dict().items() if "prefix_" in k}

forward(*args, **kwargs)

Delegate to the wrapped model. The prefix injection happens inside each block's monkey-patched MHA forward.

Source code in nnx/peft/prefix.py

def forward(self, *args, **kwargs):
    """Delegate to the wrapped model. The prefix injection happens
    inside each block's monkey-patched MHA forward."""
    return self.model(*args, **kwargs)

prefix_state_dict() -> dict

Return a state-dict containing only the prefix tensors, keyed for round-trip via :meth:load_prefix_weights.

The keys are the same as self.state_dict() filtered to the prefix entries — i.e., prefix_keys.0, prefix_values.0, prefix_keys.1, …

Source code in nnx/peft/prefix.py

def prefix_state_dict(self) -> dict:
    """Return a state-dict containing only the prefix tensors,
    keyed for round-trip via :meth:`load_prefix_weights`.

    The keys are the same as ``self.state_dict()`` filtered to the
    prefix entries — i.e., ``prefix_keys.0``, ``prefix_values.0``,
    ``prefix_keys.1``, …
    """
    return {k: v for k, v in self.state_dict().items() if "prefix_" in k}

trainable_parameters() -> Iterator[nn.Parameter]

Yield only the learned prefix tensors.

The wrapped model's parameters are frozen on construction; this is the iterator you hand to an optimizer.

Source code in nnx/peft/prefix.py

def trainable_parameters(self) -> Iterator[nn.Parameter]:
    """Yield only the learned prefix tensors.

    The wrapped model's parameters are frozen on construction; this
    is the iterator you hand to an optimizer.
    """
    yield from self.prefix_keys
    yield from self.prefix_values

nnx.peft.prefix.save_prefix_weights(tuner: PrefixTuner, path: Union[str, Path]) -> str

Save ONLY the prefix tensors of tuner to path.

Parameters:

Name	Type	Description	Default
tuner	PrefixTuner	a :class:PrefixTuner instance.	required
path	Union[str, Path]	destination file path.	required

Returns:

Type	Description
str	The path written, so calls can be chained.

Source code in nnx/peft/prefix.py

def save_prefix_weights(tuner: PrefixTuner, path: Union[str, Path]) -> str:
    """Save ONLY the prefix tensors of ``tuner`` to ``path``.

    Args:
        tuner: a :class:`PrefixTuner` instance.
        path: destination file path.

    Returns:
        The path written, so calls can be chained.
    """
    sd = tuner.prefix_state_dict()
    torch.save(sd, str(path))
    return str(path)

nnx.peft.prefix.load_prefix_weights(tuner: PrefixTuner, source: Union[str, Path, dict]) -> int

Load prefix tensors into tuner from source.

Parameters:

Name	Type	Description	Default
tuner	PrefixTuner	must already have the same prefix shape as the source (same n_prefix, n_heads, head_dim, n_layers). Otherwise load_state_dict will surface the mismatch.	required
source	Union[str, Path, dict]	a path to a file produced by :func:save_prefix_weights, or a state-dict dict directly.	required

Returns:

Type	Description
int	The number of parameter tensors loaded.

Source code in nnx/peft/prefix.py

def load_prefix_weights(tuner: PrefixTuner, source: Union[str, Path, dict]) -> int:
    """Load prefix tensors into ``tuner`` from ``source``.

    Args:
        tuner: must already have the same prefix shape as the source
            (same n_prefix, n_heads, head_dim, n_layers). Otherwise
            ``load_state_dict`` will surface the mismatch.
        source: a path to a file produced by :func:`save_prefix_weights`,
            or a state-dict dict directly.

    Returns:
        The number of parameter tensors loaded.
    """
    sd = _resolve_source_to_state_dict(source, "load_prefix_weights")
    # Filter to prefix-only keys defensively so a full-model state-dict
    # accidentally passed in doesn't blow up the loader.
    sd = {k: v for k, v in sd.items() if "prefix_" in k}
    result = tuner.load_state_dict(sd, strict=False)
    # strict=False silently drops keys that don't exist on the tuner —
    # subtract them so the return value counts tensors that landed.
    return len(sd) - len(result.unexpected_keys)

9.5. Prompt Tuning

nnx.peft.prompt.PromptTuner

Bases: Module

Wrap a :class:TransformerNN with a learnable soft prompt.

Freezes every base parameter and allocates an (n_prompt_tokens, d_model) embedding tensor. The wrapper's forward prepends the prompt to the token embeddings, runs the stack, then trims the prompt positions off the logits before returning.

Parameters:

Name	Type	Description	Default
model	TransformerNN	a :class:TransformerNN instance. Its parameters are mutated in place (set to requires_grad=False).	required
n_prompt_tokens	int	number of soft-prompt slots. Must be > 0.	20

The soft prompt is initialized with nn.init.normal_(std=0.02) — the same scale Lester et al. use as their "random init" baseline.

Source code in nnx/peft/prompt.py

class PromptTuner(nn.Module):
    """Wrap a :class:`TransformerNN` with a learnable soft prompt.

    Freezes every base parameter and allocates an
    ``(n_prompt_tokens, d_model)`` embedding tensor. The wrapper's
    forward prepends the prompt to the token embeddings, runs the
    stack, then trims the prompt positions off the logits before
    returning.

    Args:
        model: a :class:`TransformerNN` instance. Its parameters are
            mutated in place (set to ``requires_grad=False``).
        n_prompt_tokens: number of soft-prompt slots. Must be > 0.

    The soft prompt is initialized with ``nn.init.normal_(std=0.02)``
    — the same scale Lester et al. use as their "random init" baseline.
    """

    def __init__(
        self,
        model: TransformerNN,
        *,
        n_prompt_tokens: int = 20,
    ):
        super().__init__()
        if not isinstance(model, TransformerNN):
            raise TypeError(f"PromptTuner requires a TransformerNN, got {type(model).__name__}")
        if n_prompt_tokens <= 0:
            raise ValueError(f"n_prompt_tokens must be positive, got {n_prompt_tokens}")

        self.model = model
        self.n_prompt_tokens = n_prompt_tokens

        # Freeze every parameter of the wrapped TransformerNN — the
        # soft prompt is the only trainable bit going forward.
        for p in self.model.parameters():
            p.requires_grad = False

        d_model = model.params.d_model
        # (n_prompt_tokens, d_model). Normal(0, 0.02) init.
        self.soft_prompt = nn.Parameter(torch.empty(n_prompt_tokens, d_model))
        nn.init.normal_(self.soft_prompt, std=0.02)

    @property
    def effective_max_seq_len(self) -> int:
        """Window available for REAL tokens: the soft prompt occupies
        ``n_prompt_tokens`` of the wrapped model's ``max_seq_len`` slots.
        ``GenerativeNNModel.generate`` reads this so its sliding window
        never overflows the wrapped model mid-generation."""
        return self.model.params.max_seq_len - self.n_prompt_tokens

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        """Run the wrapped model with the soft prompt prepended.

        Args:
            tokens: (batch, seq) long tensor of token ids.

        Returns:
            (batch, seq, vocab_size) logits over the REAL token
            positions only. The soft-prompt positions are scaffolding
            and their logits are discarded.

        Raises:
            ValueError: if ``seq + n_prompt_tokens`` exceeds the
                wrapped model's ``max_seq_len``. The soft prompt
                consumes positions in the RoPE table just like real
                tokens do.
        """
        b, t = tokens.shape
        total = t + self.n_prompt_tokens
        if total > self.model.params.max_seq_len:
            raise ValueError(
                f"input length ({t}) + soft prompt ({self.n_prompt_tokens}) = {total} "
                f"exceeds wrapped model's max_seq_len={self.model.params.max_seq_len}"
            )

        # Manual replication of TransformerNN.forward, but with the
        # soft prompt spliced in front of the token embeddings.
        x = self.model.tok_embed(tokens)  # (B, T, d_model)
        prompt = self.soft_prompt.unsqueeze(0).expand(b, -1, -1)  # (B, n_prompt, d_model)
        x = torch.cat([prompt, x], dim=1)  # (B, n_prompt + T, d_model)

        for block in self.model.blocks:
            x, _ = block(x, use_cache=False)
        x = self.model.norm_out(x)
        logits = self.model.lm_head(x)  # (B, n_prompt + T, vocab)
        # Trim the soft-prompt positions off the returned logits — the
        # caller only ever wants predictions for real input tokens.
        return logits[:, self.n_prompt_tokens :, :]

    def trainable_parameters(self) -> Iterator[nn.Parameter]:
        """Yield only the soft-prompt tensor.

        The wrapped model's parameters are frozen on construction; this
        is the iterator you hand to an optimizer.
        """
        yield self.soft_prompt

    def prompt_state_dict(self) -> dict:
        """Return a state-dict containing only the soft-prompt tensor,
        keyed for round-trip via :meth:`load_prompt_weights`."""
        return {k: v for k, v in self.state_dict().items() if "soft_prompt" in k}

effective_max_seq_len: int property

Window available for REAL tokens: the soft prompt occupies n_prompt_tokens of the wrapped model's max_seq_len slots. GenerativeNNModel.generate reads this so its sliding window never overflows the wrapped model mid-generation.

forward(tokens: torch.Tensor) -> torch.Tensor

Run the wrapped model with the soft prompt prepended.

Parameters:

Name	Type	Description	Default
tokens	Tensor	(batch, seq) long tensor of token ids.	required

Returns:

Type	Description
Tensor	(batch, seq, vocab_size) logits over the REAL token
Tensor	positions only. The soft-prompt positions are scaffolding
Tensor	and their logits are discarded.

Raises:

Type	Description
ValueError	if seq + n_prompt_tokens exceeds the wrapped model's max_seq_len. The soft prompt consumes positions in the RoPE table just like real tokens do.

Source code in nnx/peft/prompt.py

def forward(self, tokens: torch.Tensor) -> torch.Tensor:
    """Run the wrapped model with the soft prompt prepended.

    Args:
        tokens: (batch, seq) long tensor of token ids.

    Returns:
        (batch, seq, vocab_size) logits over the REAL token
        positions only. The soft-prompt positions are scaffolding
        and their logits are discarded.

    Raises:
        ValueError: if ``seq + n_prompt_tokens`` exceeds the
            wrapped model's ``max_seq_len``. The soft prompt
            consumes positions in the RoPE table just like real
            tokens do.
    """
    b, t = tokens.shape
    total = t + self.n_prompt_tokens
    if total > self.model.params.max_seq_len:
        raise ValueError(
            f"input length ({t}) + soft prompt ({self.n_prompt_tokens}) = {total} "
            f"exceeds wrapped model's max_seq_len={self.model.params.max_seq_len}"
        )

    # Manual replication of TransformerNN.forward, but with the
    # soft prompt spliced in front of the token embeddings.
    x = self.model.tok_embed(tokens)  # (B, T, d_model)
    prompt = self.soft_prompt.unsqueeze(0).expand(b, -1, -1)  # (B, n_prompt, d_model)
    x = torch.cat([prompt, x], dim=1)  # (B, n_prompt + T, d_model)

    for block in self.model.blocks:
        x, _ = block(x, use_cache=False)
    x = self.model.norm_out(x)
    logits = self.model.lm_head(x)  # (B, n_prompt + T, vocab)
    # Trim the soft-prompt positions off the returned logits — the
    # caller only ever wants predictions for real input tokens.
    return logits[:, self.n_prompt_tokens :, :]

prompt_state_dict() -> dict

Return a state-dict containing only the soft-prompt tensor, keyed for round-trip via :meth:load_prompt_weights.

Source code in nnx/peft/prompt.py

def prompt_state_dict(self) -> dict:
    """Return a state-dict containing only the soft-prompt tensor,
    keyed for round-trip via :meth:`load_prompt_weights`."""
    return {k: v for k, v in self.state_dict().items() if "soft_prompt" in k}

trainable_parameters() -> Iterator[nn.Parameter]

Yield only the soft-prompt tensor.

The wrapped model's parameters are frozen on construction; this is the iterator you hand to an optimizer.

Source code in nnx/peft/prompt.py

def trainable_parameters(self) -> Iterator[nn.Parameter]:
    """Yield only the soft-prompt tensor.

    The wrapped model's parameters are frozen on construction; this
    is the iterator you hand to an optimizer.
    """
    yield self.soft_prompt

nnx.peft.prompt.save_prompt_weights(tuner: PromptTuner, path: Union[str, Path]) -> str

Save ONLY the soft-prompt tensor of tuner to path.

Parameters:

Name	Type	Description	Default
tuner	PromptTuner	a :class:PromptTuner instance.	required
path	Union[str, Path]	destination file path.	required

Returns:

Type	Description
str	The path written, so calls can be chained.

Source code in nnx/peft/prompt.py

def save_prompt_weights(tuner: PromptTuner, path: Union[str, Path]) -> str:
    """Save ONLY the soft-prompt tensor of ``tuner`` to ``path``.

    Args:
        tuner: a :class:`PromptTuner` instance.
        path: destination file path.

    Returns:
        The path written, so calls can be chained.
    """
    sd = tuner.prompt_state_dict()
    torch.save(sd, str(path))
    return str(path)

nnx.peft.prompt.load_prompt_weights(tuner: PromptTuner, source: Union[str, Path, dict]) -> int

Load the soft-prompt tensor into tuner from source.

Parameters:

Name	Type	Description	Default
tuner	PromptTuner	must already have the same prompt shape as the source (same n_prompt_tokens, d_model). A shape mismatch is surfaced by load_state_dict.	required
source	Union[str, Path, dict]	a path to a file produced by :func:save_prompt_weights, or a state-dict dict directly.	required

Returns:

Type	Description
int	The number of parameter tensors loaded.

Source code in nnx/peft/prompt.py

def load_prompt_weights(tuner: PromptTuner, source: Union[str, Path, dict]) -> int:
    """Load the soft-prompt tensor into ``tuner`` from ``source``.

    Args:
        tuner: must already have the same prompt shape as the source
            (same n_prompt_tokens, d_model). A shape mismatch is
            surfaced by ``load_state_dict``.
        source: a path to a file produced by :func:`save_prompt_weights`,
            or a state-dict dict directly.

    Returns:
        The number of parameter tensors loaded.
    """
    sd = _resolve_source_to_state_dict(source, "load_prompt_weights")
    sd = {k: v for k, v in sd.items() if "soft_prompt" in k}
    result = tuner.load_state_dict(sd, strict=False)
    # strict=False silently drops keys that don't exist on the tuner —
    # subtract them so the return value counts tensors that landed.
    return len(sd) - len(result.unexpected_keys)

9.6. Adapters

nnx.peft.adapters.AdapterLayer

Bases: Module

Bottleneck residual block: y = x + up(act(down(x))).

up.weight and up.bias are zero-initialized so at step 0 the layer's output equals its input exactly. Gradient flow through up and down is unblocked from the first step; only the magnitude of the residual starts at zero.

Parameters:

Name	Type	Description	Default
dim	int	input and output feature dimension. The adapter is shape-preserving.	required
bottleneck	int	hidden dimension. Typically much smaller than dim (e.g., dim=768 → bottleneck=64 in the original Houlsby setup). Lower bottleneck = fewer params, potentially less expressive.	required
activation	Callable[[], Module]	nn.Module factory (called with no args inside __init__ to produce the activation module). Defaults to torch.nn.GELU — the modern adapter choice; nn.ReLU works too.	GELU

Source code in nnx/peft/adapters.py

class AdapterLayer(nn.Module):
    """Bottleneck residual block: ``y = x + up(act(down(x)))``.

    ``up.weight`` and ``up.bias`` are zero-initialized so at step 0
    the layer's output equals its input exactly. Gradient flow
    through ``up`` and ``down`` is unblocked from the first step;
    only the magnitude of the residual starts at zero.

    Args:
        dim: input and output feature dimension. The adapter is
            shape-preserving.
        bottleneck: hidden dimension. Typically much smaller than
            ``dim`` (e.g., dim=768 → bottleneck=64 in the original
            Houlsby setup). Lower bottleneck = fewer params,
            potentially less expressive.
        activation: ``nn.Module`` factory (called with no args inside
            ``__init__`` to produce the activation module). Defaults to
            ``torch.nn.GELU`` — the modern adapter choice; ``nn.ReLU``
            works too.
    """

    def __init__(
        self,
        dim: int,
        bottleneck: int,
        activation: Callable[[], nn.Module] = nn.GELU,
    ):
        super().__init__()
        if dim <= 0:
            raise ValueError(f"dim must be positive, got {dim}")
        if bottleneck <= 0:
            raise ValueError(f"bottleneck must be positive, got {bottleneck}")

        self.down = nn.Linear(dim, bottleneck)
        self.act = activation()
        self.up = nn.Linear(bottleneck, dim)

        # Zero-init the up-projection so the adapter is the identity
        # residual at step 0. Gradients still flow through `up` from
        # the first backward pass; only the OUTPUT magnitude starts
        # at zero.
        nn.init.zeros_(self.up.weight)
        nn.init.zeros_(self.up.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.up(self.act(self.down(x)))

    def extra_repr(self) -> str:
        return f"dim={self.down.in_features}, bottleneck={self.down.out_features}"

10. Pruning (nnx.prune)

nnx.prune.magnitude.magnitude_prune(net: nn.Module, sparsity: float, *, layer_pattern: str = '*', bake: bool = True) -> int

Zero the smallest-magnitude entries of every matched layer's weight.

For each :class:nn.Linear submodule of net whose dotted name matches layer_pattern (fnmatch glob), call :func:torch.nn.utils.prune.l1_unstructured with amount=sparsity. PyTorch's implementation zeros round(sparsity · weight.numel()) entries — the ones with the smallest absolute value — per layer.

Parameters:

Name	Type	Description	Default
net	Module	root module to walk. The function mutates net in place.	required
sparsity	float	fraction of weights to zero, in [0, 1). 0.0 is a valid no-op; 1.0 is rejected here — a fully-zeroed Linear is never useful (torch itself would accept amount=1.0 and silently zero the whole weight).	required
layer_pattern	str	fnmatch glob against dotted submodule name. "*" (the default) matches every :class:nn.Linear.	'*'
bake	bool	when True (default), call :func:torch.nn.utils.prune.remove immediately after each layer is pruned. The mask is baked into a plain weight tensor, the reparameterization is dropped, and the state_dict keys stay identical to the pre-prune layout. When False, the reparameterization stays in place — the state_dict carries <name>.weight_orig + <name>.weight_mask instead of <name>.weight. Use False for iterative pruning schedules (e.g., 10% per epoch for N epochs) where successive magnitude_prune calls need to compose with the existing mask.	True

Returns:

Type	Description
int	The number of :class:nn.Linear submodules that were pruned.
int	0 if layer_pattern matched nothing.

Raises:

Type	Description
ValueError	if sparsity is outside [0, 1).

Idempotency note: calling magnitude_prune twice at the same sparsity is a no-op for the second call — l1_unstructured picks the smallest-magnitude entries, which after the first prune are exactly the already-zeroed positions. The zero count stays the same; nothing is double-pruned.

Source code in nnx/prune/magnitude.py

def magnitude_prune(
    net: nn.Module,
    sparsity: float,
    *,
    layer_pattern: str = "*",
    bake: bool = True,
) -> int:
    """Zero the smallest-magnitude entries of every matched layer's weight.

    For each :class:`nn.Linear` submodule of ``net`` whose dotted name
    matches ``layer_pattern`` (fnmatch glob), call
    :func:`torch.nn.utils.prune.l1_unstructured` with
    ``amount=sparsity``. PyTorch's implementation zeros
    ``round(sparsity · weight.numel())`` entries — the ones with the
    smallest absolute value — per layer.

    Args:
        net: root module to walk. The function mutates ``net`` in place.
        sparsity: fraction of weights to zero, in ``[0, 1)``. ``0.0``
            is a valid no-op; ``1.0`` is rejected here — a fully-zeroed
            Linear is never useful (torch itself would accept
            ``amount=1.0`` and silently zero the whole weight).
        layer_pattern: fnmatch glob against dotted submodule name.
            ``"*"`` (the default) matches every :class:`nn.Linear`.
        bake: when ``True`` (default), call
            :func:`torch.nn.utils.prune.remove` immediately after each
            layer is pruned. The mask is baked into a plain ``weight``
            tensor, the reparameterization is dropped, and the
            ``state_dict`` keys stay identical to the pre-prune layout.
            When ``False``, the reparameterization stays in place — the
            ``state_dict`` carries ``<name>.weight_orig`` +
            ``<name>.weight_mask`` instead of ``<name>.weight``. Use
            ``False`` for iterative pruning schedules (e.g., 10% per
            epoch for N epochs) where successive ``magnitude_prune``
            calls need to compose with the existing mask.

    Returns:
        The number of :class:`nn.Linear` submodules that were pruned.
        ``0`` if ``layer_pattern`` matched nothing.

    Raises:
        ValueError: if ``sparsity`` is outside ``[0, 1)``.

    **Idempotency note:** calling ``magnitude_prune`` twice at the
    same ``sparsity`` is a no-op for the second call — l1_unstructured
    picks the smallest-magnitude entries, which after the first prune
    are exactly the already-zeroed positions. The zero count stays the
    same; nothing is double-pruned.
    """
    if not (0.0 <= sparsity < 1.0):
        raise ValueError(f"magnitude_prune sparsity must be in [0, 1), got {sparsity}")

    # Two-phase: collect targets first, then mutate. Iterating
    # named_modules() while reassigning child attributes via
    # `prune.l1_unstructured` registers a forward_pre_hook that doesn't
    # alter traversal order today, but collecting up-front keeps the
    # function defensive against any future PyTorch internals change.
    targets: list[tuple[str, nn.Linear]] = []
    for name, mod in net.named_modules():
        if isinstance(mod, nn.Linear) and fnmatch.fnmatchcase(name, layer_pattern):
            targets.append((name, mod))

    for _, mod in targets:
        # l1_unstructured at amount=0.0 is a valid no-op — torch still
        # registers the reparameterization but the mask is all-ones.
        # Calling remove afterwards (bake=True) restores the original
        # plain `weight` tensor with no change.
        torch_prune.l1_unstructured(mod, name="weight", amount=sparsity)
        if bake:
            # Bake the mask into a plain weight tensor. After this call
            # the layer's state_dict layout is identical to a fresh
            # nn.Linear's (weight + bias), so pruned checkpoints stay
            # shape-compatible with unpruned-network code.
            torch_prune.remove(mod, "weight")

    return len(targets)

nnx.prune.semi_structured.semi_structured_24(net: nn.Module, *, layer_pattern: str = '*') -> int

Swap each matched :class:nn.Linear's weight with a 2:4 semi-structured sparse tensor via :func:torchao.sparsity.sparsify_.

Parameters:

Name	Type	Description	Default
net	Module	root module to walk. The function mutates net in place.	required
layer_pattern	str	fnmatch glob against dotted submodule name. "*" (the default) matches every :class:nn.Linear.	'*'

Returns:

Type	Description
int	The number of :class:nn.Linear submodules that were swapped.
int	0 if layer_pattern matched nothing (in which case the
int	underlying torchao.sparsity.sparsify_ is NOT invoked — this
int	avoids an unnecessary torchao dispatch and the CUDA-only kernel
int	error on CPU runners with no Linear targets to swap).

Raises:

Type	Description
ImportError	if torchao isn't installed. The import happens inside the function body so :mod:nnx.prune doesn't pull torchao at package-import time; users on the magnitude-only path pay no dep cost.
RuntimeError	surfaced from the underlying torch.sparse.SparseSemiStructuredTensor constructor on unsupported hardware (CPU / pre-Ampere GPU) or on weights whose inner dimension isn't a multiple of 4. The error originates in torch / torchao; we don't intercept it.

Pattern semantics: same fnmatch convention as :func:nnx.peft.apply_lora_to and :func:nnx.prune.magnitude_prune — dotted submodule names against shell wildcards. Only :class:nn.Linear submodules are eligible (Conv2d / BatchNorm / Embedding / etc. are skipped even under a wildcard pattern).

Note on weights: torchao.sparsity.sparsify_ does NOT enforce the 2:4 mask before the swap. Callers are expected to either (a) magnitude-prune the weight to a valid 2:4 pattern beforehand (via :func:magnitude_prune or a custom mask), or (b) accept whatever 2:4 approximation :func:torch.sparse.to_sparse_semi_structured picks (which keeps the top-2-by-absolute-value entries per 4-group). For training workflows, the standard recipe is to pre-mask, then train the surviving entries.

Source code in nnx/prune/semi_structured.py

def semi_structured_24(net: nn.Module, *, layer_pattern: str = "*") -> int:
    """Swap each matched :class:`nn.Linear`'s weight with a 2:4
    semi-structured sparse tensor via :func:`torchao.sparsity.sparsify_`.

    Args:
        net: root module to walk. The function mutates ``net`` in place.
        layer_pattern: fnmatch glob against dotted submodule name.
            ``"*"`` (the default) matches every :class:`nn.Linear`.

    Returns:
        The number of :class:`nn.Linear` submodules that were swapped.
        ``0`` if ``layer_pattern`` matched nothing (in which case the
        underlying ``torchao.sparsity.sparsify_`` is NOT invoked — this
        avoids an unnecessary torchao dispatch and the CUDA-only kernel
        error on CPU runners with no Linear targets to swap).

    Raises:
        ImportError: if ``torchao`` isn't installed. The import happens
            inside the function body so :mod:`nnx.prune` doesn't pull
            torchao at package-import time; users on the magnitude-only
            path pay no dep cost.
        RuntimeError: surfaced from the underlying
            ``torch.sparse.SparseSemiStructuredTensor`` constructor on
            unsupported hardware (CPU / pre-Ampere GPU) or on weights
            whose inner dimension isn't a multiple of 4. The error
            originates in torch / torchao; we don't intercept it.

    **Pattern semantics:** same fnmatch convention as
    :func:`nnx.peft.apply_lora_to` and
    :func:`nnx.prune.magnitude_prune` — dotted submodule names against
    shell wildcards. Only :class:`nn.Linear` submodules are eligible
    (Conv2d / BatchNorm / Embedding / etc. are skipped even under a
    wildcard pattern).

    **Note on weights:** ``torchao.sparsity.sparsify_`` does NOT enforce
    the 2:4 mask before the swap. Callers are expected to either
    (a) magnitude-prune the weight to a valid 2:4 pattern beforehand
    (via :func:`magnitude_prune` or a custom mask), or
    (b) accept whatever 2:4 approximation
    :func:`torch.sparse.to_sparse_semi_structured` picks (which keeps
    the top-2-by-absolute-value entries per 4-group). For training
    workflows, the standard recipe is to pre-mask, then train the
    surviving entries.
    """
    # Defer the torchao import so consumers of `nnx.prune.magnitude_prune`
    # don't pay the torchao dep cost just by importing the package.
    import torchao.sparsity as ao_sparsity

    # Two-phase: enumerate the matched-Linear set FIRST so we can
    # return the count and short-circuit when there's nothing to do.
    # The same name list also drives the filter_fn we pass to torchao.
    targets: set[str] = set()
    for name, mod in net.named_modules():
        if isinstance(mod, nn.Linear) and fnmatch.fnmatchcase(name, layer_pattern):
            targets.add(name)

    if not targets:
        # No-op fast path: avoid invoking torchao at all when the
        # pattern matched nothing. Keeps the CPU-only error path off
        # the call graph for users who want to dry-run their pattern.
        return 0

    def _filter(mod: nn.Module, fqn: str) -> bool:
        # torchao calls this for every submodule. Accept only the
        # nn.Linear instances whose dotted name made the targets set
        # — re-checking the isinstance here is defensive (a future
        # torchao change could pass non-Linear modules through the
        # filter for non-Linear sparsity workflows).
        return isinstance(mod, nn.Linear) and fqn in targets

    ao_sparsity.sparsify_(net, ao_sparsity.semi_sparse_weight(), _filter)
    return len(targets)

11. Model surgery (nnx.surgery)

Walkthrough at Model surgery. Every primitive returns a fresh nn.Module and composes with NNModel.train() for the "load checkpoint → surgery → refine" loop.

nnx.surgery.widen.widen(model: nn.Module, *, layer_name: str, new_width: int, rng_seed: Optional[int] = 0) -> nn.Module

Net2WiderNet: grow a Linear's out_features to new_width.

Returns a deep copy of model with the named layer expanded and the downstream Linear's in_features adjusted so the overall forward output is preserved exactly (within FP rounding).

Parameters:

Name	Type	Description	Default
model	Module	any :class:nn.Module. The function deep-copies it so the caller's reference survives.	required
layer_name	str	dotted name (as produced by named_modules()) of the :class:nn.Linear to widen. Must be a Linear, must have an immediately downstream Linear, otherwise raises.	required
new_width	int	desired out_features. Must be strictly greater than the current out_features.	required
rng_seed	Optional[int]	seed for the unit-duplication choices. Pass an int for deterministic surgery; None to seed the local generator non-deterministically (fresh entropy — the global torch RNG is never read or advanced). Defaults to 0 so the primitive is deterministic by default.	0

Returns:

Type	Description
Module	A new :class:nn.Module (same class as model) with the
Module	widened Linear in place. Forward output equals the original's
Module	within atol=1e-5 (typically much tighter).

Raises:

Type	Description
KeyError	if layer_name is not a submodule of model.
TypeError	if the named submodule is not :class:nn.Linear.
ValueError	if new_width is not strictly greater than the current out_features, or if no downstream Linear exists.

Source code in nnx/surgery/widen.py

def widen(
    model: nn.Module,
    *,
    layer_name: str,
    new_width: int,
    rng_seed: Optional[int] = 0,
) -> nn.Module:
    """Net2WiderNet: grow a Linear's ``out_features`` to ``new_width``.

    Returns a deep copy of ``model`` with the named layer expanded and
    the downstream Linear's ``in_features`` adjusted so the overall
    forward output is preserved exactly (within FP rounding).

    Args:
        model: any :class:`nn.Module`. The function deep-copies it so
            the caller's reference survives.
        layer_name: dotted name (as produced by ``named_modules()``) of
            the :class:`nn.Linear` to widen. Must be a Linear, must
            have an immediately downstream Linear, otherwise raises.
        new_width: desired ``out_features``. Must be strictly greater
            than the current ``out_features``.
        rng_seed: seed for the unit-duplication choices. Pass an int
            for deterministic surgery; ``None`` to seed the local
            generator non-deterministically (fresh entropy — the global
            torch RNG is never read or advanced). Defaults to ``0`` so
            the primitive is deterministic by default.

    Returns:
        A new :class:`nn.Module` (same class as ``model``) with the
        widened Linear in place. Forward output equals the original's
        within ``atol=1e-5`` (typically much tighter).

    Raises:
        KeyError: if ``layer_name`` is not a submodule of ``model``.
        TypeError: if the named submodule is not :class:`nn.Linear`.
        ValueError: if ``new_width`` is not strictly greater than the
            current ``out_features``, or if no downstream Linear exists.
    """
    new_model = copy.deepcopy(model)
    layer = get_module(new_model, layer_name)
    if not isinstance(layer, nn.Linear):
        raise TypeError(f"widen target {layer_name!r} is {type(layer).__name__}, expected nn.Linear")
    cur = layer.out_features
    if new_width <= cur:
        raise ValueError(f"new_width must be > current out_features ({cur}); got {new_width}")
    q = new_width - cur

    # Pick q indices (with replacement) of existing units to duplicate.
    # A local generator keeps the surgery reproducible without touching
    # global RNG state.
    g = torch.Generator()
    if rng_seed is not None:
        g.manual_seed(int(rng_seed))
    else:
        g.seed()
    duplicates = torch.randint(0, cur, (q,), generator=g)

    # Track replication count per original index so we can divide the
    # downstream weight columns. Each "1" is the original column; each
    # appearance in `duplicates` adds one more.
    replication_count = torch.ones(cur, dtype=layer.weight.dtype, device=layer.weight.device)
    for idx in duplicates.tolist():
        replication_count[idx] += 1.0

    # --- Expand the target layer (in-out: in → cur+q) -----------------
    new_weight = torch.cat([layer.weight.data, layer.weight.data[duplicates]], dim=0)
    # skip_init: every param is fully overwritten below, so meta-device
    # construction avoids burning ambient RNG draws on a discarded init
    # (a seeded caller pipeline would otherwise silently diverge).
    new_layer = skip_init(
        nn.Linear,
        layer.in_features,
        new_width,
        bias=layer.bias is not None,
        device=layer.weight.device,
        dtype=layer.weight.dtype,
    )
    new_layer.weight.data.copy_(new_weight)
    if layer.bias is not None:
        new_bias = torch.cat([layer.bias.data, layer.bias.data[duplicates]], dim=0)
        new_layer.bias.data.copy_(new_bias)
    set_module(new_model, layer_name, new_layer)

    # --- Adjust the downstream Linear so the forward is preserved -----
    down_name, down_layer = _find_next_linear(new_model, layer_name)
    if down_layer.in_features != cur:
        raise ValueError(
            f"downstream Linear {down_name!r} has in_features={down_layer.in_features}, "
            f"expected {cur} to match {layer_name!r}'s old out_features. "
            "The two layers are not directly connected."
        )

    # New columns mirror columns from `duplicates`; rescale every column
    # by 1 / replication_count to preserve the forward sum.
    new_down_weight = torch.cat(
        [down_layer.weight.data, down_layer.weight.data[:, duplicates]],
        dim=1,
    )
    extended_rep = torch.cat([replication_count, replication_count[duplicates]])
    new_down_weight = new_down_weight / extended_rep.unsqueeze(0)

    new_down_layer = skip_init(
        nn.Linear,
        new_width,
        down_layer.out_features,
        bias=down_layer.bias is not None,
        device=down_layer.weight.device,
        dtype=down_layer.weight.dtype,
    )
    new_down_layer.weight.data.copy_(new_down_weight)
    if down_layer.bias is not None:
        # Bias is *additive after* W·x, so it doesn't change with the
        # column rescaling — copy as-is.
        new_down_layer.bias.data.copy_(down_layer.bias.data)
    set_module(new_model, down_name, new_down_layer)

    return new_model

nnx.surgery.deepen.deepen(model: nn.Module, *, after_layer_name: str) -> nn.Module

Net2DeeperNet: insert an identity-initialized Linear after the named layer. Function-preserving on ReLU networks only.

Parameters:

Name	Type	Description	Default
model	Module	any :class:nn.Module. Deep-copied so the caller's reference survives.	required
after_layer_name	str	dotted name (as in named_modules()) of the insertion site. Either: an :class:nn.ReLU inside a parent :class:nn.Sequential — the primitive splices Linear(I) → ReLU in after it. an :class:nn.Linear inside a parent :class:nn.ModuleList whose grandparent module declares ReLU as its activation (the FeedFwdNN contract) — the primitive inserts a new identity-init Linear into the ModuleList right after.	required

Returns:

Type	Description
Module	A fresh :class:nn.Module whose forward output matches the
Module	original within atol=1e-5.

Raises:

Type	Description
KeyError	if after_layer_name is not a submodule.
TypeError	if the layer is neither a ReLU-in-Sequential nor a Linear-in-FeedFwdNN-like ModuleList.
ValueError	if the parent's activation is anything other than ReLU. Sigmoid / tanh / GELU break function-preservation.

Source code in nnx/surgery/deepen.py

def deepen(
    model: nn.Module,
    *,
    after_layer_name: str,
) -> nn.Module:
    """Net2DeeperNet: insert an identity-initialized Linear after the
    named layer. Function-preserving on ReLU networks only.

    Args:
        model: any :class:`nn.Module`. Deep-copied so the caller's
            reference survives.
        after_layer_name: dotted name (as in ``named_modules()``) of
            the insertion site. Either:

              * an :class:`nn.ReLU` inside a parent :class:`nn.Sequential`
                — the primitive splices ``Linear(I) → ReLU`` in after it.
              * an :class:`nn.Linear` inside a parent :class:`nn.ModuleList`
                whose grandparent module declares ReLU as its activation
                (the FeedFwdNN contract) — the primitive inserts a new
                identity-init Linear into the ModuleList right after.

    Returns:
        A fresh :class:`nn.Module` whose forward output matches the
        original within ``atol=1e-5``.

    Raises:
        KeyError: if ``after_layer_name`` is not a submodule.
        TypeError: if the layer is neither a ReLU-in-Sequential nor a
            Linear-in-FeedFwdNN-like ModuleList.
        ValueError: if the parent's activation is anything other than
            ReLU. Sigmoid / tanh / GELU break function-preservation.
    """
    new_model = copy.deepcopy(model)
    target = get_module(new_model, after_layer_name)

    # Locate the parent container so we can splice.
    parent_path, _, attr = after_layer_name.rpartition(".")
    parent = new_model if not parent_path else new_model.get_submodule(parent_path)

    if isinstance(target, nn.ReLU) and isinstance(parent, nn.Sequential):
        return _insert_after_relu_in_sequential(new_model, parent_path, attr)

    if isinstance(target, nn.Linear) and isinstance(parent, nn.ModuleList):
        return _insert_after_linear_in_module_list(new_model, parent_path, attr, target)

    raise TypeError(
        f"deepen: cannot insert after {after_layer_name!r} "
        f"({type(target).__name__} in a {type(parent).__name__} parent). "
        "Supported sites: nn.ReLU inside nn.Sequential, or nn.Linear "
        "inside an nn.ModuleList whose parent applies an activation."
    )

nnx.surgery.drop_layer.drop_layer(model: nn.Module, *, layer_name: Union[str, list[str]], importance: Callable[[nn.Module], float] | None = None) -> nn.Module

Replace a named layer with :class:nn.Identity.

Parameters:

Name	Type	Description	Default
model	Module	any :class:nn.Module. Deep-copied so the caller's reference survives.	required
layer_name	Union[str, list[str]]	either a dotted submodule name, or a list of dotted names to choose from. When a list is given, importance must be provided as well.	required
importance	Callable[[Module], float] | None	optional callable fn(submodule) -> float. When layer_name is a list, the candidate with the minimum importance score is dropped (lowest = least informative = safest to remove). Rejected when layer_name is a single string because there is no candidate-selection step to score.	None

Returns:

Type	Description
Module	A fresh :class:nn.Module with the chosen layer replaced by
Module	class:nn.Identity. Forward shape contract is preserved iff
Module	the dropped layer was shape-preserving (e.g. an activation or
Module	a square Linear); otherwise calling forward on the surged
Module	module will raise — by design, since silently corrupting the
Module	shape would be worse than a loud failure.

Raises:

Type	Description
KeyError	if any candidate name is missing.
ValueError	if layer_name is an empty list, or a list without importance, or a single string with importance.

Source code in nnx/surgery/drop_layer.py

def drop_layer(
    model: nn.Module,
    *,
    layer_name: Union[str, list[str]],
    importance: Callable[[nn.Module], float] | None = None,
) -> nn.Module:
    """Replace a named layer with :class:`nn.Identity`.

    Args:
        model: any :class:`nn.Module`. Deep-copied so the caller's
            reference survives.
        layer_name: either a dotted submodule name, or a list of
            dotted names to choose from. When a list is given,
            ``importance`` must be provided as well.
        importance: optional callable ``fn(submodule) -> float``. When
            ``layer_name`` is a list, the candidate with the *minimum*
            importance score is dropped (lowest = least informative =
            safest to remove). Rejected when ``layer_name`` is a single
            string because there is no candidate-selection step to score.

    Returns:
        A fresh :class:`nn.Module` with the chosen layer replaced by
        :class:`nn.Identity`. Forward shape contract is preserved iff
        the dropped layer was shape-preserving (e.g. an activation or
        a square Linear); otherwise calling forward on the surged
        module will raise — by design, since silently corrupting the
        shape would be worse than a loud failure.

    Raises:
        KeyError: if any candidate name is missing.
        ValueError: if ``layer_name`` is an empty list, or a list
            without ``importance``, or a single string with
            ``importance``.
    """
    chosen = _resolve_target(model, layer_name, importance)

    new_model = copy.deepcopy(model)
    # Verify the chosen name still resolves in the copy — it should,
    # since deepcopy preserves the module tree structure.
    get_module(new_model, chosen)
    set_module(new_model, chosen, nn.Identity())
    return new_model

nnx.surgery.low_rank.low_rank_factorize(linear: nn.Linear, *, rank: int, method: str = 'svd') -> nn.Sequential

Factor a Linear into two smaller Linears via rank-k SVD truncation.

Parameters:

Name	Type	Description	Default
linear	Linear	an :class:nn.Linear to factorize. Its weights are read but not mutated — the returned Sequential is a fresh pair of Linears.	required
rank	int	the truncation rank k. Must be in [1, min(out, in)]. When k == min(out, in) the factorization is exact.	required
method	str	"svd" (the only option in v1). Reserved for the future "activation_svd" / "fisher" variants.	'svd'

Returns:

Type	Description
Sequential	class:nn.Sequential of two Linears whose composition
Sequential	approximates the input layer. The first Linear has
Sequential	bias=False (Sx@V.T has no native bias term); the second
Sequential	Linear carries the original bias verbatim.

Raises:

Type	Description
TypeError	if linear is not :class:nn.Linear.
ValueError	if rank is out of range, or method unknown.

Source code in nnx/surgery/low_rank.py

def low_rank_factorize(
    linear: nn.Linear,
    *,
    rank: int,
    method: str = "svd",
) -> nn.Sequential:
    """Factor a Linear into two smaller Linears via rank-``k`` SVD
    truncation.

    Args:
        linear: an :class:`nn.Linear` to factorize. Its weights are
            read but not mutated — the returned Sequential is a fresh
            pair of Linears.
        rank: the truncation rank ``k``. Must be in ``[1, min(out, in)]``.
            When ``k == min(out, in)`` the factorization is exact.
        method: ``"svd"`` (the only option in v1). Reserved for the
            future ``"activation_svd"`` / ``"fisher"`` variants.

    Returns:
        :class:`nn.Sequential` of two Linears whose composition
        approximates the input layer. The first Linear has
        ``bias=False`` (Sx@V.T has no native bias term); the second
        Linear carries the original bias verbatim.

    Raises:
        TypeError: if ``linear`` is not :class:`nn.Linear`.
        ValueError: if ``rank`` is out of range, or ``method`` unknown.
    """
    if not isinstance(linear, nn.Linear):
        raise TypeError(f"low_rank_factorize requires an nn.Linear, got {type(linear).__name__}")
    if method not in _SUPPORTED_METHODS:
        raise ValueError(f"unsupported method {method!r}; supported: {_SUPPORTED_METHODS}")

    out_features = linear.out_features
    in_features = linear.in_features
    max_rank = min(out_features, in_features)
    if not (1 <= rank <= max_rank):
        raise ValueError(f"rank must be in [1, min(out_features, in_features) = {max_rank}]; got {rank}")

    # Run SVD on the float64-promoted weight to keep the reconstruction
    # numerically clean for the "exact at max rank" test. The factored
    # Linears are downcast back to the original dtype before return.
    W = linear.weight.data
    orig_dtype = W.dtype
    U, S, Vh = torch.linalg.svd(W.to(torch.float64), full_matrices=False)
    U_k = U[:, :rank]  # (out, k)
    S_k = S[:rank]  # (k,)
    Vh_k = Vh[:rank, :]  # (k, in)

    # First Linear: in → k. Its weight (k, in) is (S_k * Vh_k); no bias
    # since the original bias is added after the second matmul.
    down_weight = (S_k.unsqueeze(1) * Vh_k).to(orig_dtype)
    # device= threads the original layer's placement — without it a
    # CUDA-resident model gets CPU layers spliced in and the next
    # forward crashes with a device mismatch (widen already does this).
    # skip_init: every param is fully overwritten, so meta-device
    # construction keeps the surgery off the global RNG stream.
    down = skip_init(nn.Linear, in_features, rank, bias=False, dtype=orig_dtype, device=W.device)
    down.weight.data.copy_(down_weight)

    # Second Linear: k → out. Weight is U_k. Bias is the original.
    up_weight = U_k.to(orig_dtype)
    up = skip_init(nn.Linear, rank, out_features, bias=linear.bias is not None, dtype=orig_dtype, device=W.device)
    up.weight.data.copy_(up_weight)
    if linear.bias is not None:
        up.bias.data.copy_(linear.bias.data)

    return nn.Sequential(down, up)

nnx.surgery.embedding.expand_embedding(emb: nn.Embedding, *, new_num_embeddings: int, init: InitStrategy = 'zeros') -> tuple[nn.Embedding, torch.Tensor]

Return a larger Embedding whose first rows match emb exactly.

Parameters:

Name	Type	Description	Default
emb	Embedding	the source embedding. Its weights are read but not mutated.	required
new_num_embeddings	int	the desired num_embeddings for the returned layer. Must be strictly greater than the current.	required
init	InitStrategy	how to initialize the new rows. "zeros" — fill with zeros (default; deterministic, safe). "copy_mean" — fill each new row with the per-column mean of the original rows.	'zeros'

Returns:

Type	Description
Embedding	(new_emb, frozen_mask) where new_emb is a fresh
Tensor	class:nn.Embedding with the original rows preserved, and
tuple[Embedding, Tensor]	frozen_mask is a bool tensor of shape
tuple[Embedding, Tensor]	(new_num_embeddings,) marking the original rows (True)
tuple[Embedding, Tensor]	as candidates for freezing during refinement.

Raises:

Type	Description
TypeError	if emb is not :class:nn.Embedding.
ValueError	if new_num_embeddings is not strictly greater, or if init is unknown.

Source code in nnx/surgery/embedding.py

def expand_embedding(
    emb: nn.Embedding,
    *,
    new_num_embeddings: int,
    init: InitStrategy = "zeros",
) -> tuple[nn.Embedding, torch.Tensor]:
    """Return a larger Embedding whose first rows match ``emb`` exactly.

    Args:
        emb: the source embedding. Its weights are read but not
            mutated.
        new_num_embeddings: the desired ``num_embeddings`` for the
            returned layer. Must be strictly greater than the current.
        init: how to initialize the new rows. ``"zeros"`` — fill with
            zeros (default; deterministic, safe). ``"copy_mean"`` —
            fill each new row with the per-column mean of the original
            rows.

    Returns:
        ``(new_emb, frozen_mask)`` where ``new_emb`` is a fresh
        :class:`nn.Embedding` with the original rows preserved, and
        ``frozen_mask`` is a bool tensor of shape
        ``(new_num_embeddings,)`` marking the original rows (``True``)
        as candidates for freezing during refinement.

    Raises:
        TypeError: if ``emb`` is not :class:`nn.Embedding`.
        ValueError: if ``new_num_embeddings`` is not strictly greater,
            or if ``init`` is unknown.
    """
    if not isinstance(emb, nn.Embedding):
        raise TypeError(f"expand_embedding requires an nn.Embedding, got {type(emb).__name__}")
    old_num = emb.num_embeddings
    if new_num_embeddings <= old_num:
        raise ValueError(f"new_num_embeddings must be > current ({old_num}); got {new_num_embeddings}")
    if init not in ("zeros", "copy_mean"):
        raise ValueError(f"unknown init strategy {init!r}; expected 'zeros' or 'copy_mean'")

    dim = emb.embedding_dim
    # skip_init: every row is overwritten below (originals copied, new
    # rows zeroed / mean-filled), so meta-device construction avoids
    # consuming global RNG on a discarded normal_ init.
    new_emb = skip_init(
        nn.Embedding,
        num_embeddings=new_num_embeddings,
        embedding_dim=dim,
        padding_idx=emb.padding_idx,
        max_norm=emb.max_norm,
        norm_type=emb.norm_type,
        scale_grad_by_freq=emb.scale_grad_by_freq,
        sparse=emb.sparse,
        dtype=emb.weight.dtype,
        device=emb.weight.device,
    )
    with torch.no_grad():
        # Preserve the original rows exactly.
        new_emb.weight[:old_num].copy_(emb.weight.data)

        # Initialize the new rows by the chosen strategy.
        if init == "zeros":
            new_emb.weight[old_num:].zero_()
        elif init == "copy_mean":
            # Mean across rows of the original embedding — a per-column
            # vector — broadcast into every new row.
            row_mean = emb.weight.data.mean(dim=0, keepdim=True)  # (1, dim)
            new_emb.weight[old_num:].copy_(row_mean.expand(new_num_embeddings - old_num, dim))

    # Frozen mask: True for rows that came from the original embedding.
    frozen_mask = torch.zeros(new_num_embeddings, dtype=torch.bool, device=emb.weight.device)
    frozen_mask[:old_num] = True
    return new_emb, frozen_mask

12. Quantization (nnx.quantize)

PTQ INT8 weight-only + QAT 8da4w via torchao (the replacement for the removed torch.ao.quantization). Opt-in via pip install "thekaveh-nnx[quantize]".

nnx.quantize.ptq.quantize_int8(model: NNModel) -> NNModel

Return a new :class:NNModel with int8 weight-only quantized net.

Deep-copies model.net and applies torchao.quantization.quantize_(net, Int8WeightOnlyConfig()) to the copy. Every nn.Linear submodule of the copy has its weight parameter replaced with an :class:AffineQuantizedTensor (int8 per-channel, symmetric). Activations stay FP32 — only the weights are stored in int8.

The original model is untouched. The returned NNModel shares every other attribute (params, net_params, device, loss_fn) with the original — only net is the quantized copy.

Parameters:

Name	Type	Description	Default
model	NNModel	a trained :class:NNModel. PTQ has no training step; this function is a pure post-process.	required

Returns:

Type	Description
NNModel	a new :class:NNModel instance whose net is the quantized
NNModel	deep-copy of model.net. The new model can be used for
NNModel	predict / evaluate / to_onnx exactly like the
NNModel	original; train on the quantized model is not supported
NNModel	(QAT lands in a separate module).

Raises:

Type	Description
ImportError	if torchao is not installed. Install with pip install thekaveh-nnx[quantize].

Source code in nnx/quantize/ptq.py

def quantize_int8(model: NNModel) -> NNModel:
    """Return a new :class:`NNModel` with int8 weight-only quantized ``net``.

    Deep-copies ``model.net`` and applies
    ``torchao.quantization.quantize_(net, Int8WeightOnlyConfig())`` to
    the copy. Every ``nn.Linear`` submodule of the copy has its weight
    parameter replaced with an :class:`AffineQuantizedTensor` (int8
    per-channel, symmetric). Activations stay FP32 — only the weights
    are stored in int8.

    The original ``model`` is untouched. The returned ``NNModel`` shares
    every other attribute (``params``, ``net_params``, ``device``,
    ``loss_fn``) with the original — only ``net`` is the quantized copy.

    Args:
        model: a trained :class:`NNModel`. PTQ has no training step;
            this function is a pure post-process.

    Returns:
        a new :class:`NNModel` instance whose ``net`` is the quantized
        deep-copy of ``model.net``. The new model can be used for
        ``predict`` / ``evaluate`` / ``to_onnx`` exactly like the
        original; ``train`` on the quantized model is not supported
        (QAT lands in a separate module).

    Raises:
        ImportError: if ``torchao`` is not installed. Install with
            ``pip install thekaveh-nnx[quantize]``.
    """
    try:
        from torchao.quantization import Int8WeightOnlyConfig, quantize_
    except ImportError as e:  # pragma: no cover — opt-in extra
        raise ImportError(
            "quantize_int8 requires torchao. Install with `pip install thekaveh-nnx[quantize]` "
            "(or `pip install 'torchao>=0.17'` directly)."
        ) from e

    # Operate on a copy so the caller's original NNModel stays untouched —
    # lets users keep both around for an accuracy delta comparison.
    quantized_net = deepcopy(model.net)
    # version=2 silences the v1-deprecation UserWarning torchao emits. The
    # default granularity is PerRow (per-output-channel), which is the
    # standard weight-only INT8 layout.
    quantize_(quantized_net, Int8WeightOnlyConfig(version=2))

    # Build a sibling NNModel without re-running __init__ (which would
    # build a fresh FP32 net via self.params.net(...)). Share every other
    # attribute so loss_fn / device / params line up with the original.
    m = NNModel.__new__(NNModel)
    m.__dict__.update(model.__dict__)
    m.net = quantized_net
    return m

nnx.quantize.qat.qat_train_step_factory(base_step: Optional[TrainStepFn] = None, qat_config: str = '8da4w') -> TrainStepFn

Return a :class:TrainStepFn that runs base_step against a fake-quantized model.

The returned step is the same as base_step (or :func:default_train_step when base_step is None) — fake-quant insertion happens once, via :class:QATLifecycleCallback, on on_train_begin. The per-batch forward/backward then exercises those fake-quant ops automatically through the standard module forward.

Why split the work between a factory and a callback?

• The factory validates qat_config early (at construction time) so misconfigurations surface before the data loader spins up.
• The callback owns the lifecycle: prepare at start, convert at end. Bundling that into the per-batch step would re-check the module state every iteration and complicate gradient flow.

Both pieces are needed in :meth:NNModel.train::

callback = QATLifecycleCallback(qat_config="8da4w")
step_fn  = qat_train_step_factory(qat_config="8da4w")
model.train(params=..., callbacks=[callback], train_step_fn=step_fn)

Parameters:

Name	Type	Description	Default
base_step	Optional[TrainStepFn]	optional underlying training step to wrap. None (the default) uses :func:default_train_step — the standard supervised forward/backward. Pass a custom step here to combine QAT with e.g. knowledge distillation or mixup; the fake-quant ops live in the model graph, so any standard step picks them up transparently.	None
qat_config	str	shortcut for the torchao QAT recipe. Currently only "8da4w" is supported (int8 dynamic activations + int4 grouped weights). Validated eagerly so a typo doesn't propagate to the callback.	'8da4w'

Returns:

Name	Type	Description
a	TrainStepFn	class:TrainStepFn ready to pass to
	TrainStepFn	NNModel.train(..., train_step_fn=...).

Raises:

Type	Description
ValueError	if qat_config is not in :data:_SUPPORTED_CONFIGS.
ImportError	if torchao is not installed.

Source code in nnx/quantize/qat.py

def qat_train_step_factory(
    base_step: Optional[TrainStepFn] = None,
    qat_config: str = "8da4w",
) -> TrainStepFn:
    """Return a :class:`TrainStepFn` that runs ``base_step`` against a
    fake-quantized model.

    The returned step is the *same* as ``base_step`` (or
    :func:`default_train_step` when ``base_step`` is None) — fake-quant
    insertion happens once, via :class:`QATLifecycleCallback`, on
    ``on_train_begin``. The per-batch forward/backward then exercises
    those fake-quant ops automatically through the standard module
    forward.

    Why split the work between a factory and a callback?

    - The factory validates ``qat_config`` early (at construction time)
      so misconfigurations surface before the data loader spins up.
    - The callback owns the lifecycle: ``prepare`` at start, ``convert``
      at end. Bundling that into the per-batch step would re-check the
      module state every iteration and complicate gradient flow.

    Both pieces are needed in :meth:`NNModel.train`::

        callback = QATLifecycleCallback(qat_config="8da4w")
        step_fn  = qat_train_step_factory(qat_config="8da4w")
        model.train(params=..., callbacks=[callback], train_step_fn=step_fn)

    Args:
        base_step: optional underlying training step to wrap. ``None``
            (the default) uses :func:`default_train_step` — the standard
            supervised forward/backward. Pass a custom step here to
            combine QAT with e.g. knowledge distillation or mixup; the
            fake-quant ops live in the model graph, so any standard
            step picks them up transparently.
        qat_config: shortcut for the torchao QAT recipe. Currently only
            ``"8da4w"`` is supported (int8 dynamic activations + int4
            grouped weights). Validated eagerly so a typo doesn't
            propagate to the callback.

    Returns:
        a :class:`TrainStepFn` ready to pass to
        ``NNModel.train(..., train_step_fn=...)``.

    Raises:
        ValueError: if ``qat_config`` is not in
            :data:`_SUPPORTED_CONFIGS`.
        ImportError: if ``torchao`` is not installed.
    """
    # Validate eagerly — fail at factory construction time so the
    # caller's `model.train(...)` invocation doesn't blow up two minutes
    # into a long-running data load.
    _build_quantizer(qat_config)

    return base_step if base_step is not None else default_train_step

nnx.quantize.qat.QATLifecycleCallback

Bases: Callback

Manage the torchao prepare / convert lifecycle around training.

Add to callbacks=[...] in :meth:NNModel.train. On train begin, swaps every eligible :class:torch.nn.Linear in model.net for its fake-quantized counterpart (the model now learns to be robust to int4/int8 rounding). On train end, the fake-quantized linears are converted to actually-quantized ones — the resulting model is suitable for inference / export.

The mutation is in place on model.net: after training, model.net IS the converted model. The callback exposes the quantizer instance as self.quantizer for callers who want to pickle quantizer-specific state alongside their checkpoint, and tracks the prepare/convert phase via self.is_prepared and self.is_converted for downstream inspection.

Parameters:

Name	Type	Description	Default
qat_config	str	torchao recipe shortcut. See :func:qat_train_step_factory.	'8da4w'
groupsize	int	group size for the int4 weight quantizer. 32 is the default — small enough to apply to toy nets in tests (where hidden_dim=64) while being a real-world setting. Larger groupsizes (128, 256) give better compression at the cost of accuracy.	32

Source code in nnx/quantize/qat.py

class QATLifecycleCallback(Callback):
    """Manage the torchao ``prepare`` / ``convert`` lifecycle around training.

    Add to ``callbacks=[...]`` in :meth:`NNModel.train`. On train begin,
    swaps every eligible :class:`torch.nn.Linear` in ``model.net`` for
    its fake-quantized counterpart (the model now learns to be robust
    to int4/int8 rounding). On train end, the fake-quantized linears
    are converted to actually-quantized ones — the resulting model is
    suitable for inference / export.

    The mutation is **in place** on ``model.net``: after training,
    ``model.net`` IS the converted model. The callback exposes the
    quantizer instance as ``self.quantizer`` for callers who want to
    pickle quantizer-specific state alongside their checkpoint, and
    tracks the prepare/convert phase via ``self.is_prepared`` and
    ``self.is_converted`` for downstream inspection.

    Args:
        qat_config: torchao recipe shortcut. See
            :func:`qat_train_step_factory`.
        groupsize: group size for the int4 weight quantizer. 32 is the
            default — small enough to apply to toy nets in tests
            (where hidden_dim=64) while being a real-world setting.
            Larger groupsizes (128, 256) give better compression at
            the cost of accuracy.
    """

    def __init__(self, qat_config: str = "8da4w", *, groupsize: int = 32):
        self.qat_config = qat_config
        self.groupsize = groupsize
        # Resolve the quantizer eagerly so a torchao ImportError or a
        # bad config name surfaces at callback construction, not on
        # ``on_train_begin`` (which would otherwise crash mid-train).
        self.quantizer = _build_quantizer(qat_config, groupsize=groupsize)
        self.is_prepared: bool = False
        self.is_converted: bool = False

    def on_train_begin(self, ctx: _CallbackContext) -> None:
        """Insert fake-quant ops into ``ctx.model.net`` in place."""
        # Idempotency guard: re-running prepare on an already-prepared
        # net would corrupt the linear-wrapping. Real-world training
        # loops don't hit this, but a callback shared across two
        # back-to-back ``model.train()`` calls would.
        if self.is_prepared:
            return
        self.quantizer.prepare(ctx.model.net)
        self.is_prepared = True

    def on_train_end(self, ctx: _CallbackContext) -> None:
        """Convert fake-quant ops in ``ctx.model.net`` to true int4/int8 modules.

        After this returns, ``ctx.model.net`` produces real quantized
        outputs and is suitable for inference / ONNX export. The model
        is no longer trainable through the usual FP32 optimizer path —
        a fresh training session on the same NNModel would need a new
        QATLifecycleCallback.
        """
        # Symmetric idempotency guard, paired with ``on_train_begin``.
        if self.is_converted:
            return
        if not self.is_prepared:
            # No prepare ever ran — converting an FP32 model is a no-op
            # that would still mutate state-flags misleadingly. Bail.
            return
        self.quantizer.convert(ctx.model.net)
        self.is_converted = True

on_train_begin(ctx: _CallbackContext) -> None

Insert fake-quant ops into ctx.model.net in place.

Source code in nnx/quantize/qat.py

def on_train_begin(self, ctx: _CallbackContext) -> None:
    """Insert fake-quant ops into ``ctx.model.net`` in place."""
    # Idempotency guard: re-running prepare on an already-prepared
    # net would corrupt the linear-wrapping. Real-world training
    # loops don't hit this, but a callback shared across two
    # back-to-back ``model.train()`` calls would.
    if self.is_prepared:
        return
    self.quantizer.prepare(ctx.model.net)
    self.is_prepared = True

on_train_end(ctx: _CallbackContext) -> None

Convert fake-quant ops in ctx.model.net to true int4/int8 modules.

After this returns, ctx.model.net produces real quantized outputs and is suitable for inference / ONNX export. The model is no longer trainable through the usual FP32 optimizer path — a fresh training session on the same NNModel would need a new QATLifecycleCallback.

Source code in nnx/quantize/qat.py

def on_train_end(self, ctx: _CallbackContext) -> None:
    """Convert fake-quant ops in ``ctx.model.net`` to true int4/int8 modules.

    After this returns, ``ctx.model.net`` produces real quantized
    outputs and is suitable for inference / ONNX export. The model
    is no longer trainable through the usual FP32 optimizer path —
    a fresh training session on the same NNModel would need a new
    QATLifecycleCallback.
    """
    # Symmetric idempotency guard, paired with ``on_train_begin``.
    if self.is_converted:
        return
    if not self.is_prepared:
        # No prepare ever ran — converting an FP32 model is a no-op
        # that would still mutate state-flags misleadingly. Bail.
        return
    self.quantizer.convert(ctx.model.net)
    self.is_converted = True

13. Diffusion (nnx.diffusion)

nnx.diffusion.schedules.NoiseSchedulers

Bases: Enum

Diffusion noise-schedule factory. Enum-as-factory pattern (like :class:nnx.Nets, :class:nnx.Optims): each enum variant's __call__ constructs the underlying :class:NoiseSchedule.

Source code in nnx/diffusion/schedules.py

class NoiseSchedulers(Enum):
    """Diffusion noise-schedule factory. Enum-as-factory pattern (like
    :class:`nnx.Nets`, :class:`nnx.Optims`): each enum variant's
    ``__call__`` constructs the underlying :class:`NoiseSchedule`."""

    LINEAR = "linear"
    COSINE = "cosine"

    def __str__(self) -> str:
        return self.value

    def __repr__(self) -> str:
        return str(self)

    def __call__(
        self,
        T: int = 1000,
        *,
        beta_min: float = 1e-4,
        beta_max: float = 2e-2,
        s: float = 0.008,
    ) -> NoiseSchedule:
        """Build a :class:`NoiseSchedule`.

        Args:
            T: number of diffusion timesteps. Larger T means more steps
                during training (one t per batch) and at sampling time
                (T-many reverse passes); 1000 is the DDPM default.
            beta_min: LINEAR schedule lower endpoint; ignored for COSINE.
            beta_max: LINEAR schedule upper endpoint; ignored for COSINE.
            s: COSINE schedule offset (Improved DDPM eq. 17); ignored for LINEAR.

        Returns:
            A :class:`NoiseSchedule` on CPU. Call ``.to(device)`` after
            construction to migrate; the diffusion train step does this
            implicitly when indexing the schedule tensors with batch
            timesteps on the target device.
        """
        if T <= 0:
            raise ValueError(f"T must be positive, got {T}")
        match self:
            case NoiseSchedulers.LINEAR:
                if not (0 < beta_min < beta_max < 1):
                    raise ValueError(f"LINEAR schedule needs 0 < beta_min ({beta_min}) < beta_max ({beta_max}) < 1")
                betas = torch.linspace(beta_min, beta_max, T, dtype=torch.float32)
                return _from_betas(self, betas)
            case NoiseSchedulers.COSINE:
                if s <= 0:
                    raise ValueError(f"COSINE schedule needs s > 0, got {s}")
                # Improved DDPM eq. 17. Build ᾱ_t first, then back out the
                # betas. Clamp betas to [1e-8, 0.999] so the sqrt() and
                # posterior_variance ops can't blow up near t=T-1.
                steps = torch.arange(T + 1, dtype=torch.float64) / T
                f_t = torch.cos((steps + s) / (1 + s) * torch.pi / 2) ** 2
                alphas_cumprod = (f_t / f_t[0]).to(torch.float32)
                betas = (1.0 - alphas_cumprod[1:] / alphas_cumprod[:-1]).clamp(
                    min=1e-8,
                    max=0.999,
                )
                return _from_betas(self, betas)

__call__(T: int = 1000, *, beta_min: float = 0.0001, beta_max: float = 0.02, s: float = 0.008) -> NoiseSchedule

Build a :class:NoiseSchedule.

Parameters:

Name	Type	Description	Default
T	int	number of diffusion timesteps. Larger T means more steps during training (one t per batch) and at sampling time (T-many reverse passes); 1000 is the DDPM default.	1000
beta_min	float	LINEAR schedule lower endpoint; ignored for COSINE.	0.0001
beta_max	float	LINEAR schedule upper endpoint; ignored for COSINE.	0.02
s	float	COSINE schedule offset (Improved DDPM eq. 17); ignored for LINEAR.	0.008

Returns:

Name	Type	Description
A	NoiseSchedule	class:NoiseSchedule on CPU. Call .to(device) after
	NoiseSchedule	construction to migrate; the diffusion train step does this
	NoiseSchedule	implicitly when indexing the schedule tensors with batch
	NoiseSchedule	timesteps on the target device.

Source code in nnx/diffusion/schedules.py

def __call__(
    self,
    T: int = 1000,
    *,
    beta_min: float = 1e-4,
    beta_max: float = 2e-2,
    s: float = 0.008,
) -> NoiseSchedule:
    """Build a :class:`NoiseSchedule`.

    Args:
        T: number of diffusion timesteps. Larger T means more steps
            during training (one t per batch) and at sampling time
            (T-many reverse passes); 1000 is the DDPM default.
        beta_min: LINEAR schedule lower endpoint; ignored for COSINE.
        beta_max: LINEAR schedule upper endpoint; ignored for COSINE.
        s: COSINE schedule offset (Improved DDPM eq. 17); ignored for LINEAR.

    Returns:
        A :class:`NoiseSchedule` on CPU. Call ``.to(device)`` after
        construction to migrate; the diffusion train step does this
        implicitly when indexing the schedule tensors with batch
        timesteps on the target device.
    """
    if T <= 0:
        raise ValueError(f"T must be positive, got {T}")
    match self:
        case NoiseSchedulers.LINEAR:
            if not (0 < beta_min < beta_max < 1):
                raise ValueError(f"LINEAR schedule needs 0 < beta_min ({beta_min}) < beta_max ({beta_max}) < 1")
            betas = torch.linspace(beta_min, beta_max, T, dtype=torch.float32)
            return _from_betas(self, betas)
        case NoiseSchedulers.COSINE:
            if s <= 0:
                raise ValueError(f"COSINE schedule needs s > 0, got {s}")
            # Improved DDPM eq. 17. Build ᾱ_t first, then back out the
            # betas. Clamp betas to [1e-8, 0.999] so the sqrt() and
            # posterior_variance ops can't blow up near t=T-1.
            steps = torch.arange(T + 1, dtype=torch.float64) / T
            f_t = torch.cos((steps + s) / (1 + s) * torch.pi / 2) ** 2
            alphas_cumprod = (f_t / f_t[0]).to(torch.float32)
            betas = (1.0 - alphas_cumprod[1:] / alphas_cumprod[:-1]).clamp(
                min=1e-8,
                max=0.999,
            )
            return _from_betas(self, betas)

nnx.diffusion.schedules.NoiseSchedule dataclass

Precomputed DDPM noise schedule.

All tensors are 1D of length T and live on the same device. The factory constructs them on CPU; :meth:to returns a new schedule with every tensor migrated.

Attributes:

Name	Type	Description
kind	NoiseSchedulers	which enum variant produced this schedule (for introspection).
T	int	number of diffusion timesteps.
betas	Tensor	per-step variance, shape=(T,).
alphas	Tensor	1 - betas.
alphas_cumprod	Tensor	cumulative product of alphas (ᾱ_t in the paper).
sqrt_alphas_cumprod	Tensor	√ᾱ_t — the x_0 coefficient in q(x_t | x_0).
sqrt_one_minus_alphas_cumprod	Tensor	√(1 - ᾱ_t) — the noise coefficient.
posterior_variance	Tensor	variance of q(x_{t-1} | x_t, x_0), used by the reverse-step sampler.

Source code in nnx/diffusion/schedules.py

@dataclass(frozen=True, slots=True)
class NoiseSchedule:
    """Precomputed DDPM noise schedule.

    All tensors are 1D of length ``T`` and live on the same device. The
    factory constructs them on CPU; :meth:`to` returns a new schedule
    with every tensor migrated.

    Attributes:
        kind: which enum variant produced this schedule (for introspection).
        T: number of diffusion timesteps.
        betas: per-step variance, ``shape=(T,)``.
        alphas: ``1 - betas``.
        alphas_cumprod: cumulative product of alphas (``ᾱ_t`` in the paper).
        sqrt_alphas_cumprod: ``√ᾱ_t`` — the x_0 coefficient in q(x_t | x_0).
        sqrt_one_minus_alphas_cumprod: ``√(1 - ᾱ_t)`` — the noise coefficient.
        posterior_variance: variance of q(x_{t-1} | x_t, x_0), used by the
            reverse-step sampler.
    """

    kind: NoiseSchedulers
    T: int
    betas: torch.Tensor
    alphas: torch.Tensor
    alphas_cumprod: torch.Tensor
    sqrt_alphas_cumprod: torch.Tensor
    sqrt_one_minus_alphas_cumprod: torch.Tensor
    posterior_variance: torch.Tensor

    def to(self, device) -> NoiseSchedule:
        """Return a copy with every tensor moved to ``device``. The kind
        and T fields are unchanged."""
        return NoiseSchedule(
            kind=self.kind,
            T=self.T,
            betas=self.betas.to(device),
            alphas=self.alphas.to(device),
            alphas_cumprod=self.alphas_cumprod.to(device),
            sqrt_alphas_cumprod=self.sqrt_alphas_cumprod.to(device),
            sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod.to(device),
            posterior_variance=self.posterior_variance.to(device),
        )

to(device) -> NoiseSchedule

Return a copy with every tensor moved to device. The kind and T fields are unchanged.

Source code in nnx/diffusion/schedules.py

def to(self, device) -> NoiseSchedule:
    """Return a copy with every tensor moved to ``device``. The kind
    and T fields are unchanged."""
    return NoiseSchedule(
        kind=self.kind,
        T=self.T,
        betas=self.betas.to(device),
        alphas=self.alphas.to(device),
        alphas_cumprod=self.alphas_cumprod.to(device),
        sqrt_alphas_cumprod=self.sqrt_alphas_cumprod.to(device),
        sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod.to(device),
        posterior_variance=self.posterior_variance.to(device),
    )

nnx.diffusion.nets.DiffusionMLP

Bases: Module

Conditional MLP for low-dim diffusion: forward(x_t, t) -> ε_pred.

Architecture: sinusoidal time embed → small projection → concat with flat x_t → MLP → linear head producing a noise prediction of the same shape as x_t. Bare ReLU activations, no skip connections — a single file's worth of code, enough to learn a 2D Gaussian mixture or a small tabular distribution.

Inputs of any rank are supported by flattening dimensions ≥ 1 before the MLP and un-flattening at the output. The network is NOT a U-Net — it has no spatial structure. For image-space diffusion, the same train/sample/schedule machinery works against a user-supplied U-Net.

Source code in nnx/diffusion/nets.py

class DiffusionMLP(nn.Module):
    """Conditional MLP for low-dim diffusion: ``forward(x_t, t) -> ε_pred``.

    Architecture: sinusoidal time embed → small projection → concat with
    flat x_t → MLP → linear head producing a noise prediction of the same
    shape as x_t. Bare ReLU activations, no skip connections — a single
    file's worth of code, enough to learn a 2D Gaussian mixture or a
    small tabular distribution.

    Inputs of any rank are supported by flattening dimensions ≥ 1 before
    the MLP and un-flattening at the output. The network is *NOT* a
    U-Net — it has no spatial structure. For image-space diffusion, the
    same train/sample/schedule machinery works against a user-supplied
    U-Net.
    """

    def __init__(
        self,
        input_dim: int,
        hidden_dims: list[int] | None = None,
        time_embed_dim: int = 32,
    ):
        super().__init__()
        if input_dim <= 0:
            raise ValueError(f"input_dim must be positive, got {input_dim}")
        # Same fail-fast as input_dim: an odd/non-positive embed dim
        # otherwise only surfaces at the first forward, mid-training.
        if time_embed_dim <= 0 or time_embed_dim % 2 != 0:
            raise ValueError(f"time_embed_dim must be a positive even integer, got {time_embed_dim}")
        if hidden_dims is None:
            hidden_dims = [128, 128]
        # Same fail-fast + convention as NNParams.hidden_dims: a non-positive
        # entry otherwise builds a zero/negative-width nn.Linear silently. An
        # empty list is valid (a single input+t -> input projection).
        if not all(d > 0 for d in hidden_dims):
            raise ValueError(f"hidden_dims entries must be positive, got {hidden_dims}")
        self.input_dim = input_dim
        self.time_embed_dim = time_embed_dim

        # Project the sinusoidal embedding through one Linear+GELU before
        # concatenation — DDPM uses an MLP, but a single layer is enough
        # for low-dim cases and keeps the param count minimal.
        self.time_mlp = nn.Sequential(
            nn.Linear(time_embed_dim, time_embed_dim),
            nn.GELU(),
        )

        layer_dims = [input_dim + time_embed_dim, *hidden_dims, input_dim]
        layers = []
        for in_d, out_d in zip(layer_dims[:-1], layer_dims[1:], strict=True):
            layers.append(nn.Linear(in_d, out_d))
        self.layers = nn.ModuleList(layers)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """Predict noise added to ``x`` at timestep ``t``.

        Args:
            x: ``(B, *)`` clean shape; flattened internally to ``(B, D)``.
            t: ``(B,)`` integer timesteps.

        Returns:
            Tensor of the same shape as ``x``.
        """
        orig_shape = x.shape
        B = x.shape[0]
        x_flat = x.reshape(B, -1)
        if x_flat.shape[1] != self.input_dim:
            raise ValueError(f"DiffusionMLP expects flattened input dim {self.input_dim}, got {x_flat.shape[1]}")

        # Time conditioning: sinusoidal embed → MLP → concat with x.
        t_emb = sinusoidal_time_embed(t, self.time_embed_dim)
        t_emb = self.time_mlp(t_emb)

        h = torch.cat([x_flat, t_emb], dim=-1)
        for layer in self.layers[:-1]:
            h = torch.relu(layer(h))
        out = self.layers[-1](h)

        return out.reshape(orig_shape)

    def unpack_batch(self, batch):
        """Standard ``(X-tuple, Y)`` adapter so this net plays nicely with
        the NNx dataloader contract. ``Y`` is unused by diffusion — every
        consumer that calls ``unpack_batch`` discards it."""
        if isinstance(batch, (list, tuple)):
            x, y = batch[0], batch[1] if len(batch) > 1 else None
        else:
            x, y = batch, None
        return (x,), y

forward(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor

Predict noise added to x at timestep t.

Parameters:

Name	Type	Description	Default
x	Tensor	(B, *) clean shape; flattened internally to (B, D).	required
t	Tensor	(B,) integer timesteps.	required

Returns:

Type	Description
Tensor	Tensor of the same shape as x.

Source code in nnx/diffusion/nets.py

def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Predict noise added to ``x`` at timestep ``t``.

    Args:
        x: ``(B, *)`` clean shape; flattened internally to ``(B, D)``.
        t: ``(B,)`` integer timesteps.

    Returns:
        Tensor of the same shape as ``x``.
    """
    orig_shape = x.shape
    B = x.shape[0]
    x_flat = x.reshape(B, -1)
    if x_flat.shape[1] != self.input_dim:
        raise ValueError(f"DiffusionMLP expects flattened input dim {self.input_dim}, got {x_flat.shape[1]}")

    # Time conditioning: sinusoidal embed → MLP → concat with x.
    t_emb = sinusoidal_time_embed(t, self.time_embed_dim)
    t_emb = self.time_mlp(t_emb)

    h = torch.cat([x_flat, t_emb], dim=-1)
    for layer in self.layers[:-1]:
        h = torch.relu(layer(h))
    out = self.layers[-1](h)

    return out.reshape(orig_shape)

unpack_batch(batch)

Standard (X-tuple, Y) adapter so this net plays nicely with the NNx dataloader contract. Y is unused by diffusion — every consumer that calls unpack_batch discards it.

Source code in nnx/diffusion/nets.py

def unpack_batch(self, batch):
    """Standard ``(X-tuple, Y)`` adapter so this net plays nicely with
    the NNx dataloader contract. ``Y`` is unused by diffusion — every
    consumer that calls ``unpack_batch`` discards it."""
    if isinstance(batch, (list, tuple)):
        x, y = batch[0], batch[1] if len(batch) > 1 else None
    else:
        x, y = batch, None
    return (x,), y

nnx.diffusion.nets.sinusoidal_time_embed(t: torch.Tensor, dim: int) -> torch.Tensor

Standard transformer-style sinusoidal positional embedding, applied to scalar timesteps so the denoising network can condition on t.

Parameters:

Name	Type	Description	Default
t	Tensor	integer or float tensor of shape (B,) — per-sample timesteps.	required
dim	int	embedding dimension. Half of it carries sin frequencies, half carries cos; dim must be even.	required

Returns:

Type	Description
Tensor	Tensor of shape (B, dim).

Source code in nnx/diffusion/nets.py

def sinusoidal_time_embed(t: torch.Tensor, dim: int) -> torch.Tensor:
    """Standard transformer-style sinusoidal positional embedding,
    applied to scalar timesteps so the denoising network can condition
    on ``t``.

    Args:
        t: integer or float tensor of shape ``(B,)`` — per-sample timesteps.
        dim: embedding dimension. Half of it carries sin frequencies,
            half carries cos; ``dim`` must be even.

    Returns:
        Tensor of shape ``(B, dim)``.
    """
    if dim % 2 != 0:
        raise ValueError(f"sinusoidal_time_embed dim must be even, got {dim}")
    half = dim // 2
    # Inverse-frequency scaling, matching the original Transformer paper
    # and ho:DDPM. half == 1 (dim=2) degenerates to a single unit
    # frequency — special-cased so the divisor isn't zero.
    decay = math.log(10000.0) / (half - 1) if half > 1 else 0.0
    freqs = torch.exp(-decay * torch.arange(half, dtype=torch.float32, device=t.device))
    args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
    return torch.cat([args.sin(), args.cos()], dim=-1)

nnx.diffusion.training.diffusion_train_step_factory(schedule: NoiseSchedule) -> TrainStepFn

Build a DDPM noise-prediction :class:TrainStepFn.

Each call to the returned step fn:

1. Samples a random per-sample timestep t ~ Uniform[0, T).
2. Samples Gaussian noise ε ~ N(0, I) matching x_0's shape.
3. Computes x_t = √ᾱ_t · x_0 + √(1 - ᾱ_t) · ε (forward diffusion).
4. Calls model.net(x_t, t) to predict ε_pred.
5. Backprops the MSE between ε_pred and ε, steps the optimizer.

Loss is reported as both .loss and .error on the returned EDP so BEST checkpoint tracking and the ReduceLROnPlateau scheduler have a metric to lock onto. The standard supervised classification metrics (accuracy/f1/...) are not meaningful for a generative paradigm and stay zero.

Parameters:

Name	Type	Description	Default
schedule	NoiseSchedule	a :class:NoiseSchedule from :class:NoiseSchedulers. Built on any device; the step fn lazily migrates the indexed tensors to model.device per call.	required

Returns:

Type	Description
TrainStepFn	A function suitable for NNModel.train(..., train_step_fn=...).

Source code in nnx/diffusion/training.py

def diffusion_train_step_factory(schedule: NoiseSchedule) -> TrainStepFn:
    """Build a DDPM noise-prediction :class:`TrainStepFn`.

    Each call to the returned step fn:

      1. Samples a random per-sample timestep ``t ~ Uniform[0, T)``.
      2. Samples Gaussian noise ``ε ~ N(0, I)`` matching x_0's shape.
      3. Computes ``x_t = √ᾱ_t · x_0 + √(1 - ᾱ_t) · ε`` (forward diffusion).
      4. Calls ``model.net(x_t, t)`` to predict ``ε_pred``.
      5. Backprops the MSE between ``ε_pred`` and ``ε``, steps the optimizer.

    Loss is reported as both ``.loss`` and ``.error`` on the returned
    EDP so BEST checkpoint tracking and the ReduceLROnPlateau scheduler
    have a metric to lock onto. The standard supervised classification
    metrics (accuracy/f1/...) are not meaningful for a generative
    paradigm and stay zero.

    Args:
        schedule: a :class:`NoiseSchedule` from :class:`NoiseSchedulers`.
            Built on any device; the step fn lazily migrates the
            indexed tensors to ``model.device`` per call.

    Returns:
        A function suitable for ``NNModel.train(..., train_step_fn=...)``.
    """

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        # Standard NNx dataloader contract: batch is (X, Y) or a single
        # tensor; Y is unused by diffusion. Use unpack_batch when the net
        # supplies one, fall back to direct unpacking otherwise.
        if hasattr(m.net, "unpack_batch"):
            (x_0,), _ = m.net.unpack_batch(ctx.batch)
        elif isinstance(ctx.batch, (list, tuple)):
            x_0 = ctx.batch[0]
        else:
            x_0 = ctx.batch
        x_0 = x_0.to(m.device)

        B = x_0.shape[0]
        t = torch.randint(0, schedule.T, (B,), device=m.device)
        eps = torch.randn_like(x_0)

        sqrt_a = _extract(schedule.sqrt_alphas_cumprod, t, x_0.shape)
        sqrt_1ma = _extract(schedule.sqrt_one_minus_alphas_cumprod, t, x_0.shape)
        x_t = sqrt_a * x_0 + sqrt_1ma * eps

        eps_pred = m.net(x_t, t)
        loss = F.mse_loss(eps_pred, eps)
        loss_val = finalize_step(loss, ctx, paradigm="diffusion")

        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=0.0,
            precision=0.0,
            loss=loss_val,
            error=loss_val,
        )

    return step

nnx.diffusion.sampling.sample(model: NNModel, schedule: NoiseSchedule, shape: tuple[int, ...], *, device: Optional[torch.device] = None, generator: Optional[torch.Generator] = None) -> torch.Tensor

Run T reverse-diffusion steps and return samples drawn from the distribution the model was trained on.

Parameters:

Name	Type	Description	Default
model	NNModel	an :class:NNModel whose .net is the trained denoising network (e.g., :class:DiffusionMLP or any forward(x, t) -> ε module).	required
schedule	NoiseSchedule	the same :class:NoiseSchedule used during training. Indexed tensors are moved to device lazily.	required
shape	tuple[int, ...]	full tensor shape to generate, e.g., (256, 2) for 256 2D samples.	required
device	Optional[device]	target device. Defaults to model.device.	None
generator	Optional[Generator]	optional torch.Generator for reproducible sampling (pass one built with torch.Generator(device).manual_seed(...)).	None

Returns:

Type	Description
Tensor	A tensor of shape shape carrying the generated samples.

Source code in nnx/diffusion/sampling.py

@torch.no_grad()
def sample(
    model: NNModel,
    schedule: NoiseSchedule,
    shape: tuple[int, ...],
    *,
    device: Optional[torch.device] = None,
    generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
    """Run T reverse-diffusion steps and return samples drawn from the
    distribution the model was trained on.

    Args:
        model: an :class:`NNModel` whose ``.net`` is the trained
            denoising network (e.g., :class:`DiffusionMLP` or any
            ``forward(x, t) -> ε`` module).
        schedule: the same :class:`NoiseSchedule` used during training.
            Indexed tensors are moved to ``device`` lazily.
        shape: full tensor shape to generate, e.g., ``(256, 2)`` for
            256 2D samples.
        device: target device. Defaults to ``model.device``.
        generator: optional torch.Generator for reproducible sampling
            (pass one built with ``torch.Generator(device).manual_seed(...)``).

    Returns:
        A tensor of shape ``shape`` carrying the generated samples.
    """
    if device is None:
        device = model.device

    # Migrate the schedule's indexed tensors once so the per-step loop
    # doesn't re-transfer them.
    sched = schedule.to(device)
    # Snapshot training-mode for non-destructive restore (matches
    # NNModel.predict / evaluate, GenerativeNNModel.generate,
    # nnx.viz.activation_map, and nnx.lr_finder).
    was_training = model.net.training
    model.net.eval()

    try:
        x = torch.randn(*shape, device=device, generator=generator)

        for t_int in reversed(range(sched.T)):
            # Batch-broadcast the timestep so the network's t-conditioning
            # works whether called with B==1 or B==shape[0].
            t = torch.full((shape[0],), t_int, dtype=torch.long, device=device)

            eps_pred = model.net(x, t)

            beta_t = sched.betas[t_int]
            alpha_t = sched.alphas[t_int]
            alpha_bar_t = sched.alphas_cumprod[t_int]

            # Posterior mean (DDPM eq. 11): predict x_0 implicitly through eps,
            # then the posterior of q(x_{t-1} | x_t, x_0).
            coef = beta_t / (1.0 - alpha_bar_t).sqrt()
            mean = (1.0 / alpha_t.sqrt()) * (x - coef * eps_pred)

            if t_int > 0:
                # Standard DDPM: re-inject noise scaled by the posterior std
                # at every step except the final t=0 (the boundary case is a
                # deterministic mean — adding noise would not match the training
                # objective).
                noise = (
                    torch.randn_like(x)
                    if generator is None
                    else torch.randn(
                        x.shape,
                        generator=generator,
                        device=device,
                        dtype=x.dtype,
                    )
                )
                x = mean + sched.posterior_variance[t_int].sqrt() * noise
            else:
                x = mean

        return x
    finally:
        if was_training:
            model.net.train()

14. Training paradigms (nnx.paradigms)

Each factory returns a TrainStepFn for the train_step_fn= hook on NNModel.train. The training loop, checkpoint cadence, callbacks, and persistence are unchanged — only the per-batch update is swapped.

14.1. Knowledge distillation

nnx.paradigms.distillation.kd_train_step_factory(teacher: NNModel, *, alpha: float = 0.5, temperature: float = 4.0) -> TrainStepFn

Build a knowledge-distillation :class:TrainStepFn.

Parameters:

Name	Type	Description	Default
teacher	NNModel	a fully-trained :class:NNModel whose net produces logits of the same shape as the student's. The teacher's parameters are frozen (requires_grad=False) and its net is set to eval mode on factory call.	required
alpha	float	weight on the distillation (soft) loss. The hard-label loss gets 1 − α. α=1.0 is pure distillation; α=0.0 collapses to standard supervised training (the teacher is loaded but unused). 0.5 is the common default.	0.5
temperature	float	softmax temperature applied to BOTH student and teacher logits before the KL. Higher T flattens the distribution and exposes more dark knowledge; the × T² factor in front of the KL keeps gradient magnitude comparable to the hard-label term across T. 4.0 is the classical Hinton choice.	4.0

Returns:

Type	Description
TrainStepFn	A TrainStepFn suitable for NNModel.train(..., train_step_fn=...).

Raises:

Type	Description
ValueError	if alpha is not in [0, 1], or temperature ≤ 0.

Source code in nnx/paradigms/distillation.py

def kd_train_step_factory(
    teacher: NNModel,
    *,
    alpha: float = 0.5,
    temperature: float = 4.0,
) -> TrainStepFn:
    """Build a knowledge-distillation :class:`TrainStepFn`.

    Args:
        teacher: a fully-trained :class:`NNModel` whose net produces
            logits of the same shape as the student's. The teacher's
            parameters are frozen (``requires_grad=False``) and its
            net is set to eval mode on factory call.
        alpha: weight on the distillation (soft) loss. The hard-label
            loss gets ``1 − α``. ``α=1.0`` is pure distillation;
            ``α=0.0`` collapses to standard supervised training (the
            teacher is loaded but unused). 0.5 is the common default.
        temperature: softmax temperature applied to BOTH student and
            teacher logits before the KL. Higher T flattens the
            distribution and exposes more dark knowledge; the
            ``× T²`` factor in front of the KL keeps gradient
            magnitude comparable to the hard-label term across T.
            4.0 is the classical Hinton choice.

    Returns:
        A ``TrainStepFn`` suitable for ``NNModel.train(..., train_step_fn=...)``.

    Raises:
        ValueError: if ``alpha`` is not in [0, 1], or ``temperature`` ≤ 0.
    """
    if not (0.0 <= alpha <= 1.0):
        raise ValueError(f"alpha must be in [0, 1], got {alpha}")
    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")

    # Freeze the teacher and pin to eval mode. The student's training
    # never touches the teacher; this just guards against accidental
    # gradient flow if the caller wires them into a shared module later.
    teacher.net.eval()
    for p in teacher.net.parameters():
        p.requires_grad = False

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        # unpack_batch returns ((X,), Y); the singleton destructure asserts
        # a single-input net and binds X to that one tensor.
        (X,), Y = m.net.unpack_batch(ctx.batch)
        X = X.to(m.device)
        Y = Y.to(m.device)

        student_logits = m.net(X)
        with torch.no_grad():
            # Teacher might live on a different device; migrate the
            # input into its frame for the forward pass. Cheap when
            # student.device == teacher.device.
            teacher_logits = teacher.net(X.to(teacher.device)).to(m.device)

        # KL(teacher || student) with both softened by T — the standard
        # Hinton direction (see softened_kl for the F.kl_div contract).
        soft_loss = softened_kl(student_logits, teacher_logits, temperature)
        hard_loss = m.loss_fn(student_logits, Y)
        loss = alpha * soft_loss + (1.0 - alpha) * hard_loss
        loss_val = finalize_step(loss, ctx, paradigm="distillation")

        return classification_edp(
            Y=Y,
            Y_hat=student_logits.argmax(dim=-1),
            loss=loss_val,
            extra_metrics=ctx.extra_metrics,
        )

    return step

nnx.paradigms.distillation.feature_kd_train_step_factory(teacher: NNModel, *, auxiliary_layers: dict[str, str], alpha: float = 0.5, beta: float = 0.5, temperature: float = 4.0) -> TrainStepFn

Build a FitNets-style feature-distillation :class:TrainStepFn.

Extends :func:kd_train_step_factory with an additional MSE term matching named intermediate-layer activations between the (frozen) teacher and the trainable student. Forward hooks register on the pairs in auxiliary_layers; collected activations feed an elementwise MSE that's mixed into the loss via beta::

L = α · KL_soft · T² + β · MSE(student_act, teacher_act) + (1 − α) · L_hard

Parameters:

Name	Type	Description	Default
teacher	NNModel	a fully-trained :class:NNModel whose net produces logits of the same shape as the student's. The teacher's parameters are frozen (requires_grad=False) and its net is set to eval mode on factory call — same guarantee as :func:kd_train_step_factory.	required
auxiliary_layers	dict[str, str]	dict mapping teacher_layer_name -> student_layer_name for each (teacher, student) pair to match. Names are dotted paths resolved via :meth:torch.nn.Module.get_submodule against the teacher / student net. Must be non-empty. The teacher and student activations at each pair must share shape — if they don't, the factory raises ValueError on the first forward (the projector FeatureRegressor from FitNets is intentionally deferred).	required
alpha	float	weight on the soft (logit-KL) term. The hard-label loss gets 1 − α. 0.5 is the common default.	0.5
beta	float	weight on the feature-MSE term. 0.5 is the common starting point; tune downward if it dominates the logit term, upward to bias the student toward matching internal representations.	0.5
temperature	float	softmax temperature for the logit-KL term — identical contract to :func:kd_train_step_factory.	4.0

Returns:

Type	Description
TrainStepFn	A TrainStepFn suitable for ``NNModel.train(...,
TrainStepFn	train_step_fn=...)``.

Raises:

Type	Description
ValueError	if alpha or beta is not in [0, 1], if temperature ≤ 0, or if auxiliary_layers is empty. On the first batch, if any paired teacher/student activation shapes disagree.

Source code in nnx/paradigms/distillation.py

def feature_kd_train_step_factory(
    teacher: NNModel,
    *,
    auxiliary_layers: dict[str, str],
    alpha: float = 0.5,
    beta: float = 0.5,
    temperature: float = 4.0,
) -> TrainStepFn:
    """Build a FitNets-style feature-distillation :class:`TrainStepFn`.

    Extends :func:`kd_train_step_factory` with an additional MSE term
    matching named intermediate-layer activations between the (frozen)
    teacher and the trainable student. Forward hooks register on the
    pairs in ``auxiliary_layers``; collected activations feed an
    elementwise MSE that's mixed into the loss via ``beta``::

        L = α · KL_soft · T² + β · MSE(student_act, teacher_act) + (1 − α) · L_hard

    Args:
        teacher: a fully-trained :class:`NNModel` whose net produces
            logits of the same shape as the student's. The teacher's
            parameters are frozen (``requires_grad=False``) and its
            net is set to eval mode on factory call — same guarantee
            as :func:`kd_train_step_factory`.
        auxiliary_layers: dict mapping ``teacher_layer_name ->
            student_layer_name`` for each (teacher, student) pair to
            match. Names are dotted paths resolved via
            :meth:`torch.nn.Module.get_submodule` against the teacher
            / student ``net``. Must be non-empty. The teacher and
            student activations at each pair must share shape — if
            they don't, the factory raises ``ValueError`` on the
            first forward (the projector ``FeatureRegressor`` from
            FitNets is intentionally deferred).
        alpha: weight on the soft (logit-KL) term. The hard-label
            loss gets ``1 − α``. 0.5 is the common default.
        beta: weight on the feature-MSE term. 0.5 is the common
            starting point; tune downward if it dominates the logit
            term, upward to bias the student toward matching internal
            representations.
        temperature: softmax temperature for the logit-KL term —
            identical contract to :func:`kd_train_step_factory`.

    Returns:
        A ``TrainStepFn`` suitable for ``NNModel.train(...,
        train_step_fn=...)``.

    Raises:
        ValueError: if ``alpha`` or ``beta`` is not in [0, 1], if
            ``temperature`` ≤ 0, or if ``auxiliary_layers`` is empty.
            On the first batch, if any paired teacher/student
            activation shapes disagree.
    """
    if not (0.0 <= alpha <= 1.0):
        raise ValueError(f"alpha must be in [0, 1], got {alpha}")
    if not (0.0 <= beta <= 1.0):
        raise ValueError(f"beta must be in [0, 1], got {beta}")
    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")
    if not auxiliary_layers:
        raise ValueError(
            "auxiliary_layers must be a non-empty mapping of teacher_layer_name -> student_layer_name; got empty dict"
        )

    # Freeze the teacher and pin to eval mode — same guarantee as
    # kd_train_step_factory. Student training never touches the
    # teacher; this just guards against accidental gradient flow.
    teacher.net.eval()
    for p in teacher.net.parameters():
        p.requires_grad = False

    # Resolve named submodules eagerly so a typo in the user's mapping
    # raises a clear error on factory call (not deep inside the first
    # forward where the AttributeError would be cryptic).
    teacher_layers: dict[str, torch.nn.Module] = {}
    for t_name in auxiliary_layers:
        try:
            teacher_layers[t_name] = teacher.net.get_submodule(t_name)
        except AttributeError as e:
            raise ValueError(
                f"auxiliary_layers: teacher has no submodule named {t_name!r} (get_submodule raised: {e})"
            ) from e

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        (X,), Y = m.net.unpack_batch(ctx.batch)
        X = X.to(m.device)
        Y = Y.to(m.device)

        # Resolve student layers per-step (cheap, and keeps us robust
        # if the user swapped m.net between factory build and first
        # step). Same typo-surfacing as the teacher resolution above.
        student_acts: dict[str, torch.Tensor] = {}
        teacher_acts: dict[str, torch.Tensor] = {}
        handles: list[torch.utils.hooks.RemovableHandle] = []
        try:
            for t_name, s_name in auxiliary_layers.items():
                try:
                    s_layer = m.net.get_submodule(s_name)
                except AttributeError as e:
                    raise ValueError(
                        f"auxiliary_layers: student has no submodule named {s_name!r} (get_submodule raised: {e})"
                    ) from e

                # Capture by default-arg to bind the loop's current names
                # rather than the closure's late-binding.
                def _t_hook(_module, _inputs, output, _key=t_name):
                    teacher_acts[_key] = output

                def _s_hook(_module, _inputs, output, _key=t_name):
                    student_acts[_key] = output

                handles.append(teacher_layers[t_name].register_forward_hook(_t_hook))
                handles.append(s_layer.register_forward_hook(_s_hook))

            student_logits = m.net(X)
            with torch.no_grad():
                # Teacher might live on a different device; migrate the
                # input into its frame for the forward pass.
                teacher_logits = teacher.net(X.to(teacher.device)).to(m.device)
        finally:
            for h in handles:
                h.remove()

        # Verify every paired activation has matching shape. The
        # FeatureRegressor (a small projection) is deferred — for now
        # we surface a clear error so the user doesn't get cryptic
        # broadcast results from F.mse_loss.
        feature_loss = torch.zeros((), device=m.device)
        for t_name in auxiliary_layers:
            t_act = teacher_acts[t_name].to(m.device)
            s_act = student_acts[t_name]
            if t_act.shape != s_act.shape:
                raise ValueError(
                    f"auxiliary_layers: shape mismatch at pair "
                    f"{t_name!r} -> {auxiliary_layers[t_name]!r} — "
                    f"teacher activation has shape {tuple(t_act.shape)} "
                    f"but student activation has shape {tuple(s_act.shape)}. "
                    "The v1 factory requires shape-matched paired layers; "
                    "the FeatureRegressor projector is deferred."
                )
            feature_loss = feature_loss + F.mse_loss(s_act, t_act)
        # Average across paired layers so beta's scale is invariant
        # to how many pairs the user provided.
        feature_loss = feature_loss / len(auxiliary_layers)

        soft_loss = softened_kl(student_logits, teacher_logits, temperature)
        hard_loss = m.loss_fn(student_logits, Y)
        loss = alpha * soft_loss + beta * feature_loss + (1.0 - alpha) * hard_loss
        loss_val = finalize_step(loss, ctx, paradigm="feature_distillation")

        return classification_edp(
            Y=Y,
            Y_hat=student_logits.argmax(dim=-1),
            loss=loss_val,
            extra_metrics=ctx.extra_metrics,
        )

    return step

nnx.paradigms.born_again.born_again_train(model: NNModel, *, generations: int = 3, train_params: NNTrainParams, **kd_kwargs: Any) -> list[NNRun]

Iterate G generations of self-distillation on a single model.

Generation 0 trains plain (no teacher) — standard supervised loss. Each subsequent generation uses a deep-copied, frozen, eval-mode snapshot of the model after the prior generation completed as the teacher for a Hinton-style KD step (via :func:kd_train_step_factory).

The same in-place model is reused across generations; only the teacher snapshot is duplicated. This matches the original paper's setup and keeps memory usage to two copies of the network at any one time (the live student + the frozen teacher snapshot).

Parameters:

Name	Type	Description	Default
model	NNModel	the :class:NNModel to train. Mutated in place across generations; the final state corresponds to the LAST generation. Restore from a checkpoint if you need an intermediate generation's weights.	required
generations	int	how many generations to run. generations=1 is a plain supervised run (no KD) — kept as a degenerate case so callers can sweep generations including the baseline. Must be ≥ 1.	3
train_params	NNTrainParams	passed unchanged to every :meth:NNModel.train call. The same run.id would be computed every generation if nothing else changed, but in practice each generation mutates the model's weights between calls, so the underlying artifacts (idps / phase checkpoints) diverge generation to generation even with the same id. Caveat: the shared id also means each generation's BEST tracking seeds from the previous generation's on-disk BEST — a generation that never beats its predecessor leaves BEST pointing at the EARLIER generation's weights. Run each generation from a fresh runs/ root (fresh cwd) if you need per-generation artifacts or independent BEST tracking.	required
**kd_kwargs	Any	forwarded to :func:kd_train_step_factory for generations ≥ 1 (alpha, temperature). Ignored on generation 0 (no teacher).	{}

Returns:

Type	Description
list[NNRun]	A list of :class:NNRun objects, one per generation, in order.
list[NNRun]	runs[0] is the plain run; runs[k] for k > 0 is the
list[NNRun]	KD run that used generation k-1's model as teacher.

Raises:

Type	Description
ValueError	if generations < 1.

Source code in nnx/paradigms/born_again.py

def born_again_train(
    model: NNModel,
    *,
    generations: int = 3,
    train_params: NNTrainParams,
    **kd_kwargs: Any,
) -> list[NNRun]:
    """Iterate G generations of self-distillation on a single model.

    Generation 0 trains plain (no teacher) — standard supervised loss.
    Each subsequent generation uses a deep-copied, frozen, eval-mode
    snapshot of the model *after* the prior generation completed as the
    teacher for a Hinton-style KD step (via :func:`kd_train_step_factory`).

    The same in-place ``model`` is reused across generations; only the
    teacher snapshot is duplicated. This matches the original paper's
    setup and keeps memory usage to two copies of the network at any
    one time (the live student + the frozen teacher snapshot).

    Args:
        model: the :class:`NNModel` to train. Mutated in place across
            generations; the final state corresponds to the LAST
            generation. Restore from a checkpoint if you need an
            intermediate generation's weights.
        generations: how many generations to run. ``generations=1`` is
            a plain supervised run (no KD) — kept as a degenerate case
            so callers can sweep generations including the baseline.
            Must be ``≥ 1``.
        train_params: passed unchanged to every :meth:`NNModel.train`
            call. The same ``run.id`` would be computed every generation
            if nothing else changed, but in practice each generation
            mutates the model's weights between calls, so the underlying
            artifacts (idps / phase checkpoints) diverge generation to
            generation even with the same id. Caveat: the shared id
            also means each generation's BEST tracking seeds from the
            previous generation's on-disk BEST — a generation that
            never beats its predecessor leaves BEST pointing at the
            EARLIER generation's weights. Run each generation from a
            fresh ``runs/`` root (fresh cwd) if you need per-generation
            artifacts or independent BEST tracking.
        **kd_kwargs: forwarded to :func:`kd_train_step_factory` for
            generations ≥ 1 (``alpha``, ``temperature``). Ignored on
            generation 0 (no teacher).

    Returns:
        A list of :class:`NNRun` objects, one per generation, in order.
        ``runs[0]`` is the plain run; ``runs[k]`` for ``k > 0`` is the
        KD run that used generation ``k-1``'s model as teacher.

    Raises:
        ValueError: if ``generations < 1``.
    """
    if generations < 1:
        raise ValueError(f"generations must be >= 1, got {generations}")

    runs: list[NNRun] = []
    teacher: NNModel | None = None
    for g in range(generations):
        if teacher is None:
            run = model.train(params=train_params)
        else:
            step_fn = kd_train_step_factory(teacher=teacher, **kd_kwargs)
            run = model.train(params=train_params, train_step_fn=step_fn)
        runs.append(run)

        if g == generations - 1:
            # The last generation has no successor — deep-copying a
            # never-used teacher would only double peak memory at exit.
            break

        # Snapshot the just-trained model as the next generation's
        # teacher. deepcopy duplicates the net's parameters into a
        # detached graph; freezing + eval-mode is then enforced here
        # (and again inside kd_train_step_factory as belt-and-braces).
        teacher = copy.deepcopy(model)
        teacher.net.eval()
        for p in teacher.net.parameters():
            p.requires_grad = False

    return runs

14.2. Contrastive

nnx.paradigms.contrastive.simclr_train_step_factory(*, temperature: float = 0.5) -> TrainStepFn

Build a SimCLR :class:TrainStepFn.

Parameters:

Name	Type	Description	Default
temperature	float	temperature in :func:nt_xent_loss. 0.5 default.	0.5

Returns:

Type	Description
TrainStepFn	A TrainStepFn for NNModel.train(..., train_step_fn=...).
TrainStepFn	The training loader MUST yield batches of two augmented views
TrainStepFn	per source sample — typically (view1, view2) tensors, or
TrainStepFn	((view1, view2), y_unused) when reusing a labelled dataset.
TrainStepFn	model.net is invoked once per view (no batch-doubling) so
TrainStepFn	BatchNorm statistics see one view at a time; users who want
TrainStepFn	all-at-once normalization can stack the views and forward once.
TrainStepFn	Sharp edge: a labeled (X, Y) batch from a standard
TrainStepFn	TensorDataset will silently be interpreted as
TrainStepFn	(view1=X, view2=Y) and produce a shape-mismatch in
TrainStepFn	func:nt_xent_loss. Use a paired-view dataset whose
TrainStepFn	__getitem__ returns (view1, view2) instead.

Raises:

Type	Description
ValueError	if temperature <= 0.

Source code in nnx/paradigms/contrastive.py

def simclr_train_step_factory(*, temperature: float = 0.5) -> TrainStepFn:
    """Build a SimCLR :class:`TrainStepFn`.

    Args:
        temperature: temperature in :func:`nt_xent_loss`. 0.5 default.

    Returns:
        A ``TrainStepFn`` for ``NNModel.train(..., train_step_fn=...)``.
        The training loader MUST yield batches of two augmented views
        per source sample — typically ``(view1, view2)`` tensors, or
        ``((view1, view2), y_unused)`` when reusing a labelled dataset.
        ``model.net`` is invoked once per view (no batch-doubling) so
        BatchNorm statistics see one view at a time; users who want
        all-at-once normalization can stack the views and forward once.

        **Sharp edge:** a labeled ``(X, Y)`` batch from a standard
        ``TensorDataset`` will silently be interpreted as
        ``(view1=X, view2=Y)`` and produce a shape-mismatch in
        :func:`nt_xent_loss`. Use a paired-view dataset whose
        ``__getitem__`` returns ``(view1, view2)`` instead.

    Raises:
        ValueError: if ``temperature`` <= 0.
    """
    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        x1, x2 = _unpack_views(ctx.batch)
        x1 = x1.to(m.device)
        x2 = x2.to(m.device)

        z1 = m.net(x1)
        z2 = m.net(x2)

        loss = nt_xent_loss(z1, z2, temperature=temperature)
        loss_val = finalize_step(loss, ctx, paradigm="simclr")

        # No classification metric for a contrastive paradigm — report
        # the loss in both slots so BEST tracking and ReduceLROnPlateau
        # have a single signal to lock onto.
        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=0.0,
            precision=0.0,
            loss=loss_val,
            error=loss_val,
        )

    return step

nnx.paradigms.contrastive.nt_xent_loss(z1: torch.Tensor, z2: torch.Tensor, *, temperature: float = 0.5) -> torch.Tensor

SimCLR's Normalized Temperature-scaled cross-entropy loss.

Parameters:

Name	Type	Description	Default
z1	Tensor	(B, D) embeddings of the first view of each sample.	required
z2	Tensor	(B, D) embeddings of the second view.	required
temperature	float	divisor on the cosine similarity. Lower T sharpens the distribution; 0.5 is the SimCLR default. Must be > 0.	0.5

Returns:

Type	Description
Tensor	Scalar loss tensor (mean across the 2B positions in the batch).

Raises:

Type	Description
ValueError	if shapes mismatch, temperature ≤ 0, or the batch has fewer than 2 pairs (no negatives to contrast).

Source code in nnx/paradigms/contrastive.py

def nt_xent_loss(
    z1: torch.Tensor,
    z2: torch.Tensor,
    *,
    temperature: float = 0.5,
) -> torch.Tensor:
    """SimCLR's Normalized Temperature-scaled cross-entropy loss.

    Args:
        z1: ``(B, D)`` embeddings of the first view of each sample.
        z2: ``(B, D)`` embeddings of the second view.
        temperature: divisor on the cosine similarity. Lower T sharpens
            the distribution; 0.5 is the SimCLR default. Must be > 0.

    Returns:
        Scalar loss tensor (mean across the 2B positions in the batch).

    Raises:
        ValueError: if shapes mismatch, ``temperature`` ≤ 0, or the
            batch has fewer than 2 pairs (no negatives to contrast).
    """
    if z1.shape != z2.shape:
        raise ValueError(f"z1 / z2 shape mismatch: {z1.shape} vs {z2.shape}")
    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")
    if z1.shape[0] < 2:
        # With a single pair there are no negatives: after the diagonal
        # mask each row has exactly one finite logit (its positive), so
        # the loss is identically 0.0 with zero gradient — a silent
        # no-op train step. Raise instead.
        raise ValueError(f"nt_xent_loss needs a batch of >= 2 pairs (got {z1.shape[0]}) — no negatives to contrast.")

    B = z1.shape[0]
    # L2-normalize so the matmul gives cosine similarity directly.
    z1 = F.normalize(z1, dim=-1)
    z2 = F.normalize(z2, dim=-1)

    # Stack into (2B, D) and form the (2B, 2B) similarity matrix.
    z = torch.cat([z1, z2], dim=0)
    sim = torch.matmul(z, z.t()) / temperature

    # Mask out self-similarity (the diagonal) — every row's logits
    # otherwise have a +1/T entry for the row's own embedding, which
    # would always be the most-similar position and crush the loss.
    self_mask = torch.eye(2 * B, dtype=torch.bool, device=sim.device)
    sim = sim.masked_fill(self_mask, float("-inf"))

    # Positives: row i in [0, B) has positive at column i+B; row j in
    # [B, 2B) has positive at column j-B. Cross-entropy of the
    # softmax over rejected self-positions, targeting those columns.
    targets = torch.cat([torch.arange(B, 2 * B, device=sim.device), torch.arange(0, B, device=sim.device)])
    return F.cross_entropy(sim, targets)

14.3. Augmentation

nnx.paradigms.augmentation.mixup_train_step_factory(*, alpha: float = 0.4) -> TrainStepFn

Build a Mixup :class:TrainStepFn.

Parameters:

Name	Type	Description	Default
alpha	float	Beta-distribution shape parameter. λ ~ Beta(α, α); α=1.0 yields a uniform mix, lower values concentrate λ near 0 or 1 (closer to no-mixing). 0.4 is the classification default; image-task papers often use 0.2-1.0. Must be positive.	0.4

Returns:

Type	Description
TrainStepFn	A TrainStepFn for NNModel.train(..., train_step_fn=...).
TrainStepFn	Reports a Mixup-weighted error and the mixed loss. The
TrainStepFn	loss honors the model's loss_fn (so this works for any
TrainStepFn	classification loss, not just CrossEntropy).

Raises:

Type	Description
ValueError	if alpha <= 0.

Source code in nnx/paradigms/augmentation.py

def mixup_train_step_factory(*, alpha: float = 0.4) -> TrainStepFn:
    """Build a Mixup :class:`TrainStepFn`.

    Args:
        alpha: Beta-distribution shape parameter. ``λ ~ Beta(α, α)``;
            α=1.0 yields a uniform mix, lower values concentrate λ
            near 0 or 1 (closer to no-mixing). 0.4 is the
            classification default; image-task papers often use 0.2-1.0.
            Must be positive.

    Returns:
        A ``TrainStepFn`` for ``NNModel.train(..., train_step_fn=...)``.
        Reports a Mixup-weighted ``error`` and the mixed loss. The
        loss honors the model's ``loss_fn`` (so this works for any
        classification loss, not just CrossEntropy).

    Raises:
        ValueError: if ``alpha`` <= 0.
    """
    if alpha <= 0:
        raise ValueError(f"alpha must be positive, got {alpha}")
    # Seed from the torch RNG so set_seed(...) controls the lambda draws
    # (and CutMix boxes): a bare default_rng() self-seeds from OS entropy
    # and ignores the run's reproducibility contract. Drawing the seed
    # through torch also keeps two factories' streams distinct.
    rng = np.random.default_rng(int(torch.randint(0, 2**31 - 1, (1,)).item()))

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        X, Y = _unpack_supervised(ctx)
        B = X.shape[0]

        lam = float(rng.beta(alpha, alpha))
        perm = torch.randperm(B, device=m.device)
        X_mixed = lam * X + (1.0 - lam) * X[perm]
        Y_a = Y
        Y_b = Y[perm]

        Y_hat_logits = m.net(X_mixed)
        loss = lam * m.loss_fn(Y_hat_logits, Y_a) + (1.0 - lam) * m.loss_fn(Y_hat_logits, Y_b)
        loss_val = finalize_step(loss, ctx, paradigm="mixup")

        Y_hat = Y_hat_logits.argmax(dim=-1)
        acc = _weighted_acc(Y_hat, Y_a, Y_b, lam)
        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=acc,
            precision=0.0,
            loss=loss_val,
            error=float(1.0 - acc),
        )

    return step

nnx.paradigms.augmentation.cutmix_train_step_factory(*, alpha: float = 1.0) -> TrainStepFn

Build a CutMix :class:TrainStepFn for 4D image batches.

Parameters:

Name	Type	Description	Default
alpha	float	Beta-distribution shape parameter for the area ratio. λ ~ Beta(α, α); controls the size of the swapped rectangle. 1.0 is the original paper default. Must be positive.	1.0

Returns:

Type	Description
TrainStepFn	A TrainStepFn for image classification (4D (B, C, H, W)
TrainStepFn	inputs). Raises at step time on lower-rank input — CutMix's
TrainStepFn	spatial cut isn't well-defined without H and W.

Raises:

Type	Description
ValueError	if alpha <= 0.

Source code in nnx/paradigms/augmentation.py

def cutmix_train_step_factory(*, alpha: float = 1.0) -> TrainStepFn:
    """Build a CutMix :class:`TrainStepFn` for 4D image batches.

    Args:
        alpha: Beta-distribution shape parameter for the area ratio.
            ``λ ~ Beta(α, α)``; controls the size of the swapped
            rectangle. 1.0 is the original paper default.
            Must be positive.

    Returns:
        A ``TrainStepFn`` for image classification (4D ``(B, C, H, W)``
        inputs). Raises at step time on lower-rank input — CutMix's
        spatial cut isn't well-defined without H and W.

    Raises:
        ValueError: if ``alpha`` <= 0.
    """
    if alpha <= 0:
        raise ValueError(f"alpha must be positive, got {alpha}")
    # Seed from the torch RNG so set_seed(...) controls the lambda draws
    # (and CutMix boxes): a bare default_rng() self-seeds from OS entropy
    # and ignores the run's reproducibility contract. Drawing the seed
    # through torch also keeps two factories' streams distinct.
    rng = np.random.default_rng(int(torch.randint(0, 2**31 - 1, (1,)).item()))

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        X, Y = _unpack_supervised(ctx)
        if X.dim() != 4:
            raise ValueError(
                f"CutMix requires 4D image input (B, C, H, W), got rank {X.dim()}. "
                "Use Mixup for tabular / lower-rank data."
            )
        B, _, H, W = X.shape

        lam = float(rng.beta(alpha, alpha))
        # Box dimensions: side lengths scale with √(1−λ) so the area
        # ratio is (1−λ). cx, cy are the box center, sampled uniformly.
        cut_ratio = (1.0 - lam) ** 0.5
        cut_w = int(W * cut_ratio)
        cut_h = int(H * cut_ratio)
        cx = int(rng.integers(0, W))
        cy = int(rng.integers(0, H))
        x1 = max(cx - cut_w // 2, 0)
        x2 = min(cx + cut_w // 2, W)
        y1 = max(cy - cut_h // 2, 0)
        y2 = min(cy + cut_h // 2, H)

        perm = torch.randperm(B, device=m.device)
        X_cut = X.clone()
        X_cut[:, :, y1:y2, x1:x2] = X[perm, :, y1:y2, x1:x2]
        # Re-derive λ from the actual cut area — clipping at the edges
        # makes the realized box smaller than the nominal Beta draw.
        lam = 1.0 - ((x2 - x1) * (y2 - y1) / (W * H))
        Y_a = Y
        Y_b = Y[perm]

        Y_hat_logits = m.net(X_cut)
        loss = lam * m.loss_fn(Y_hat_logits, Y_a) + (1.0 - lam) * m.loss_fn(Y_hat_logits, Y_b)
        loss_val = finalize_step(loss, ctx, paradigm="cutmix")

        Y_hat = Y_hat_logits.argmax(dim=-1)
        acc = _weighted_acc(Y_hat, Y_a, Y_b, lam)
        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=acc,
            precision=0.0,
            loss=loss_val,
            error=float(1.0 - acc),
        )

    return step

14.4. Mixture-of-Experts

MoELinear is the drop-in layer (documented in §4); moe_train_step_factory adds the Switch-style load-balancing aux loss to the supervised step.

nnx.paradigms.moe.moe_train_step_factory(*, aux_loss_weight: float = 0.01) -> TrainStepFn

Build an MoE-aware supervised :class:TrainStepFn.

The returned step performs the standard supervised forward (loss = m.loss_fn(net(X), Y)) and then adds the Switch-style load-balancing penalty summed across every :class:MoELinear layer in model.net, weighted by aux_loss_weight. Backward, grad-clip, and optimizer step go through :func:nnx._step_helpers.finalize_step for the same NaN-guard + grad-clip tail as the other paradigm factories.

Parameters:

Name	Type	Description	Default
aux_loss_weight	float	weight on the aux loss term (α in the Switch formulation). Must be non-negative. 0.0 turns the factory into a plain supervised step (the aux loss is still computed by each MoE forward but contributes 0 to backward). Defaults to 0.01 — the Switch paper's tutorial value; small enough not to dominate the main loss, large enough to prevent expert collapse.	0.01

Returns:

Type	Description
TrainStepFn	A TrainStepFn for :meth:NNModel.train. Works on any
TrainStepFn	single-input supervised net that contains ≥ 0
TrainStepFn	class:MoELinear layers; if there are no MoE layers, the
TrainStepFn	aux loss is 0 and the step is exactly supervised.

Raises:

Type	Description
ValueError	if aux_loss_weight < 0.

Source code in nnx/paradigms/moe.py

def moe_train_step_factory(*, aux_loss_weight: float = 0.01) -> TrainStepFn:
    """Build an MoE-aware supervised :class:`TrainStepFn`.

    The returned step performs the standard supervised forward
    (``loss = m.loss_fn(net(X), Y)``) and then *adds* the
    Switch-style load-balancing penalty summed across every
    :class:`MoELinear` layer in ``model.net``, weighted by
    ``aux_loss_weight``. Backward, grad-clip, and optimizer step go
    through :func:`nnx._step_helpers.finalize_step` for the same
    NaN-guard + grad-clip tail as the other paradigm factories.

    Args:
        aux_loss_weight: weight on the aux loss term (``α`` in the
            Switch formulation). Must be non-negative. ``0.0`` turns
            the factory into a plain supervised step (the aux loss is
            still computed by each MoE forward but contributes 0 to
            backward). Defaults to ``0.01`` — the Switch paper's
            tutorial value; small enough not to dominate the main
            loss, large enough to prevent expert collapse.

    Returns:
        A ``TrainStepFn`` for :meth:`NNModel.train`. Works on any
        single-input supervised net that contains ≥ 0
        :class:`MoELinear` layers; if there are no MoE layers, the
        aux loss is 0 and the step is exactly supervised.

    Raises:
        ValueError: if ``aux_loss_weight < 0``.
    """
    if aux_loss_weight < 0:
        raise ValueError(f"aux_loss_weight must be non-negative, got {aux_loss_weight}")

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        # Single-input destructuring — same contract as the other
        # supervised paradigm factories (Mixup / CutMix / KD).
        # Multi-input nets need a custom step.
        (X,), Y = m.net.unpack_batch(ctx.batch)
        X = X.to(m.device)
        Y = Y.to(m.device)

        # Clear stale aux losses BEFORE the forward: a MoELinear that's
        # registered but not exercised by this batch (conditional branch,
        # frozen tower) would otherwise contribute the previous step's
        # tensor — whose graph the previous backward already freed —
        # crashing with "backward through the graph a second time".
        for module in m.net.modules():
            if isinstance(module, MoELinear):
                module.last_aux_loss = None

        # The supervised forward populates each MoELinear's
        # ``.last_aux_loss`` as a side effect.
        Y_hat_logits = m.net(X)
        supervised_loss = m.loss_fn(Y_hat_logits, Y)

        # Sum the per-layer aux losses across every MoE layer in the
        # net. ``.last_aux_loss`` is set by every MoELinear forward;
        # we collect them post-forward. If there are no MoELinear
        # layers, the sum is the scalar 0 — the factory degenerates
        # to a plain supervised step. We construct the zero on the
        # same device as the supervised loss so the add below stays
        # device-coherent regardless of model placement.
        aux_loss: torch.Tensor = torch.zeros((), device=supervised_loss.device)
        for module in m.net.modules():
            if isinstance(module, MoELinear) and module.last_aux_loss is not None:
                aux_loss = aux_loss + module.last_aux_loss

        loss = supervised_loss + aux_loss_weight * aux_loss
        loss_val = finalize_step(loss, ctx, paradigm="moe")

        return classification_edp(
            Y=Y,
            Y_hat=Y_hat_logits.argmax(dim=-1),
            loss=loss_val,
            extra_metrics=ctx.extra_metrics,
        )

    return step

14.5. I-JEPA

Walkthrough at I-JEPA. The ViTNN encoder is documented in §4.

nnx.paradigms.jepa.jepa_train_step_factory(target_encoder: nn.Module, predictor: nn.Module, mask_fn: Callable[[int, torch.device], tuple[torch.Tensor, torch.Tensor]], *, ema_momentum: float = 0.996) -> TrainStepFn

Build an I-JEPA :class:TrainStepFn.

Per step:

1. Sample (context_mask, target_mask) for the batch via mask_fn(n_patches, device). Both are 1-D BoolTensor[n_patches] and complementary — every patch is either context or target.
2. Forward each input image through model.net with the context mask, producing (B, T_ctx + 1, d_model) context embeddings (CLS at index 0).
3. Forward the full image (no mask) through target_encoder under no_grad to produce target embeddings. Slice out the positions in target_mask only.
4. Predict (B, T_tgt, d_model) from context via predictor.
5. MSE loss against the target embeddings.
6. :func:finalize_step — NaN guard, optimizer step, grad clip.
7. :func:update_ema — EMA-update the target encoder from model.net.

Parameters:

Name	Type	Description	Default
target_encoder	Module	an EMA copy of model.net. Build via :func:build_target_encoder. The factory freezes it again on call and pins to eval() mode.	required
predictor	Module	a :class:JEPAPredictor (or any module with the same forward(context_embeds, context_positions, target_positions) contract). The predictor's parameters are not frozen — the optimizer's param_groups need to include them; the simplest path is to register the predictor as a submodule of model.net (the ViTNN) before constructing the optimizer.	required
mask_fn	Callable[[int, device], tuple[Tensor, Tensor]]	callable (n_patches, device) -> (context_mask, target_mask) where both are 1-D BoolTensor[n_patches]. Sampled freshly once per step and shared across the batch. The bundled :func:random_block_mask helper is the common choice; passing a fixed mask is fine for tests.	required
ema_momentum	float	EMA decay used by :func:update_ema. Default 0.996 (reference I-JEPA).	0.996

Returns:

Type	Description
TrainStepFn	A TrainStepFn for NNModel.train(..., train_step_fn=...).

Raises:

Type	Description
ValueError	when ema_momentum is outside [0, 1).

Source code in nnx/paradigms/jepa.py

def jepa_train_step_factory(
    target_encoder: nn.Module,
    predictor: nn.Module,
    mask_fn: Callable[[int, torch.device], tuple[torch.Tensor, torch.Tensor]],
    *,
    ema_momentum: float = 0.996,
) -> TrainStepFn:
    """Build an I-JEPA :class:`TrainStepFn`.

    Per step:

      1. Sample ``(context_mask, target_mask)`` for the batch via
         ``mask_fn(n_patches, device)``. Both are 1-D
         ``BoolTensor[n_patches]`` and **complementary** — every
         patch is either context or target.
      2. Forward each input image through ``model.net`` with the
         context mask, producing ``(B, T_ctx + 1, d_model)`` context
         embeddings (CLS at index 0).
      3. Forward the full image (no mask) through ``target_encoder``
         under ``no_grad`` to produce target embeddings. Slice out
         the positions in ``target_mask`` only.
      4. Predict ``(B, T_tgt, d_model)`` from context via
         ``predictor``.
      5. MSE loss against the target embeddings.
      6. :func:`finalize_step` — NaN guard, optimizer step, grad clip.
      7. :func:`update_ema` — EMA-update the target encoder from
         ``model.net``.

    Args:
        target_encoder: an EMA copy of ``model.net``. Build via
            :func:`build_target_encoder`. The factory **freezes** it
            again on call and pins to ``eval()`` mode.
        predictor: a :class:`JEPAPredictor` (or any module with the
            same ``forward(context_embeds, context_positions,
            target_positions)`` contract). The predictor's parameters
            are *not* frozen — the optimizer's ``param_groups`` need
            to include them; the simplest path is to register the
            predictor as a submodule of ``model.net`` (the ViTNN)
            before constructing the optimizer.
        mask_fn: callable ``(n_patches, device) -> (context_mask,
            target_mask)`` where both are 1-D ``BoolTensor[n_patches]``.
            Sampled freshly **once per step** and shared across the
            batch. The bundled :func:`random_block_mask` helper is the
            common choice; passing a fixed mask is fine for tests.
        ema_momentum: EMA decay used by :func:`update_ema`. Default
            0.996 (reference I-JEPA).

    Returns:
        A ``TrainStepFn`` for ``NNModel.train(..., train_step_fn=...)``.

    Raises:
        ValueError: when ``ema_momentum`` is outside ``[0, 1)``.
    """
    if not (0.0 <= ema_momentum < 1.0):
        raise ValueError(f"ema_momentum must be in [0, 1), got {ema_momentum}")

    # Defensive freeze in case the caller built the target encoder by
    # hand without going through ``build_target_encoder``.
    target_encoder.eval()
    for p in target_encoder.parameters():
        p.requires_grad = False

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        predictor.train()
        m.net.zero_grad()
        # The predictor may live outside `model.net` (its params added to the
        # optimizer directly, per the docstring's alternate path) — then
        # `m.net.zero_grad()` above does not reach it. Zero it explicitly so
        # its gradients don't accumulate across steps. No-op when it is a
        # submodule of `model.net` (the recommended path).
        predictor.zero_grad()

        # Standard dataloader contract: (X, Y) or single tensor. Y is
        # ignored — JEPA is self-supervised.
        if hasattr(m.net, "unpack_batch"):
            (x,), _ = m.net.unpack_batch(ctx.batch)
        elif isinstance(ctx.batch, (list, tuple)):
            x = ctx.batch[0]
        else:
            x = ctx.batch
        x = x.to(m.device)

        # Sample the per-step mask. Both masks are 1-D length n_patches.
        n_patches = m.net.n_patches
        context_mask_1d, target_mask_1d = mask_fn(n_patches, m.device)
        if context_mask_1d.shape != (n_patches,) or target_mask_1d.shape != (n_patches,):
            raise ValueError(
                f"mask_fn must return two BoolTensors of shape ({n_patches},); "
                f"got {tuple(context_mask_1d.shape)} and {tuple(target_mask_1d.shape)}"
            )
        if not torch.equal(context_mask_1d, ~target_mask_1d):
            raise ValueError(
                "context_mask and target_mask must be complementary (every patch is either context or target)."
            )
        B = x.shape[0]
        context_mask = _broadcast_mask(context_mask_1d, B)

        # Forward context through trainable encoder.
        context_embeds = m.net(x, mask=context_mask)  # (B, T_ctx + 1, d_model)

        # Build position indices the predictor needs. CLS is position 0;
        # ViTNN.patch_positions() owns the 1..n_patches CLS shift, so
        # boolean-masking it yields the kept/target position indices.
        patch_positions = m.net.patch_positions()
        kept_patch_positions = patch_positions[context_mask_1d]
        context_positions = torch.cat([torch.zeros(1, dtype=torch.long, device=m.device), kept_patch_positions])
        target_positions = patch_positions[target_mask_1d]

        # Forward FULL image through EMA target (no grad), then slice
        # out the target-position embeddings only. Skip CLS at index 0.
        with torch.no_grad():
            target_full = target_encoder(x)  # (B, n_patches + 1, d_model)
            # target_positions are in the [1, n_patches] range.
            target_embeds = target_full[:, target_positions, :]

        # Predict target embeddings from context.
        predicted_embeds = predictor(context_embeds, context_positions, target_positions)

        loss = F.mse_loss(predicted_embeds, target_embeds)
        loss_val = finalize_step(loss, ctx, paradigm="jepa")
        # EMA-update target encoder AFTER the optimizer has stepped
        # — finalize_step ran the optimizer, so model.net now reflects
        # the post-step weights and the EMA tracks the freshest signal.
        update_ema(m.net, target_encoder, ema_momentum)

        # No classification metric for a self-supervised paradigm; report
        # the loss in both slots so BEST tracking + ReduceLROnPlateau
        # have a signal.
        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=0.0,
            precision=0.0,
            loss=loss_val,
            error=loss_val,
        )

    return step

nnx.paradigms.jepa.JEPAPredictor

Bases: Module

Tiny ViT-like predictor: (context_embeds, target_positions) -> predicted_target_embeds.

Architecture: project context_embeds to predictor_dim, concatenate learnable mask tokens (one per target position) plus that position's positional embedding, run a few ViT blocks, project back to embed_dim, return the predictions at the target positions only.

Kept deliberately small — the reference I-JEPA predictor is also much narrower than the encoder. For our CIFAR-shape demo, two blocks at predictor_dim = embed_dim // 2 is enough plumbing to verify the loss decreases without dominating wall-clock time.

Source code in nnx/paradigms/jepa.py

class JEPAPredictor(nn.Module):
    """Tiny ViT-like predictor: ``(context_embeds, target_positions)
    -> predicted_target_embeds``.

    Architecture: project context_embeds to ``predictor_dim``,
    concatenate learnable mask tokens (one per target position) plus
    that position's positional embedding, run a few ViT blocks, project
    back to ``embed_dim``, return the predictions at the target
    positions only.

    Kept deliberately small — the reference I-JEPA predictor is also
    much narrower than the encoder. For our CIFAR-shape demo, two
    blocks at ``predictor_dim = embed_dim // 2`` is enough plumbing
    to verify the loss decreases without dominating wall-clock time.
    """

    def __init__(
        self,
        *,
        embed_dim: int,
        n_patches: int,
        predictor_dim: Optional[int] = None,
        n_layers: int = 2,
        n_heads: int = 2,
        ffn_mult: int = 4,
    ):
        super().__init__()
        # Local import keeps the paradigm module independent of the
        # vit_nn module ordering (the only reason ViTBlock lives where
        # it does is its use of RMSNorm + SwiGLU from the transformer
        # building blocks).
        from ..nn.net.vit_nn import ViTBlock

        # Positive-dimension guards: JEPAPredictor is a public top-level export
        # (nnx.JEPAPredictor) that builds ViTBlocks from raw numeric kwargs, so
        # n_layers=0 / ffn_mult=0 / embed_dim=0 would otherwise build a silently
        # degenerate predictor — the same [[params-boundary-validation]] footgun.
        # Validate predictor_dim before the None-resolution below.
        for _name, _value in (
            ("embed_dim", embed_dim),
            ("n_patches", n_patches),
            ("n_layers", n_layers),
            ("n_heads", n_heads),
            ("ffn_mult", ffn_mult),
        ):
            if _value <= 0:
                raise ValueError(f"JEPAPredictor requires {_name} > 0, got {_value}")
        if predictor_dim is not None and predictor_dim <= 0:
            raise ValueError(f"JEPAPredictor requires predictor_dim > 0 when set, got {predictor_dim}")

        if predictor_dim is None:
            predictor_dim = max(8, embed_dim // 2)
        self.embed_dim = embed_dim
        self.predictor_dim = predictor_dim
        self.n_patches = n_patches

        # In/out projections so the predictor can run at a smaller
        # internal width than the encoder.
        self.in_proj = nn.Linear(embed_dim, predictor_dim, bias=False)
        self.out_proj = nn.Linear(predictor_dim, embed_dim, bias=False)
        # Per-patch positional embedding shared with the encoder *in
        # spirit* but learned independently — the predictor doesn't
        # have access to the encoder's tied weights.
        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, predictor_dim))
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        # Learnable mask token — same value at every target position.
        # The predictor disambiguates "which patch are we predicting?"
        # via the additive positional embedding.
        self.mask_token = nn.Parameter(torch.zeros(1, 1, predictor_dim))
        nn.init.trunc_normal_(self.mask_token, std=0.02)

        self.blocks = nn.ModuleList(
            [ViTBlock(d_model=predictor_dim, n_heads=n_heads, ffn_mult=ffn_mult) for _ in range(n_layers)]
        )

    def forward(
        self,
        context_embeds: torch.Tensor,
        context_positions: torch.Tensor,
        target_positions: torch.Tensor,
    ) -> torch.Tensor:
        """Predict embeddings at ``target_positions`` from
        ``context_embeds``.

        Args:
            context_embeds: ``(B, T_ctx, embed_dim)``. The CLS token
                produced by the encoder is included as the first entry
                (position 0).
            context_positions: ``LongTensor[T_ctx]`` — positions of
                the kept context tokens *including* CLS at index 0.
            target_positions: ``LongTensor[T_tgt]`` — positions of the
                target patches to predict (1..n_patches).

        Returns:
            ``(B, T_tgt, embed_dim)`` predicted target embeddings.
        """
        B = context_embeds.shape[0]
        x_ctx = self.in_proj(context_embeds)  # (B, T_ctx, predictor_dim)
        # Add the predictor's own positional embedding to context tokens.
        x_ctx = x_ctx + self.pos_embed[:, context_positions, :]
        # Build mask tokens at target positions.
        T_tgt = target_positions.shape[0]
        x_tgt = self.mask_token.expand(B, T_tgt, -1) + self.pos_embed[:, target_positions, :]
        # Concatenate and pass through the small transformer.
        x = torch.cat([x_ctx, x_tgt], dim=1)
        for block in self.blocks:
            x = block(x)
        # Slice out only the target-position predictions and project back.
        T_ctx = x_ctx.shape[1]
        x_tgt_out = x[:, T_ctx:, :]
        return self.out_proj(x_tgt_out)

forward(context_embeds: torch.Tensor, context_positions: torch.Tensor, target_positions: torch.Tensor) -> torch.Tensor

Predict embeddings at target_positions from context_embeds.

Parameters:

Name	Type	Description	Default
context_embeds	Tensor	(B, T_ctx, embed_dim). The CLS token produced by the encoder is included as the first entry (position 0).	required
context_positions	Tensor	LongTensor[T_ctx] — positions of the kept context tokens including CLS at index 0.	required
target_positions	Tensor	LongTensor[T_tgt] — positions of the target patches to predict (1..n_patches).	required

Returns:

Type	Description
Tensor	(B, T_tgt, embed_dim) predicted target embeddings.

Source code in nnx/paradigms/jepa.py

def forward(
    self,
    context_embeds: torch.Tensor,
    context_positions: torch.Tensor,
    target_positions: torch.Tensor,
) -> torch.Tensor:
    """Predict embeddings at ``target_positions`` from
    ``context_embeds``.

    Args:
        context_embeds: ``(B, T_ctx, embed_dim)``. The CLS token
            produced by the encoder is included as the first entry
            (position 0).
        context_positions: ``LongTensor[T_ctx]`` — positions of
            the kept context tokens *including* CLS at index 0.
        target_positions: ``LongTensor[T_tgt]`` — positions of the
            target patches to predict (1..n_patches).

    Returns:
        ``(B, T_tgt, embed_dim)`` predicted target embeddings.
    """
    B = context_embeds.shape[0]
    x_ctx = self.in_proj(context_embeds)  # (B, T_ctx, predictor_dim)
    # Add the predictor's own positional embedding to context tokens.
    x_ctx = x_ctx + self.pos_embed[:, context_positions, :]
    # Build mask tokens at target positions.
    T_tgt = target_positions.shape[0]
    x_tgt = self.mask_token.expand(B, T_tgt, -1) + self.pos_embed[:, target_positions, :]
    # Concatenate and pass through the small transformer.
    x = torch.cat([x_ctx, x_tgt], dim=1)
    for block in self.blocks:
        x = block(x)
    # Slice out only the target-position predictions and project back.
    T_ctx = x_ctx.shape[1]
    x_tgt_out = x[:, T_ctx:, :]
    return self.out_proj(x_tgt_out)

nnx.paradigms.jepa.build_target_encoder(source: nn.Module) -> nn.Module

Deep-copy source, freeze every parameter, return the copy.

The target encoder is updated only via :func:update_ema after each optimizer step. Freezing here is belt-and-braces — even if a user accidentally hands the target into an optimizer that scans parameters(), requires_grad=False keeps the gradients off and the optimizer's state empty for those tensors.

Source code in nnx/paradigms/jepa.py

def build_target_encoder(source: nn.Module) -> nn.Module:
    """Deep-copy ``source``, freeze every parameter, return the copy.

    The target encoder is updated **only** via :func:`update_ema` after
    each optimizer step. Freezing here is belt-and-braces — even if a
    user accidentally hands the target into an optimizer that scans
    ``parameters()``, ``requires_grad=False`` keeps the gradients off
    and the optimizer's state empty for those tensors.
    """
    target = copy.deepcopy(source)
    for p in target.parameters():
        p.requires_grad = False
    target.eval()
    return target

nnx.paradigms.jepa.update_ema(source: nn.Module, target: nn.Module, momentum: float) -> None

In-place EMA update: target ← momentum * target + (1 - momentum) * source.

Called once per training step from inside the JEPA train_step_fn. Runs under torch.no_grad so the EMA tensors do not become part of the autograd graph — the target encoder is supposed to be a detached snapshot.

Parameters:

Name	Type	Description	Default
source	Module	the trainable module (i.e., model.net).	required
target	Module	the EMA copy returned by :func:build_target_encoder. Mutated in place.	required
momentum	float	EMA decay in [0, 1). Higher = slower target tracking. I-JEPA's reference recipe uses 0.996 with a cosine schedule up to 1.0 over training; the factory's default matches.	required

Raises:

Type	Description
ValueError	when momentum is outside [0, 1).
KeyError	when a target parameter has no same-named source parameter (the name-keyed update contract).

Source code in nnx/paradigms/jepa.py

def update_ema(source: nn.Module, target: nn.Module, momentum: float) -> None:
    """In-place EMA update: ``target ← momentum * target + (1 - momentum) * source``.

    Called once per training step from inside the JEPA train_step_fn.
    Runs under ``torch.no_grad`` so the EMA tensors do not become part
    of the autograd graph — the target encoder is supposed to be a
    detached snapshot.

    Args:
        source: the trainable module (i.e., ``model.net``).
        target: the EMA copy returned by :func:`build_target_encoder`.
            Mutated in place.
        momentum: EMA decay in ``[0, 1)``. Higher = slower target
            tracking. I-JEPA's reference recipe uses 0.996 with a
            cosine schedule up to 1.0 over training; the factory's
            default matches.

    Raises:
        ValueError: when ``momentum`` is outside ``[0, 1)``.
        KeyError: when a target parameter has no same-named source
            parameter (the name-keyed update contract).
    """
    if not (0.0 <= momentum < 1.0):
        raise ValueError(f"EMA momentum must be in [0, 1), got {momentum}")
    # Walk by named-parameter so the EMA continues to work when the
    # *source* has been augmented with extra submodules (typical
    # idiom: register the JEPA predictor under ``model.net`` so a
    # single optimizer picks it up). The target encoder is a frozen
    # snapshot of the original encoder; any source param whose name
    # is absent from the target is something the EMA was never
    # responsible for and is silently skipped.
    src_by_name = dict(source.named_parameters())
    with torch.no_grad():
        for name, pt in target.named_parameters():
            ps = src_by_name.get(name)
            if ps is None:
                raise KeyError(
                    f"EMA target param {name!r} has no counterpart in the source; "
                    "did you swap the source network's structure after building "
                    "the target encoder?"
                )
            pt.mul_(momentum).add_(ps, alpha=1.0 - momentum)

nnx.paradigms.jepa.random_block_mask(*, n_patches: int, grid_size: int, block_scale: tuple[float, float] = (0.15, 0.2), block_aspect: tuple[float, float] = (0.75, 1.5), generator: Optional[torch.Generator] = None, device: Optional[torch.device] = None) -> tuple[torch.Tensor, torch.Tensor]

Sample one I-JEPA-style rectangular block mask on a patch grid.

Returns (context_mask, target_mask) where:

• context_mask: BoolTensor[n_patches] — True at positions kept by the context encoder (i.e., NOT in the target block).
• target_mask: BoolTensor[n_patches] — True at positions the predictor is asked to predict (i.e., inside the target block, exactly the complement of context_mask).

The block is a single rectangle of randomly-sampled width/height drawn from block_scale × n_patches with an aspect ratio in block_aspect. Reference I-JEPA samples 4 target blocks per image; this helper samples 1 — enough for the verify-the-plumbing example we ship. Users can compose multiple calls if they want the 4-block recipe.

Parameters:

Name	Type	Description	Default
n_patches	int	total number of patch tokens. Must equal grid_size**2.	required
grid_size	int	width (= height) of the patch grid. The rectangular block is sampled in this coordinate system.	required
block_scale	tuple[float, float]	(min, max) fraction of n_patches covered by the block. Default (0.15, 0.2) mirrors I-JEPA.	(0.15, 0.2)
block_aspect	tuple[float, float]	(min, max) width/height ratio.	(0.75, 1.5)
generator	Optional[Generator]	optional torch.Generator for reproducibility.	None
device	Optional[device]	device on which the masks are placed. None → default tensor device (CPU).	None

Returns:

Type	Description
tuple[Tensor, Tensor]	A pair of BoolTensors, both 1-D length n_patches.

Raises:

Type	Description
ValueError	when grid_size**2 != n_patches, or when the sampled block would be empty / larger than the grid.

Source code in nnx/paradigms/jepa.py

def random_block_mask(
    *,
    n_patches: int,
    grid_size: int,
    block_scale: tuple[float, float] = (0.15, 0.2),
    block_aspect: tuple[float, float] = (0.75, 1.5),
    generator: Optional[torch.Generator] = None,
    device: Optional[torch.device] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Sample one I-JEPA-style rectangular block mask on a patch grid.

    Returns ``(context_mask, target_mask)`` where:

      * ``context_mask: BoolTensor[n_patches]`` — True at positions
        kept by the context encoder (i.e., NOT in the target block).
      * ``target_mask: BoolTensor[n_patches]`` — True at positions
        the predictor is asked to predict (i.e., inside the target
        block, exactly the complement of context_mask).

    The block is a single rectangle of randomly-sampled width/height
    drawn from ``block_scale`` × n_patches with an aspect ratio in
    ``block_aspect``. Reference I-JEPA samples 4 target blocks per
    image; this helper samples 1 — enough for the verify-the-plumbing
    example we ship. Users can compose multiple calls if they want
    the 4-block recipe.

    Args:
        n_patches: total number of patch tokens. Must equal
            ``grid_size**2``.
        grid_size: width (= height) of the patch grid. The
            rectangular block is sampled in this coordinate system.
        block_scale: ``(min, max)`` fraction of ``n_patches`` covered
            by the block. Default ``(0.15, 0.2)`` mirrors I-JEPA.
        block_aspect: ``(min, max)`` width/height ratio.
        generator: optional ``torch.Generator`` for reproducibility.
        device: device on which the masks are placed. ``None`` →
            default tensor device (CPU).

    Returns:
        A pair of ``BoolTensor``s, both 1-D length ``n_patches``.

    Raises:
        ValueError: when ``grid_size**2 != n_patches``, or when the
            sampled block would be empty / larger than the grid.
    """
    if grid_size * grid_size != n_patches:
        raise ValueError(f"grid_size**2 ({grid_size**2}) != n_patches ({n_patches})")

    # Sample block scale and aspect. ``torch.rand`` accepts a
    # generator directly; ``Tensor.uniform_`` does not, which is why
    # we go through rand here instead of uniform_.
    u_scale = torch.rand(1, generator=generator).item()
    u_aspect = torch.rand(1, generator=generator).item()
    scale = block_scale[0] + (block_scale[1] - block_scale[0]) * u_scale
    aspect = block_aspect[0] + (block_aspect[1] - block_aspect[0]) * u_aspect
    block_area = scale * n_patches
    h = max(1, int(round(math.sqrt(block_area / aspect))))
    w = max(1, int(round(math.sqrt(block_area * aspect))))
    h = min(h, grid_size)
    w = min(w, grid_size)
    # Sample top-left corner. The block must fit inside the grid.
    top = int(torch.randint(0, grid_size - h + 1, (1,), generator=generator).item())
    left = int(torch.randint(0, grid_size - w + 1, (1,), generator=generator).item())

    grid = torch.zeros(grid_size, grid_size, dtype=torch.bool, device=device)
    grid[top : top + h, left : left + w] = True
    target_mask = grid.flatten()  # (n_patches,)
    context_mask = ~target_mask
    return context_mask, target_mask

14.6. DPO

Walkthrough at DPO.

nnx.paradigms.dpo.dpo_train_step_factory(ref_model: NNModel, *, beta: float = 0.1, pad_token_id: Optional[int] = None) -> TrainStepFn

Build a Direct Preference Optimization :class:TrainStepFn.

Parameters:

Name	Type	Description	Default
ref_model	NNModel	a frozen reference policy — typically a copy of the SFT checkpoint that the trainable policy was initialized from. Its net is set to eval mode and every parameter has requires_grad cleared on factory call. Must share vocab_size and tokenization with the policy.	required
beta	float	temperature on the implicit reward. Larger beta makes the loss sharper (closer to a hard preference); smaller beta keeps the policy closer to the reference. The original DPO paper uses 0.1 as the default; values in [0.01, 0.5] are common. Must be > 0.	0.1
pad_token_id	Optional[int]	the id the dataset used to right-pad chosen / rejected responses (NNPreferenceDataset.pad_token_id). When set, padded positions are excluded from the response log-prob sums. Without it, pad tokens are scored too — the pad terms don't cancel between policy/reference or chosen/rejected (different contexts), biasing the objective and training the policy to emit pads after short responses. None is only appropriate when every response genuinely fills max_response_len. Two caveats: masking is by token-id equality, so a genuine occurrence of the pad id inside a response is dropped too (pick a dedicated pad id); and prompt-side padding remains visible to the model (no attention mask) — a pre-existing modeling bias this knob doesn't address.	None

Returns:

Type	Description
TrainStepFn	A TrainStepFn for NNModel.train(..., train_step_fn=...).
TrainStepFn	The training loader MUST yield batches of three
TrainStepFn	torch.LongTensor of shape (B, T_*):: (prompt_ids, chosen_ids, rejected_ids)
TrainStepFn	— typically from :class:nnx.NNPreferenceDataset. All three
TrainStepFn	tensors must already be padded / right-aligned by the dataset.

Raises:

Type	Description
ValueError	if beta ≤ 0.

Source code in nnx/paradigms/dpo.py

def dpo_train_step_factory(
    ref_model: NNModel,
    *,
    beta: float = 0.1,
    pad_token_id: Optional[int] = None,
) -> TrainStepFn:
    """Build a Direct Preference Optimization :class:`TrainStepFn`.

    Args:
        ref_model: a frozen reference policy — typically a copy of the
            SFT checkpoint that the trainable policy was initialized
            from. Its ``net`` is set to eval mode and every parameter
            has ``requires_grad`` cleared on factory call. Must share
            ``vocab_size`` and tokenization with the policy.
        beta: temperature on the implicit reward. Larger ``beta`` makes
            the loss sharper (closer to a hard preference); smaller
            ``beta`` keeps the policy closer to the reference. The
            original DPO paper uses 0.1 as the default; values in
            ``[0.01, 0.5]`` are common. Must be > 0.
        pad_token_id: the id the dataset used to right-pad chosen /
            rejected responses (``NNPreferenceDataset.pad_token_id``).
            When set, padded positions are excluded from the response
            log-prob sums. Without it, pad tokens are scored too — the
            pad terms don't cancel between policy/reference or
            chosen/rejected (different contexts), biasing the objective
            and training the policy to emit pads after short responses.
            ``None`` is only appropriate when every response genuinely
            fills ``max_response_len``. Two caveats: masking is by
            token-id equality, so a genuine occurrence of the pad id
            inside a response is dropped too (pick a dedicated pad id);
            and prompt-side padding remains visible to the model (no
            attention mask) — a pre-existing modeling bias this knob
            doesn't address.

    Returns:
        A ``TrainStepFn`` for ``NNModel.train(..., train_step_fn=...)``.
        The training loader MUST yield batches of three
        ``torch.LongTensor`` of shape ``(B, T_*)``::

            (prompt_ids, chosen_ids, rejected_ids)

        — typically from :class:`nnx.NNPreferenceDataset`. All three
        tensors must already be padded / right-aligned by the dataset.

    Raises:
        ValueError: if ``beta`` ≤ 0.
    """
    if beta <= 0:
        raise ValueError(f"beta must be positive, got {beta}")

    # Freeze the reference and pin to eval mode. The policy's training
    # never touches the reference; this just guards against accidental
    # gradient flow if the caller wires them into a shared module later.
    ref_model.net.eval()
    for p in ref_model.net.parameters():
        p.requires_grad = False

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        prompt_ids, chosen_ids, rejected_ids = _unpack_preference_batch(ctx.batch)
        prompt_ids = prompt_ids.to(m.device)
        chosen_ids = chosen_ids.to(m.device)
        rejected_ids = rejected_ids.to(m.device)

        # Build the (prompt + response) sequences. We compute the
        # response log-prob conditional on the prompt by summing the
        # per-token log-probs over the response positions only.
        chosen_seq = torch.cat([prompt_ids, chosen_ids], dim=1)
        rejected_seq = torch.cat([prompt_ids, rejected_ids], dim=1)
        prompt_len = prompt_ids.shape[1]

        # Policy log-probs (with gradient).
        policy_chosen_logp = _response_logprob(m.net, chosen_seq, prompt_len, pad_token_id=pad_token_id)
        policy_rejected_logp = _response_logprob(m.net, rejected_seq, prompt_len, pad_token_id=pad_token_id)

        # Reference log-probs — no gradient, but move tensors through
        # ref_model's device frame in case it lives elsewhere.
        with torch.no_grad():
            ref_chosen_logp = _response_logprob(
                ref_model.net,
                chosen_seq.to(ref_model.device),
                prompt_len,
                pad_token_id=pad_token_id,
            ).to(m.device)
            ref_rejected_logp = _response_logprob(
                ref_model.net,
                rejected_seq.to(ref_model.device),
                prompt_len,
                pad_token_id=pad_token_id,
            ).to(m.device)

        # DPO loss: −log σ(β · ((logπ_w − logπ_l) − (logπref_w − logπref_l))).
        # Equivalent to −log σ(β · (Δ_w − Δ_l)) with Δ_* the
        # policy-minus-reference margins.
        pi_logratios = policy_chosen_logp - policy_rejected_logp
        ref_logratios = ref_chosen_logp - ref_rejected_logp
        loss = -F.logsigmoid(beta * (pi_logratios - ref_logratios)).mean()

        loss_val = finalize_step(loss, ctx, paradigm="dpo")

        # No classification target — surface the loss in both slots so
        # BEST tracking + ReduceLROnPlateau have a single signal.
        # Track the chosen-minus-rejected log-prob gap as a side metric
        # via `error`: more negative = bigger preference margin learned.
        gap = float((policy_chosen_logp - policy_rejected_logp).detach().mean())
        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=0.0,
            precision=0.0,
            loss=loss_val,
            error=-gap,
        )

    return step

15. Embeddings (nnx.embeddings)

End-to-end walkthrough at Embeddings. Opt-in via pip install "thekaveh-nnx[embeddings]".

nnx.embeddings.contrastive_trainer.ContrastiveTextDataset

Bases: Dataset

Wraps (anchor, positive) string pairs as a torch Dataset.

Each __getitem__ returns a 2-tuple of strings (anchor, positive). The default collate from :class:torch.utils.data.DataLoader would attempt to stack these into tensors and crash; pair this dataset with :func:pair_collate (or pass it directly to :func:train_contrastive which wires the collate for you).

Parameters:

Name	Type	Description	Default
pairs	list[tuple[str, str]]	list of (anchor, positive) string tuples. Empty input raises :class:ValueError. Note that :func:train_contrastive additionally requires >= 2 pairs (NT-Xent needs a negative); a 1-pair dataset is accepted here only for embedding/inference-style uses.	required

Raises:

Type	Description
ValueError	if pairs is empty or any entry isn't a 2-tuple of strings.

Source code in nnx/embeddings/contrastive_trainer.py

class ContrastiveTextDataset(Dataset):
    """Wraps ``(anchor, positive)`` string pairs as a torch ``Dataset``.

    Each ``__getitem__`` returns a 2-tuple of strings (``anchor``,
    ``positive``). The default collate from :class:`torch.utils.data.DataLoader`
    would attempt to stack these into tensors and crash; pair this
    dataset with :func:`pair_collate` (or pass it directly to
    :func:`train_contrastive` which wires the collate for you).

    Args:
        pairs: list of ``(anchor, positive)`` string tuples. Empty
            input raises :class:`ValueError`. Note that
            :func:`train_contrastive` additionally requires >= 2 pairs
            (NT-Xent needs a negative); a 1-pair dataset is accepted
            here only for embedding/inference-style uses.

    Raises:
        ValueError: if ``pairs`` is empty or any entry isn't a 2-tuple
            of strings.
    """

    def __init__(self, pairs: list[tuple[str, str]]):
        if not pairs:
            raise ValueError("ContrastiveTextDataset requires at least one pair")
        for i, p in enumerate(pairs):
            if not isinstance(p, (tuple, list)) or len(p) != 2:
                raise ValueError(f"pair {i} is not a 2-tuple: {p!r}")
            a, b = p
            if not isinstance(a, str) or not isinstance(b, str):
                raise ValueError(f"pair {i} contains non-string entries: ({type(a).__name__}, {type(b).__name__})")
        self.pairs = list(pairs)

    def __len__(self) -> int:
        return len(self.pairs)

    def __getitem__(self, idx: int) -> tuple[str, str]:
        return self.pairs[idx]

nnx.embeddings.contrastive_trainer.train_contrastive(backbone: Any, dataset: Union[ContrastiveTextDataset, list[tuple[str, str]]], *, n_epochs: int = 3, batch_size: int = 16, lr: float = 2e-05, temperature: float = 0.05, device: Optional[Union[str, torch.device]] = None, shuffle: bool = True, grad_clip_norm: Optional[float] = 1.0, weight_decay: float = 0.0, optimizer_cls: type = torch.optim.AdamW, verbose: bool = False) -> Any

Train backbone on (anchor, positive) pairs via NT-Xent.

High-level wrapper around :func:nt_xent_loss. Builds a :class:DataLoader with :func:pair_collate, instantiates an optimizer over the backbone's trainable parameters, and runs n_epochs of contrastive updates. The backbone is updated in-place AND returned for chaining (e.g., directly into :func:nnx.embeddings.export_to_faiss).

For more elaborate setups — callbacks, custom schedulers, multi- optimizer training, run.id persistence under runs/<id>/ — use :func:text_contrastive_train_step_factory with the standard :meth:NNModel.train driver instead.

Parameters:

Name	Type	Description	Default
backbone	Any	text encoder. Either a :class:sentence_transformers.SentenceTransformer or any nn.Module whose forward(list[str]) -> Tensor[B, D]. Parameters with requires_grad=False are excluded from the optimizer (so :func:nnx.freeze composes cleanly).	required
dataset	Union[ContrastiveTextDataset, list[tuple[str, str]]]	a :class:ContrastiveTextDataset or a plain list of (anchor, positive) string tuples (we'll wrap it).	required
n_epochs	int	number of full passes. Default 3 — contrastive fine-tuning of a pretrained encoder typically needs few.	3
batch_size	int	pairs per batch. NT-Xent's in-batch-negatives scaling means bigger is usually better; 16-64 is typical for CPU sanity runs, hundreds for GPU.	16
lr	float	optimizer learning rate. Default 2e-5 (the canonical SBERT fine-tune LR).	2e-05
temperature	float	NT-Xent temperature. Default 0.05 (sharper than SimCLR's image default — text embedders work in a much higher-dim cosine space where small temperature helps).	0.05
device	Optional[Union[str, device]]	target device. None infers from the backbone (its .device if present, else its first parameter's device, else CPU).	None
shuffle	bool	shuffle the dataset each epoch. Default True.	True
grad_clip_norm	Optional[float]	global L2 grad-clip norm. None to disable; must be positive otherwise (a non-positive norm zeros every gradient). Default 1.0 — text encoders are sensitive to gradient spikes early in fine-tuning.	1.0
weight_decay	float	AdamW weight decay. Default 0.0.	0.0
optimizer_cls	type	optimizer constructor. Default :class:torch.optim.AdamW. Receives (trainable_params, lr=lr, weight_decay=weight_decay).	AdamW
verbose	bool	print per-epoch mean loss. Default False.	False

Returns:

Type	Description
Any	The (in-place-mutated) backbone.

Raises:

Type	Description
ValueError	on a dataset of fewer than 2 pairs, batch_size < 2, non-positive epochs, non-positive temperature, or a non-positive grad_clip_norm — NT-Xent needs at least one negative, so both the dataset and every batch must carry = 2 pairs.
FloatingPointError	when the contrastive loss goes non-finite mid-training (check lr / temperature / input normalization).

Source code in nnx/embeddings/contrastive_trainer.py

def train_contrastive(
    backbone: Any,
    dataset: Union[ContrastiveTextDataset, list[tuple[str, str]]],
    *,
    n_epochs: int = 3,
    batch_size: int = 16,
    lr: float = 2e-5,
    temperature: float = 0.05,
    device: Optional[Union[str, torch.device]] = None,
    shuffle: bool = True,
    grad_clip_norm: Optional[float] = 1.0,
    weight_decay: float = 0.0,
    optimizer_cls: type = torch.optim.AdamW,
    verbose: bool = False,
) -> Any:
    """Train ``backbone`` on ``(anchor, positive)`` pairs via NT-Xent.

    High-level wrapper around :func:`nt_xent_loss`. Builds a
    :class:`DataLoader` with :func:`pair_collate`, instantiates an
    optimizer over the backbone's trainable parameters, and runs
    ``n_epochs`` of contrastive updates. The backbone is updated
    in-place AND returned for chaining (e.g., directly into
    :func:`nnx.embeddings.export_to_faiss`).

    For more elaborate setups — callbacks, custom schedulers, multi-
    optimizer training, run.id persistence under ``runs/<id>/`` — use
    :func:`text_contrastive_train_step_factory` with the standard
    :meth:`NNModel.train` driver instead.

    Args:
        backbone: text encoder. Either a
            :class:`sentence_transformers.SentenceTransformer` or any
            ``nn.Module`` whose ``forward(list[str]) -> Tensor[B, D]``.
            Parameters with ``requires_grad=False`` are excluded from
            the optimizer (so :func:`nnx.freeze` composes cleanly).
        dataset: a :class:`ContrastiveTextDataset` or a plain list of
            ``(anchor, positive)`` string tuples (we'll wrap it).
        n_epochs: number of full passes. Default 3 — contrastive
            fine-tuning of a pretrained encoder typically needs few.
        batch_size: pairs per batch. NT-Xent's in-batch-negatives
            scaling means bigger is usually better; 16-64 is typical
            for CPU sanity runs, hundreds for GPU.
        lr: optimizer learning rate. Default 2e-5 (the canonical SBERT
            fine-tune LR).
        temperature: NT-Xent temperature. Default 0.05 (sharper than
            SimCLR's image default — text embedders work in a much
            higher-dim cosine space where small temperature helps).
        device: target device. ``None`` infers from the backbone (its
            ``.device`` if present, else its first parameter's device,
            else CPU).
        shuffle: shuffle the dataset each epoch. Default True.
        grad_clip_norm: global L2 grad-clip norm. ``None`` to disable;
            must be positive otherwise (a non-positive norm zeros every
            gradient). Default 1.0 — text encoders are sensitive to
            gradient spikes early in fine-tuning.
        weight_decay: AdamW weight decay. Default 0.0.
        optimizer_cls: optimizer constructor. Default
            :class:`torch.optim.AdamW`. Receives
            ``(trainable_params, lr=lr, weight_decay=weight_decay)``.
        verbose: print per-epoch mean loss. Default False.

    Returns:
        The (in-place-mutated) ``backbone``.

    Raises:
        ValueError: on a dataset of fewer than 2 pairs, batch_size < 2,
            non-positive epochs, non-positive temperature, or a
            non-positive ``grad_clip_norm`` — NT-Xent needs at least one
            negative, so both the dataset and every batch must carry
            >= 2 pairs.
        FloatingPointError: when the contrastive loss goes non-finite
            mid-training (check lr / temperature / input normalization).
    """
    if n_epochs <= 0:
        raise ValueError(f"n_epochs must be positive, got {n_epochs}")
    if batch_size < 2:
        # NT-Xent needs at least one negative per batch (nt_xent_loss
        # itself raises on B < 2 as the backstop; rejecting here gives
        # the error before any training work happens).
        raise ValueError(f"batch_size must be >= 2 for NT-Xent (got {batch_size}) — each batch needs a negative.")
    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")
    if grad_clip_norm is not None and grad_clip_norm <= 0:
        # `None` disables clipping; a non-positive norm would otherwise slip
        # past the `is not None` guard below and clip_grad_norm_(..., 0.0)
        # scales every gradient by 0.0/total_norm == 0, silently zeroing all
        # grads so the backbone never learns. Same footgun NNOptimParams
        # guards against; fail fast here before any training work happens.
        raise ValueError(f"grad_clip_norm must be positive or None to disable, got {grad_clip_norm}")

    if isinstance(dataset, list):
        dataset = ContrastiveTextDataset(dataset)
    if len(dataset) < 2:
        # 0 pairs: nothing to train. 1 pair: the only possible batch has
        # no negative, which nt_xent_loss rejects — fail here with the
        # dataset-level message instead.
        raise ValueError(f"train_contrastive needs >= 2 pairs (got {len(dataset)}) — NT-Xent requires negatives.")

    device = _resolve_device(backbone, device)
    backbone.to(device)

    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        collate_fn=pair_collate,
        # Drop the trailing batch ONLY when it would have size 1 — a
        # single pair has no negative (nt_xent_loss raises on B < 2).
        # Larger partial batches carry real contrastive signal and are
        # kept. Together with the batch_size >= 2 and len(dataset) >= 2
        # guards above, this makes size-1 batches unreachable; if a
        # future loader change reintroduced one, nt_xent_loss fails
        # loudly rather than silently no-op'ing.
        drop_last=(len(dataset) > batch_size and len(dataset) % batch_size == 1),
    )

    trainable = [p for p in backbone.parameters() if p.requires_grad]
    if not trainable:
        raise ValueError(
            "backbone has no trainable parameters. Did you accidentally freeze "
            "everything? Use nnx.unfreeze(backbone, '*') to thaw it."
        )
    optimizer = optimizer_cls(trainable, lr=lr, weight_decay=weight_decay)

    backbone.train()
    for epoch in range(n_epochs):
        epoch_losses: list[float] = []
        for anchors, positives in loader:
            optimizer.zero_grad()
            z1 = _encode(backbone, anchors, device)
            z2 = _encode(backbone, positives, device)
            loss = nt_xent_loss(z1, z2, temperature=temperature)
            loss_val = float(loss.detach())
            if not torch.isfinite(loss).item():
                raise FloatingPointError(
                    f"non-finite contrastive loss ({loss_val!r}) at epoch "
                    f"{epoch}. Check lr / temperature / input normalization."
                )
            loss.backward()
            if grad_clip_norm is not None:
                torch.nn.utils.clip_grad_norm_(trainable, grad_clip_norm)
            optimizer.step()
            epoch_losses.append(loss_val)
        if verbose:
            mean = sum(epoch_losses) / max(1, len(epoch_losses))
            print(f"[epoch {epoch + 1}/{n_epochs}] mean NT-Xent loss = {mean:.4f}")

    return backbone

nnx.embeddings.contrastive_trainer.embed_texts(backbone: Any, texts: list[str], *, batch_size: int = 64, device: Optional[Union[str, torch.device]] = None, normalize: bool = True) -> torch.Tensor

Encode texts with backbone and return a (N, D) tensor.

Runs in torch.no_grad() + eval() mode — this is the inference helper, not the training one. The trainer drives :func:_encode directly so gradients flow.

Parameters:

Name	Type	Description	Default
backbone	Any	text encoder — a sentence-transformers model or any nn.Module whose forward(list[str]) -> Tensor[B, D].	required
texts	list[str]	input strings. May be empty (returns a (0, ?) placeholder — the embedding dim isn't known until the first forward).	required
batch_size	int	how many texts per forward pass. Default 64.	64
device	Optional[Union[str, device]]	target device. None uses the backbone's device (sentence-transformers exposes one; plain Modules don't, in which case we fall back to the first parameter's device, or CPU when the backbone has no parameters).	None
normalize	bool	if True, L2-normalize each row so dot products with the result are cosine similarities. Default True because FAISS's IndexFlatIP interprets the inner product as a similarity score and the standard cosine-by-IP trick is normalize-then-IP.	True

Returns:

Type	Description
Tensor	A (N, D) torch.Tensor on device. Detached from
Tensor	any autograd graph.

Source code in nnx/embeddings/contrastive_trainer.py

def embed_texts(
    backbone: Any,
    texts: list[str],
    *,
    batch_size: int = 64,
    device: Optional[Union[str, torch.device]] = None,
    normalize: bool = True,
) -> torch.Tensor:
    """Encode ``texts`` with ``backbone`` and return a ``(N, D)`` tensor.

    Runs in ``torch.no_grad()`` + ``eval()`` mode — this is the
    inference helper, not the training one. The trainer drives
    :func:`_encode` directly so gradients flow.

    Args:
        backbone: text encoder — a sentence-transformers model or any
            ``nn.Module`` whose ``forward(list[str]) -> Tensor[B, D]``.
        texts: input strings. May be empty (returns a ``(0, ?)``
            placeholder — the embedding dim isn't known until the
            first forward).
        batch_size: how many texts per forward pass. Default 64.
        device: target device. ``None`` uses the backbone's device
            (sentence-transformers exposes one; plain Modules don't, in
            which case we fall back to the first parameter's device,
            or CPU when the backbone has no parameters).
        normalize: if True, L2-normalize each row so dot products with
            the result are cosine similarities. Default True because
            FAISS's ``IndexFlatIP`` interprets the inner product as a
            similarity score and the standard cosine-by-IP trick is
            normalize-then-IP.

    Returns:
        A ``(N, D)`` ``torch.Tensor`` on ``device``. Detached from
        any autograd graph.
    """
    device = _resolve_device(backbone, device)

    if not texts:
        # No way to know D without a forward; return an empty 2D tensor
        # so callers don't get a confusing 1D shape.
        return torch.empty((0, 0), device=device)

    # Snapshot training-mode for non-destructive restore (matches the
    # convention used by NNModel.predict / evaluate,
    # GenerativeNNModel.generate, diffusion.sample,
    # nnx.viz.activation_map, and nnx.lr_finder).
    was_training = backbone.training
    backbone.eval()
    chunks: list[torch.Tensor] = []
    try:
        with torch.no_grad():
            for i in range(0, len(texts), batch_size):
                chunk = texts[i : i + batch_size]
                emb = _encode(backbone, chunk, device)
                if normalize:
                    emb = F.normalize(emb, dim=-1)
                chunks.append(emb.detach())
    finally:
        if was_training:
            backbone.train()
    return torch.cat(chunks, dim=0)

nnx.embeddings.contrastive_trainer.text_contrastive_train_step_factory(*, temperature: float = 0.5) -> TrainStepFn

Build a :class:TrainStepFn for text-pair contrastive training.

This is the text-aware sibling of :func:nnx.simclr_train_step_factory. The training loader must yield (anchors: list[str], positives: list[str]) batches — typically by pairing :class:ContrastiveTextDataset with :func:pair_collate.

The step runs:

1. Encode anchors through model.net → z1.
2. Encode positives through model.net → z2.
3. NT-Xent loss across the (2B, 2B) similarity matrix.
4. Standard :func:finalize_step tail (NaN guard, grad clip, optimizer step).

Parameters:

Name	Type	Description	Default
temperature	float	NT-Xent temperature. Lower sharpens; 0.5 is the SimCLR default. Must be > 0.	0.5

Returns:

Type	Description
TrainStepFn	A TrainStepFn suitable for NNModel.train(..., train_step_fn=...).

Raises:

Type	Description
ValueError	at factory-build time if temperature ≤ 0.

Source code in nnx/embeddings/contrastive_trainer.py

def text_contrastive_train_step_factory(*, temperature: float = 0.5) -> TrainStepFn:
    """Build a :class:`TrainStepFn` for text-pair contrastive training.

    This is the text-aware sibling of
    :func:`nnx.simclr_train_step_factory`. The training loader must
    yield ``(anchors: list[str], positives: list[str])`` batches —
    typically by pairing :class:`ContrastiveTextDataset` with
    :func:`pair_collate`.

    The step runs:

      1. Encode anchors through ``model.net`` → ``z1``.
      2. Encode positives through ``model.net`` → ``z2``.
      3. NT-Xent loss across the ``(2B, 2B)`` similarity matrix.
      4. Standard :func:`finalize_step` tail (NaN guard, grad clip,
         optimizer step).

    Args:
        temperature: NT-Xent temperature. Lower sharpens; 0.5 is the
            SimCLR default. Must be > 0.

    Returns:
        A ``TrainStepFn`` suitable for ``NNModel.train(..., train_step_fn=...)``.

    Raises:
        ValueError: at factory-build time if ``temperature`` ≤ 0.
    """
    if temperature <= 0:
        raise ValueError(f"temperature must be positive, got {temperature}")

    def step(ctx: TrainStepContext) -> NNEvaluationDataPoint:
        m = ctx.model
        m.net.train()
        m.net.zero_grad()

        batch = ctx.batch
        if not (isinstance(batch, (tuple, list)) and len(batch) == 2):
            raise ValueError(
                "text contrastive step expects a batch of "
                "(anchors: list[str], positives: list[str]). Got "
                f"{type(batch).__name__} with {len(batch) if hasattr(batch, '__len__') else '?'} entries."
            )
        anchors, positives = batch
        if not (isinstance(anchors, list) and isinstance(positives, list)):
            raise ValueError(
                "text contrastive step expects list[str] views — pair "
                "ContrastiveTextDataset with pair_collate when constructing "
                "the DataLoader. Got "
                f"anchors={type(anchors).__name__}, positives={type(positives).__name__}."
            )

        z1 = _encode(m.net, anchors, m.device)
        z2 = _encode(m.net, positives, m.device)

        loss = nt_xent_loss(z1, z2, temperature=temperature)
        loss_val = finalize_step(loss, ctx, paradigm="embeddings.contrastive")

        return NNEvaluationDataPoint(
            f1=0.0,
            recall=0.0,
            accuracy=0.0,
            precision=0.0,
            loss=loss_val,
            error=loss_val,
        )

    return step

nnx.embeddings.faiss_export.export_to_faiss(backbone: Any, corpus: list[str], out_path: Union[str, Path], *, batch_size: int = 64, index_type: str = 'IndexFlatIP', normalize: Optional[bool] = None, device: Optional[Union[str, torch.device]] = None) -> str

Embed corpus with backbone and write a FAISS index file.

The default IndexFlatIP + normalize=True combination is cosine similarity: L2-normalize the embeddings, then use inner product as the score. This is the standard FAISS-cosine recipe (FAISS itself doesn't ship a cosine index; the normalize-then-IP pattern is canonical).

The corpus order is preserved in the index — index.search's returned ids are positions into corpus. The caller is responsible for keeping a parallel list / DataFrame of original document ids or metadata.

Parameters:

Name	Type	Description	Default
backbone	Any	text encoder. Either a :class:sentence_transformers.SentenceTransformer or any nn.Module whose forward(list[str]) -> Tensor[B, D].	required
corpus	list[str]	list of strings to embed. Order is the index's id space. Empty raises :class:ValueError — FAISS rejects 0-length adds.	required
out_path	Union[str, Path]	destination file path. The parent directory must exist. The file is written via FAISS's native write_index (atomic depends on the underlying FS).	required
batch_size	int	forward-pass batch size. Default 64.	64
index_type	str	FAISS index family to build. One of "IndexFlatIP" (default), "IndexFlatL2", "IndexHNSWFlat".	'IndexFlatIP'
normalize	Optional[bool]	whether to L2-normalize each embedding before insertion. None (the default) auto-selects: True for IndexFlatIP (cosine via IP), False for everything else. Pass an explicit bool to override.	None
device	Optional[Union[str, device]]	target device for the encode pass. None infers from the backbone.	None

Returns:

Type	Description
str	The string path written. Same value as str(out_path) —
str	returned for call-chain convenience.

Raises:

Type	Description
ImportError	if faiss isn't installed (lazy import; only this call requires it).
ValueError	empty corpus, unknown index_type.

Source code in nnx/embeddings/faiss_export.py

def export_to_faiss(
    backbone: Any,
    corpus: list[str],
    out_path: Union[str, Path],
    *,
    batch_size: int = 64,
    index_type: str = "IndexFlatIP",
    normalize: Optional[bool] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> str:
    """Embed ``corpus`` with ``backbone`` and write a FAISS index file.

    The default ``IndexFlatIP`` + ``normalize=True`` combination is
    cosine similarity: L2-normalize the embeddings, then use inner
    product as the score. This is the standard FAISS-cosine recipe
    (FAISS itself doesn't ship a cosine index; the normalize-then-IP
    pattern is canonical).

    The corpus order is preserved in the index — ``index.search``'s
    returned ids are positions into ``corpus``. The caller is
    responsible for keeping a parallel list / DataFrame of original
    document ids or metadata.

    Args:
        backbone: text encoder. Either a
            :class:`sentence_transformers.SentenceTransformer` or any
            ``nn.Module`` whose ``forward(list[str]) -> Tensor[B, D]``.
        corpus: list of strings to embed. Order is the index's id space.
            Empty raises :class:`ValueError` — FAISS rejects 0-length
            adds.
        out_path: destination file path. The parent directory must
            exist. The file is written via FAISS's native
            ``write_index`` (atomic depends on the underlying FS).
        batch_size: forward-pass batch size. Default 64.
        index_type: FAISS index family to build. One of
            ``"IndexFlatIP"`` (default), ``"IndexFlatL2"``,
            ``"IndexHNSWFlat"``.
        normalize: whether to L2-normalize each embedding before
            insertion. ``None`` (the default) auto-selects: True for
            ``IndexFlatIP`` (cosine via IP), False for everything else.
            Pass an explicit bool to override.
        device: target device for the encode pass. ``None`` infers
            from the backbone.

    Returns:
        The string path written. Same value as ``str(out_path)`` —
        returned for call-chain convenience.

    Raises:
        ImportError: if ``faiss`` isn't installed (lazy import; only
            this call requires it).
        ValueError: empty corpus, unknown ``index_type``.
    """
    if not corpus:
        raise ValueError("corpus is empty — FAISS rejects zero-row adds")

    faiss = _import_faiss()

    if normalize is None:
        # Cosine-via-IP is the default reason callers pick IndexFlatIP,
        # so auto-normalize there. L2 / HNSW shouldn't be normalized
        # unless the caller specifically wants unit-sphere geometry.
        normalize = index_type == "IndexFlatIP"

    emb = embed_texts(
        backbone,
        corpus,
        batch_size=batch_size,
        device=device,
        normalize=normalize,
    )
    # FAISS expects float32 contiguous arrays on CPU.
    emb_np = emb.detach().cpu().to(torch.float32).contiguous().numpy()
    dim = emb_np.shape[1]

    index = _build_index(faiss, dim, index_type)
    index.add(emb_np)

    out_path = str(out_path)
    faiss.write_index(index, out_path)
    return out_path

nnx.embeddings.faiss_export.export_to_safetensors(backbone: Any, out_path: Union[str, Path]) -> str

Persist backbone.state_dict() to disk for downstream reload.

Prefers the safetensors format (canonical for HuggingFace Hub artifacts and sentence-transformers ≥3) when the :mod:safetensors package is importable. Falls back to plain :func:torch.save when it isn't, so the function still works on a vanilla pip install thekaveh-nnx without the embeddings extra. In the fallback case out_path is written as a pickle blob; the caller's reloader needs to use :func:torch.load.

Parameters:

Name	Type	Description	Default
backbone	Any	anything with a state_dict() method. Sentence-transformers, raw nn.Module, even a plain OrderedDict of tensors.	required
out_path	Union[str, Path]	destination file path. Conventionally suffixed .safetensors for the primary path; .pt for the torch.save fallback. We don't enforce the suffix — that's cosmetic.	required

Returns:

Type	Description
str	The string path written.

Source code in nnx/embeddings/faiss_export.py

def export_to_safetensors(backbone: Any, out_path: Union[str, Path]) -> str:
    """Persist ``backbone.state_dict()`` to disk for downstream reload.

    Prefers the ``safetensors`` format (canonical for HuggingFace Hub
    artifacts and sentence-transformers ≥3) when the
    :mod:`safetensors` package is importable. Falls back to plain
    :func:`torch.save` when it isn't, so the function still works on
    a vanilla ``pip install thekaveh-nnx`` without the embeddings extra. In
    the fallback case ``out_path`` is written as a pickle blob; the
    caller's reloader needs to use :func:`torch.load`.

    Args:
        backbone: anything with a ``state_dict()`` method.
            Sentence-transformers, raw ``nn.Module``, even a plain
            ``OrderedDict`` of tensors.
        out_path: destination file path. Conventionally suffixed
            ``.safetensors`` for the primary path; ``.pt`` for the
            torch.save fallback. We don't enforce the suffix — that's
            cosmetic.

    Returns:
        The string path written.
    """
    out_path = str(out_path)
    sd = backbone.state_dict() if hasattr(backbone, "state_dict") else backbone

    try:
        from safetensors.torch import save_file  # type: ignore[import-not-found]
    except ImportError:
        # No safetensors — plain torch.save fallback. Pickle blob,
        # so reload with torch.load(weights_only=True) on trusted files.
        torch.save(sd, out_path)
        return out_path

    # safetensors rejects tensors that share storage (e.g. tied
    # embedding/head weights — TransformerNN's default, and common in
    # HF models) and requires contiguous memory. `.contiguous()` is a
    # no-op on an already-contiguous shared tensor, so `.clone()` is
    # what actually breaks the aliasing; detach + CPU keep the file
    # device-portable.
    cleaned: dict[str, torch.Tensor] = {}
    for k, v in sd.items():
        if not isinstance(v, torch.Tensor):
            # Skip non-tensor state (rare — schedulers / metadata).
            # The user can save those separately if they need them.
            continue
        cleaned[k] = v.detach().cpu().contiguous().clone()

    save_file(cleaned, out_path)
    return out_path

16. Interop (nnx.interop)

16.1. GGUF + Ollama

End-to-end walkthrough at GGUF & Ollama. Opt-in via pip install "thekaveh-nnx[gguf-write]".

nnx.interop.gguf.writer.write_gguf(transformer_nn: TransformerNN, tokenizer: NNTokenizerParams, out_path: str | os.PathLike, *, architecture: str = 'nnx_transformer', quantization: str = 'F16', model_name: Optional[str] = None) -> str

Write a TransformerNN + tokenizer to a single .gguf file.

Parameters:

Name	Type	Description	Default
transformer_nn	TransformerNN	A nnx.TransformerNN instance. The forward path's tensors are exported under llama.cpp's tensor-naming convention (see tensor_name_map.map_tensors).	required
tokenizer	NNTokenizerParams	An nnx.NNTokenizerParams (or any object with a .tokenizer attribute exposing .get_vocab() and .get_vocab_size()). Tokens + merges are emitted under the GGUF tokenizer keys.	required
out_path	str | PathLike	Destination .gguf path.	required
architecture	str	general.architecture metadata value. Defaults to "nnx_transformer" — readable by patched llama.cpp forks. Pass "llama" to claim LLaMA-arch compatibility for stock llama.cpp readers (works because we match the LLaMA tensor naming + RMSNorm/SwiGLU/RoPE choices, but users should verify their target reader version).	'nnx_transformer'
quantization	str	One of "F32", "F16", "BF16". Sub-F16 quantizations require the C++ llama-quantize binary — see the ImportError message for the shell-out recipe.	'F16'
model_name	Optional[str]	general.name metadata. Defaults to a "nnx_transformer_LxD" shape-derived name.	None

Returns:

Type	Description
str	The absolute path of the written file as a string.

Raises:

Type	Description
ImportError	when gguf is not installed, or when a quantization is requested that requires llama-quantize.
ValueError	when an unknown quantization label is passed.

Source code in nnx/interop/gguf/writer.py

def write_gguf(
    transformer_nn: TransformerNN,
    tokenizer: NNTokenizerParams,
    out_path: str | os.PathLike,
    *,
    architecture: str = "nnx_transformer",
    quantization: str = "F16",
    model_name: Optional[str] = None,
) -> str:
    """Write a TransformerNN + tokenizer to a single ``.gguf`` file.

    Args:
        transformer_nn: A ``nnx.TransformerNN`` instance. The forward
            path's tensors are exported under llama.cpp's tensor-naming
            convention (see ``tensor_name_map.map_tensors``).
        tokenizer: An ``nnx.NNTokenizerParams`` (or any object with a
            ``.tokenizer`` attribute exposing ``.get_vocab()`` and
            ``.get_vocab_size()``). Tokens + merges are emitted under
            the GGUF tokenizer keys.
        out_path: Destination ``.gguf`` path.
        architecture: ``general.architecture`` metadata value. Defaults
            to ``"nnx_transformer"`` — readable by patched llama.cpp
            forks. Pass ``"llama"`` to claim LLaMA-arch compatibility
            for stock llama.cpp readers (works because we match the
            LLaMA tensor naming + RMSNorm/SwiGLU/RoPE choices, but
            users should verify their target reader version).
        quantization: One of ``"F32"``, ``"F16"``, ``"BF16"``. Sub-F16
            quantizations require the C++ ``llama-quantize`` binary —
            see the ``ImportError`` message for the shell-out recipe.
        model_name: ``general.name`` metadata. Defaults to a
            ``"nnx_transformer_LxD"`` shape-derived name.

    Returns:
        The absolute path of the written file as a string.

    Raises:
        ImportError: when ``gguf`` is not installed, or when a
            quantization is requested that requires ``llama-quantize``.
        ValueError: when an unknown quantization label is passed.
    """
    import torch

    gguf = _require_gguf()
    _validate_quantization(quantization)

    # Local import — avoids importing torch / numpy at module top so
    # ``import nnx.interop`` is cheap even when the writer isn't used.
    from .tensor_name_map import map_tensors

    out_path = str(Path(out_path).expanduser().resolve())
    # Create the parent directory — GGUFWriter open()s the path
    # directly, so a fresh cwd writing "out/model.gguf" otherwise dies
    # with FileNotFoundError (the ollama exporter already mkdirs).
    Path(out_path).parent.mkdir(parents=True, exist_ok=True)

    params = transformer_nn.params
    # NNParams declares ``n_heads`` as Optional[int]; NNTransformerParams
    # validates it's non-None in ``__post_init__``. Re-narrow here so
    # pyright / mypy can prove the division below is safe.
    assert params.n_heads is not None, "TransformerNN params must have n_heads set"
    n_heads = params.n_heads
    if model_name is None:
        model_name = f"nnx_transformer_L{params.n_layers}_D{params.d_model}"

    # Compute the SwiGLU hidden dim. The convention matches
    # ``SwiGLU.__init__`` exactly so a reader reconstructing the
    # architecture from metadata picks the same shape we shipped.
    feed_forward_length = int(2 * params.ffn_mult * params.d_model / 3)
    head_dim = params.d_model // n_heads

    writer = gguf.GGUFWriter(out_path, arch=architecture)

    # try/finally: GGUFWriter opens the file lazily inside
    # write_header_to_file, so an error in the flush sequence (or in a
    # future gguf version that opens earlier) would otherwise leak the
    # handle alongside the partial .gguf. close() is idempotent.
    try:
        # ---------- general metadata ----------
        # ``GGUFWriter(arch=...)`` already records ``general.architecture``;
        # don't call ``add_architecture()`` again (it overwrites with the
        # same value and emits a duplicate-key warning).
        writer.add_name(model_name)
        writer.add_file_type(_gguf_file_type_for(gguf, quantization))

        # ---------- architecture metadata (llama.cpp keys) ----------
        # These are the keys llama.cpp reads when reconstructing a model;
        # populating them lets a patched reader rebuild the right shapes
        # without consulting NNx-specific config.
        writer.add_context_length(params.max_seq_len)
        writer.add_block_count(params.n_layers)
        writer.add_embedding_length(params.d_model)
        writer.add_feed_forward_length(feed_forward_length)
        writer.add_head_count(n_heads)
        writer.add_head_count_kv(n_heads)  # NNx is MHA (not GQA) — kv heads == q heads
        writer.add_layer_norm_rms_eps(1e-6)  # matches RMSNorm default in transformer_layers.py
        writer.add_rope_freq_base(params.rope_base)
        # The RoPE dimension count is the per-head dim (rotary is applied
        # per head). llama.cpp expects this exact number.
        writer.add_rope_dimension_count(head_dim)

        # ---------- tokenizer ----------
        # ``model_name`` here is the *tokenizer* model name ("llama", "gpt2",
        # "bpe", ...). NNx ships a HF tokenizers.Tokenizer trained with the
        # BPE model — emit it as "gpt2" so llama.cpp's BPE path picks it up.
        writer.add_tokenizer_model("gpt2")
        writer.add_tokenizer_pre("default")  # pre-tokenizer is Whitespace; "default" is the catch-all
        _add_tokenizer_vocab(writer, tokenizer)

        # ---------- tensors ----------
        target_dtype = _quantization_torch_dtype(quantization)
        tensors = map_tensors(transformer_nn)
        for name, arr in tensors.items():
            # The map returns numpy arrays already in F32. Cast via torch
            # so we get a single code path for F16 / BF16 / F32 (numpy
            # doesn't have a native bfloat16 dtype; torch + ml_dtypes is
            # how the upstream gguf package handles it).
            t = torch.from_numpy(arr).to(target_dtype)
            if target_dtype == torch.bfloat16:
                # GGUFWriter accepts bfloat16 via numpy when the user
                # passes raw_dtype=BF16; we route through that contract
                # because numpy proper has no bfloat16 scalar.
                np_arr = t.view(torch.uint16).numpy()
                writer.add_tensor(name, np_arr, raw_dtype=gguf.GGMLQuantizationType.BF16)
            else:
                writer.add_tensor(name, t.numpy())

        # ---------- flush ----------
        writer.write_header_to_file()
        writer.write_kv_data_to_file()
        writer.write_tensors_to_file()
    finally:
        writer.close()

    return out_path

nnx.interop.gguf.tensor_name_map.map_tensors(net: TransformerNN) -> dict[str, np.ndarray]

Walk a TransformerNN and return {gguf_name: numpy_array}.

The caller (write_gguf) then iterates this dict and calls GGUFWriter.add_tensor for each entry. Splitting the iteration here (rather than inlining it into the writer) keeps the naming convention testable in isolation — see test_gguf_writer.py.

Parameters:

Name	Type	Description	Default
net	TransformerNN	A TransformerNN instance.	required

Returns:

Type	Description
dict[str, ndarray]	Dict gguf_name -> numpy.ndarray. Q/K/V are emitted as three
dict[str, ndarray]	separate tensors even though the NNx side stores them fused.
dict[str, ndarray]	When net.params.tie_embeddings is True, output.weight
dict[str, ndarray]	is omitted (llama.cpp re-uses token_embd.weight for tied
dict[str, ndarray]	models).

Source code in nnx/interop/gguf/tensor_name_map.py

def map_tensors(net: TransformerNN) -> dict[str, np.ndarray]:
    """Walk a ``TransformerNN`` and return ``{gguf_name: numpy_array}``.

    The caller (``write_gguf``) then iterates this dict and calls
    ``GGUFWriter.add_tensor`` for each entry. Splitting the iteration
    here (rather than inlining it into the writer) keeps the naming
    convention testable in isolation — see ``test_gguf_writer.py``.

    Args:
        net: A ``TransformerNN`` instance.

    Returns:
        Dict ``gguf_name -> numpy.ndarray``. Q/K/V are emitted as three
        separate tensors even though the NNx side stores them fused.
        When ``net.params.tie_embeddings`` is True, ``output.weight``
        is omitted (llama.cpp re-uses ``token_embd.weight`` for tied
        models).
    """
    out: dict[str, np.ndarray] = {}

    # Token embedding — always present, always named ``token_embd.weight``.
    out["token_embd.weight"] = _to_numpy(_weight(net.tok_embed))

    d_model = net.params.d_model

    for i, block in enumerate(net.blocks):
        # Pre-attention norm (RMSNorm).
        out[f"blk.{i}.attn_norm.weight"] = _to_numpy(_weight(block.norm1))

        # Split the fused QKV projection. Layout in NNx:
        #   w_qkv.weight has shape (3 * d_model, d_model) and the
        #   forward path does w_qkv(x).chunk(3, dim=-1) on the OUTPUT,
        #   i.e. the leading dim is the concatenated [Q | K | V] block.
        # That maps cleanly to a row-wise slice of the weight tensor.
        w_qkv = _weight(block.attn.w_qkv)  # (3*d, d)
        out[f"blk.{i}.attn_q.weight"] = _to_numpy(w_qkv[0:d_model, :])
        out[f"blk.{i}.attn_k.weight"] = _to_numpy(w_qkv[d_model : 2 * d_model, :])
        out[f"blk.{i}.attn_v.weight"] = _to_numpy(w_qkv[2 * d_model : 3 * d_model, :])

        # Attention output projection.
        out[f"blk.{i}.attn_output.weight"] = _to_numpy(_weight(block.attn.w_o))

        # Pre-FFN norm.
        out[f"blk.{i}.ffn_norm.weight"] = _to_numpy(_weight(block.norm2))

        # SwiGLU: NNx names them w1 (gate) / w2 (down) / w3 (up);
        # llama.cpp calls the same matrices ffn_gate / ffn_down / ffn_up.
        out[f"blk.{i}.ffn_gate.weight"] = _to_numpy(_weight(block.ffn.w1))
        out[f"blk.{i}.ffn_down.weight"] = _to_numpy(_weight(block.ffn.w2))
        out[f"blk.{i}.ffn_up.weight"] = _to_numpy(_weight(block.ffn.w3))

    # Final output norm — llama.cpp calls this ``output_norm.weight``.
    out["output_norm.weight"] = _to_numpy(_weight(net.norm_out))

    # LM head. When embeddings are tied (NNx default), llama.cpp re-uses
    # ``token_embd.weight`` and we deliberately omit ``output.weight`` so
    # the file size halves and the reader doesn't see two copies of the
    # same matrix.
    if not net.params.tie_embeddings:
        out["output.weight"] = _to_numpy(_weight(net.lm_head))

    return out

nnx.interop.ollama.export_ollama_modelfile(transformer_nn: TransformerNN, tokenizer: NNTokenizerParams, out_dir: str | os.PathLike, *, system: str = '', parameters: Optional[dict] = None, template: Optional[str] = None, quantization: str = 'F16', model_name: Optional[str] = None) -> str

Emit model.gguf + Modelfile into out_dir.

Parameters:

Name	Type	Description	Default
transformer_nn	TransformerNN	An NNx TransformerNN instance — the model to export.	required
tokenizer	NNTokenizerParams	Corresponding NNTokenizerParams.	required
out_dir	str | PathLike	Output directory. Created if it doesn't exist.	required
system	str	Optional system prompt; emitted as a SYSTEM ... block (triple-quoted) when non-empty. Must not contain a triple-quote or end with a double-quote (Modelfile block delimiters — validated, raises ValueError); same constraint applies to template.	''
parameters	Optional[dict]	Optional dict of Ollama runtime parameters ({"temperature": 0.8, "top_k": 40} etc.). Each entry becomes a PARAMETER <key> <value> line. Lists become multiple PARAMETER <key> <item> lines (Ollama's convention for things like stop).	None
template	Optional[str]	Optional chat template. Emitted as a TEMPLATE ... block (triple-quoted) when set.	None
quantization	str	Forwarded to :func:write_gguf. Defaults to F16.	'F16'
model_name	Optional[str]	Forwarded to :func:write_gguf as model_name.	None

Returns:

Type	Description
str	Absolute path to the emitted Modelfile.

Source code in nnx/interop/ollama.py

def export_ollama_modelfile(
    transformer_nn: TransformerNN,
    tokenizer: NNTokenizerParams,
    out_dir: str | os.PathLike,
    *,
    system: str = "",
    parameters: Optional[dict] = None,
    template: Optional[str] = None,
    quantization: str = "F16",
    model_name: Optional[str] = None,
) -> str:
    """Emit ``model.gguf`` + ``Modelfile`` into ``out_dir``.

    Args:
        transformer_nn: An NNx ``TransformerNN`` instance — the model
            to export.
        tokenizer: Corresponding ``NNTokenizerParams``.
        out_dir: Output directory. Created if it doesn't exist.
        system: Optional system prompt; emitted as a ``SYSTEM ...``
            block (triple-quoted) when non-empty. Must not contain a
            triple-quote or end with a double-quote (Modelfile block
            delimiters — validated, raises ``ValueError``); same
            constraint applies to ``template``.
        parameters: Optional dict of Ollama runtime parameters
            (``{"temperature": 0.8, "top_k": 40}`` etc.). Each entry
            becomes a ``PARAMETER <key> <value>`` line. Lists become
            multiple ``PARAMETER <key> <item>`` lines (Ollama's
            convention for things like ``stop``).
        template: Optional chat template. Emitted as a
            ``TEMPLATE ...`` block (triple-quoted) when set.
        quantization: Forwarded to :func:`write_gguf`. Defaults to F16.
        model_name: Forwarded to :func:`write_gguf` as ``model_name``.

    Returns:
        Absolute path to the emitted ``Modelfile``.
    """
    _validate_modelfile_inputs(system, template, parameters)

    # Local import — keeps the ``nnx.interop`` boot path light when
    # only one of the two exporters is used, and matches the lazy-import
    # style in writer.py.
    from .gguf import write_gguf

    out_dir = Path(out_dir).expanduser().resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    gguf_path = out_dir / "model.gguf"
    write_gguf(
        transformer_nn,
        tokenizer,
        out_path=gguf_path,
        quantization=quantization,
        model_name=model_name,
    )

    modelfile_path = out_dir / "Modelfile"
    # encoding="utf-8" explicit so a non-ASCII SYSTEM prompt / TEMPLATE
    # (Asian-language fine-tunes, emoji prompts) round-trips correctly
    # on Windows pre-PEP-686, where Path.write_text would otherwise fall
    # back to the locale code page (cp1252) and silently mojibake.
    # Matches the convention every other text-mode write in src/nnx
    # carries (`feedback_utf8_explicit_text_opens`).
    modelfile_path.write_text(
        _render_modelfile(system=system, parameters=parameters, template=template),
        encoding="utf-8",
    )
    return str(modelfile_path)

17. HuggingFace Hub + safetensors

Opt-in via pip install "thekaveh-nnx[hub]". Two integration surfaces:

• safetensors checkpoints — NNCheckpoint.to_file(..., format="safetensors") and NNCheckpoint.from_file(..., format="safetensors") (see §3 NNCheckpoint) read and write checkpoints in the safetensors format alongside the default pickle path. Loadable by outside-Python tools (ComfyUI, vLLM, AutoGPTQ).
• Hub publish / load — NNModel mixes in huggingface_hub.PyTorchModelHubMixin, so save_pretrained(local_dir), push_to_hub(repo_id), and NNModel.from_pretrained(repo_id) work directly on a trained model. The mixin methods are inherited and live on NNModel itself — see §2.1.

Walkthrough at HuggingFace Hub.

18. Generation (nnx.generation)

LogitsProcessor chain for autoregressive sampling. Used by GenerativeNNModel.generate() (§2.2). Pure-torch — no optional deps.

nnx.generation.LogitsProcessor

Bases: Protocol

Callable protocol: logits, token_history -> adjusted_logits.

token_history is a flat list of int token ids generated so far (across batch dim 0 — we assume a single-sequence batch in GenerativeNNModel.generate). Processors that don't care about history (temperature, top-k, top-p) simply ignore the arg.

Source code in nnx/generation/logits_processors.py

class LogitsProcessor(Protocol):
    """Callable protocol: ``logits, token_history -> adjusted_logits``.

    ``token_history`` is a flat list of int token ids generated so far
    (across batch dim 0 — we assume a single-sequence batch in
    ``GenerativeNNModel.generate``). Processors that don't care about
    history (temperature, top-k, top-p) simply ignore the arg.
    """

    def __call__(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor: ...

nnx.generation.LogitsChain dataclass

A typed, ordered sequence of LogitsProcessors.

Build via LogitsChain.builder() for the safe / discoverable path; or construct directly from a list for advanced cases. The .apply() method runs the processors against a logits tensor in order, returning the adjusted tensor.

Source code in nnx/generation/logits_chain.py

@dataclass(frozen=True, kw_only=True, slots=True)
class LogitsChain:
    """A typed, ordered sequence of LogitsProcessors.

    Build via `LogitsChain.builder()` for the safe / discoverable
    path; or construct directly from a list for advanced cases. The
    `.apply()` method runs the processors against a logits tensor in
    order, returning the adjusted tensor.
    """

    processors: list[LogitsProcessor] = field(default_factory=list)

    def apply(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor:
        """Run every processor in `self.processors` in order. Thin
        wrapper around `apply_chain`."""
        return apply_chain(logits, token_history=token_history, processors=self.processors)

    @classmethod
    def builder(cls) -> LogitsChainBuilder:
        """Return a fluent builder. See `LogitsChainBuilder`."""
        return LogitsChainBuilder()

apply(logits: torch.Tensor, token_history: list[int]) -> torch.Tensor

Run every processor in self.processors in order. Thin wrapper around apply_chain.

Source code in nnx/generation/logits_chain.py

def apply(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor:
    """Run every processor in `self.processors` in order. Thin
    wrapper around `apply_chain`."""
    return apply_chain(logits, token_history=token_history, processors=self.processors)

builder() -> LogitsChainBuilder classmethod

Return a fluent builder. See LogitsChainBuilder.

Source code in nnx/generation/logits_chain.py

@classmethod
def builder(cls) -> LogitsChainBuilder:
    """Return a fluent builder. See `LogitsChainBuilder`."""
    return LogitsChainBuilder()

nnx.generation.LogitsChainBuilder

Fluent builder for a LogitsChain.

Method order at the call site doesn't matter — .build() sorts the standard processors into NNx's canonical order (matching generate()'s inline-kwargs chain; see the module docstring for why temperature is deliberately last): RepetitionPenalty → TopKFilter → TopPFilter → TemperatureScaling. Custom processors (added via .custom(processor)) are appended in the order they were added, after the canonical group.

Source code in nnx/generation/logits_chain.py

class LogitsChainBuilder:
    """Fluent builder for a `LogitsChain`.

    Method order at the call site doesn't matter — `.build()` sorts
    the standard processors into NNx's canonical order (matching
    `generate()`'s inline-kwargs chain; see the module docstring for
    why temperature is deliberately last):
    `RepetitionPenalty → TopKFilter → TopPFilter → TemperatureScaling`.
    Custom processors (added via `.custom(processor)`) are appended in
    the order they were added, after the canonical group.
    """

    def __init__(self) -> None:
        self._standard: dict[type[LogitsProcessor], LogitsProcessor] = {}
        self._custom: list[LogitsProcessor] = []

    def repetition_penalty(self, penalty: float) -> LogitsChainBuilder:
        """Add a RepetitionPenalty processor with the given penalty."""
        self._standard[RepetitionPenalty] = RepetitionPenalty(penalty=penalty)
        return self

    def top_k(self, k: int) -> LogitsChainBuilder:
        """Add a TopKFilter with the given k."""
        self._standard[TopKFilter] = TopKFilter(top_k=k)
        return self

    def top_p(self, p: float) -> LogitsChainBuilder:
        """Add a TopPFilter (nucleus sampling) with the given p."""
        self._standard[TopPFilter] = TopPFilter(top_p=p)
        return self

    def temperature(self, t: float) -> LogitsChainBuilder:
        """Add a TemperatureScaling processor with the given temperature."""
        self._standard[TemperatureScaling] = TemperatureScaling(temperature=t)
        return self

    def custom(self, processor: LogitsProcessor) -> LogitsChainBuilder:
        """Append a user-supplied LogitsProcessor after the canonical
        group. Useful for logit-bias / forbidden-token / domain-specific
        adjustments. Multiple `.custom(...)` calls append in order."""
        self._custom.append(processor)
        return self

    def build(self) -> LogitsChain:
        """Construct the LogitsChain with processors in canonical order.

        Standard processors that were chained are emitted in the
        fixed `_CANONICAL_ORDER`; custom processors come after, in
        the order they were added.
        """
        ordered: list[LogitsProcessor] = []
        for proc_type in _CANONICAL_ORDER:
            if proc_type in self._standard:
                ordered.append(self._standard[proc_type])
        ordered.extend(self._custom)
        return LogitsChain(processors=ordered)

build() -> LogitsChain

Construct the LogitsChain with processors in canonical order.

Standard processors that were chained are emitted in the fixed _CANONICAL_ORDER; custom processors come after, in the order they were added.

Source code in nnx/generation/logits_chain.py

def build(self) -> LogitsChain:
    """Construct the LogitsChain with processors in canonical order.

    Standard processors that were chained are emitted in the
    fixed `_CANONICAL_ORDER`; custom processors come after, in
    the order they were added.
    """
    ordered: list[LogitsProcessor] = []
    for proc_type in _CANONICAL_ORDER:
        if proc_type in self._standard:
            ordered.append(self._standard[proc_type])
    ordered.extend(self._custom)
    return LogitsChain(processors=ordered)

custom(processor: LogitsProcessor) -> LogitsChainBuilder

Append a user-supplied LogitsProcessor after the canonical group. Useful for logit-bias / forbidden-token / domain-specific adjustments. Multiple .custom(...) calls append in order.

Source code in nnx/generation/logits_chain.py

def custom(self, processor: LogitsProcessor) -> LogitsChainBuilder:
    """Append a user-supplied LogitsProcessor after the canonical
    group. Useful for logit-bias / forbidden-token / domain-specific
    adjustments. Multiple `.custom(...)` calls append in order."""
    self._custom.append(processor)
    return self

repetition_penalty(penalty: float) -> LogitsChainBuilder

Add a RepetitionPenalty processor with the given penalty.

Source code in nnx/generation/logits_chain.py

def repetition_penalty(self, penalty: float) -> LogitsChainBuilder:
    """Add a RepetitionPenalty processor with the given penalty."""
    self._standard[RepetitionPenalty] = RepetitionPenalty(penalty=penalty)
    return self

temperature(t: float) -> LogitsChainBuilder

Add a TemperatureScaling processor with the given temperature.

Source code in nnx/generation/logits_chain.py

def temperature(self, t: float) -> LogitsChainBuilder:
    """Add a TemperatureScaling processor with the given temperature."""
    self._standard[TemperatureScaling] = TemperatureScaling(temperature=t)
    return self

top_k(k: int) -> LogitsChainBuilder

Add a TopKFilter with the given k.

Source code in nnx/generation/logits_chain.py

def top_k(self, k: int) -> LogitsChainBuilder:
    """Add a TopKFilter with the given k."""
    self._standard[TopKFilter] = TopKFilter(top_k=k)
    return self

top_p(p: float) -> LogitsChainBuilder

Add a TopPFilter (nucleus sampling) with the given p.

Source code in nnx/generation/logits_chain.py

def top_p(self, p: float) -> LogitsChainBuilder:
    """Add a TopPFilter (nucleus sampling) with the given p."""
    self._standard[TopPFilter] = TopPFilter(top_p=p)
    return self

nnx.generation.TemperatureScaling

Divide logits by temperature before sampling.

temperature == 0 is a special case: the chain reduces to greedy decoding (argmax). We map argmax positions to +inf and others to -inf so the downstream sampler picks deterministically without branching on the temperature value.

Source code in nnx/generation/logits_processors.py

class TemperatureScaling:
    """Divide logits by ``temperature`` before sampling.

    ``temperature == 0`` is a special case: the chain reduces to greedy
    decoding (argmax). We map argmax positions to +inf and others to
    -inf so the downstream sampler picks deterministically without
    branching on the temperature value.
    """

    def __init__(self, temperature: float):
        if temperature < 0:
            raise ValueError(f"temperature must be >= 0, got {temperature}")
        self.temperature = temperature

    def __call__(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor:
        if self.temperature == 0.0:
            out = torch.full_like(logits, float("-inf"))
            argmax = logits.argmax(dim=-1, keepdim=True)
            out.scatter_(dim=-1, index=argmax, value=float("inf"))
            return out
        return logits / self.temperature

nnx.generation.TopKFilter

Keep only the top-k logits per row; set the rest to -inf.

-inf survives the temperature divide (still -inf) and gets mapped to 0 probability mass by softmax, so the order top-k → temperature or temperature → top-k both work; we don't enforce an ordering.

Source code in nnx/generation/logits_processors.py

class TopKFilter:
    """Keep only the top-k logits per row; set the rest to -inf.

    -inf survives the temperature divide (still -inf) and gets mapped
    to 0 probability mass by softmax, so the order top-k → temperature
    or temperature → top-k both work; we don't enforce an ordering.
    """

    def __init__(self, top_k: int):
        if top_k <= 0:
            raise ValueError(f"top_k must be > 0, got {top_k}")
        self.top_k = top_k

    def __call__(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor:
        k = min(self.top_k, logits.size(-1))
        topk_vals, _ = torch.topk(logits, k=k, dim=-1)
        # threshold[i] = smallest of the top-k values in row i.
        threshold = topk_vals[..., -1:].expand_as(logits)
        return torch.where(logits >= threshold, logits, torch.full_like(logits, float("-inf")))

nnx.generation.TopPFilter

Nucleus (top-p) sampling: keep the smallest set of tokens whose cumulative probability exceeds top_p.

Edge case: if a single token already has probability >= top_p, only that token is retained.

Source code in nnx/generation/logits_processors.py

class TopPFilter:
    """Nucleus (top-p) sampling: keep the smallest set of tokens whose
    cumulative probability exceeds ``top_p``.

    Edge case: if a single token already has probability >= top_p, only
    that token is retained.
    """

    def __init__(self, top_p: float):
        if not (0.0 < top_p <= 1.0):
            raise ValueError(f"top_p must be in (0, 1], got {top_p}")
        self.top_p = top_p

    def __call__(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor:
        sorted_logits, sorted_idx = torch.sort(logits, dim=-1, descending=True)
        cum_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
        # Tokens to remove: cum_probs strictly greater than top_p,
        # shifted by 1 so the token that pushes cum_probs over the
        # threshold is itself kept.
        sorted_remove = cum_probs > self.top_p
        sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
        sorted_remove[..., 0] = False  # always keep the top token
        # Scatter back to the original position order.
        remove = torch.zeros_like(sorted_remove)
        remove.scatter_(dim=-1, index=sorted_idx, src=sorted_remove)
        return logits.masked_fill(remove, float("-inf"))

nnx.generation.RepetitionPenalty

Penalize previously-seen tokens (HF-style).

For each token id i in token_history: * if logits[..., i] > 0: divide by penalty (decreases mass). * if logits[..., i] < 0: multiply by penalty (increases magnitude → further decreases relative mass after softmax).

A penalty of 1.0 is a no-op (the back-compat default).

Source code in nnx/generation/logits_processors.py

class RepetitionPenalty:
    """Penalize previously-seen tokens (HF-style).

    For each token id ``i`` in ``token_history``:
      * if ``logits[..., i] > 0``: divide by penalty (decreases mass).
      * if ``logits[..., i] < 0``: multiply by penalty (increases
        magnitude → further decreases relative mass after softmax).

    A penalty of 1.0 is a no-op (the back-compat default).
    """

    def __init__(self, penalty: float):
        if penalty < 1.0:
            raise ValueError(f"repetition_penalty must be >= 1.0, got {penalty}")
        self.penalty = penalty

    def __call__(self, logits: torch.Tensor, token_history: list[int]) -> torch.Tensor:
        if not token_history or self.penalty == 1.0:
            return logits
        out = logits.clone()
        # We assume batch size 1 throughout GenerativeNNModel.generate;
        # multi-batch generation would need a per-row history.
        idx = torch.tensor(sorted(set(token_history)), dtype=torch.long, device=logits.device)
        # Apply penalty per-token: divide positive, multiply negative.
        selected = out[..., idx]
        positive = selected > 0
        adjusted = torch.where(positive, selected / self.penalty, selected * self.penalty)
        out[..., idx] = adjusted
        return out

nnx.generation.apply_chain(logits: torch.Tensor, *, token_history: list[int], processors: list[LogitsProcessor]) -> torch.Tensor

Apply every processor in order. No-op when processors is empty.

Source code in nnx/generation/logits_processors.py

def apply_chain(
    logits: torch.Tensor,
    *,
    token_history: list[int],
    processors: list[LogitsProcessor],
) -> torch.Tensor:
    """Apply every processor in order. No-op when ``processors`` is empty."""
    for proc in processors:
        logits = proc(logits, token_history)
    return logits

nnx.generation.sample_next_token(logits: torch.Tensor, *, generator: Optional[torch.Generator] = None) -> int

Draw one token id from softmax(logits).

Parameters:

Name	Type	Description	Default
logits	Tensor	shape (1, vocab) — single-sequence sample (the LM path's batch-1 generate scope).	required
generator	Optional[Generator]	optional torch.Generator for reproducible seeded sampling. When None, sampling uses the default RNG (still affected by torch.manual_seed at the call site).	None

Returns:

Type	Description
int	An int token id.

Source code in nnx/generation/sampling.py

def sample_next_token(
    logits: torch.Tensor,
    *,
    generator: Optional[torch.Generator] = None,
) -> int:
    """Draw one token id from ``softmax(logits)``.

    Args:
        logits: shape (1, vocab) — single-sequence sample (the LM
            path's batch-1 generate scope).
        generator: optional torch.Generator for reproducible seeded
            sampling. When None, sampling uses the default RNG (still
            affected by torch.manual_seed at the call site).

    Returns:
        An int token id.
    """
    if logits.dim() != 2 or logits.size(0) != 1:
        raise ValueError(f"sample_next_token expects shape (1, vocab); got {tuple(logits.shape)}")
    # Greedy short-circuit: if any +inf appears, the temperature=0 path
    # set it as the argmax marker — pick the first such position. (We
    # don't try to softmax this: softmax(+inf) is NaN when multiple
    # +infs are present, but TemperatureScaling guarantees exactly one
    # +inf per row.)
    pos_inf_mask = torch.isposinf(logits)
    if pos_inf_mask.any():
        return int(pos_inf_mask[0].nonzero(as_tuple=False)[0].item())
    # All-finite-but-degenerate (all -inf after top-k/top-p collapse):
    # fall back to argmax so we don't crash inside multinomial.
    if torch.isinf(logits).all():
        return int(logits.argmax(dim=-1).item())

    probs = torch.softmax(logits, dim=-1)
    # multinomial requires a finite, positive-sum probability vector.
    # Three degenerate cases collapse to argmax on the original logits:
    #   * NaN-in-logits → softmax produces NaN → sum is NaN (not 0.0,
    #     so the original `== 0.0` check missed this — multinomial then
    #     crashed with "probability tensor contains either inf, nan").
    #   * sum is 0.0 (e.g., the entire row underflowed to zero in
    #     reduced-precision contexts).
    #   * sum is +/-inf (shouldn't happen given the +inf guard above,
    #     but cheap belt-and-braces).
    total = probs.sum()
    if not torch.isfinite(total) or total.item() <= 0.0:
        return int(logits.argmax(dim=-1).item())
    next_id = torch.multinomial(probs, num_samples=1, generator=generator)
    return int(next_id.item())

19. Visualization

19.1. Run-output viz (nnx.vis_utils)

nnx.vis_utils

Run-output visualization helpers.

VisUtils collects the Plotly-based visualizations for run outputs — the artifacts produced after NNModel.train() has completed: training curves, confusion matrices, classification reports, t-SNE projections of held-out logits, etc. It is the sibling of nnx.viz (model-internals visualization — weight histograms, activation maps, gradient flow, Netron export); the two subpackages are deliberately independent and do not share code.

Every method is a @staticmethod returning either a plotly.graph_objects.Figure (for plots) or a pandas.DataFrame (for tables). The class itself carries layout constants (TITLE_SIZE, LABEL_SIZE, FIG_SIZE, MARGIN_SIZE) shared across methods, plus an opt-in RENDERER override for environments where Plotly's default renderer doesn't work (e.g., when serving from a headless container).

Convenience module-level aliases re-export the most common methods at the bottom of this file so callers can write nnx.vis_utils.confusion_matrix(...) instead of nnx.vis_utils.VisUtils.confusion_matrix(...).

confusion_matrix = VisUtils.confusion_matrix module-attribute

classification_report = VisUtils.classification_report module-attribute

multi_line_plot = VisUtils.multi_line_plot module-attribute

scatter_plot = VisUtils.scatter_plot module-attribute

two_dim_tsne_checkpoint_logits = VisUtils.two_dim_tsne_checkpoint_logits module-attribute

19.2. Model-internals viz (nnx.viz)

Opt-in via pip install "thekaveh-nnx[viz]" (pulls torchinfo + captum) and pip install "thekaveh-nnx[viz-interactive]" (adds the netron browser viewer for nnx.viz.netron_export(..., launch=True)).

nnx.viz.activation.activation_map(model: Union[nn.Module, NNModel], x: torch.Tensor, layer_name: str, *, max_channels: int = 16, cols: int = 4, fig_width: int = 900, cell_size: int = 180) -> go.Figure

Capture the activation of layer_name for input x and render it.

Registers a forward hook on the named submodule, runs model(x) under torch.no_grad(), then removes the hook and turns the captured tensor into a Plotly heatmap layout:

• 4D (N, C, H, W) activations: grid of up to max_channels per-channel heatmaps from the first sample (N=0).
• 2D (N, F) activations: single (N, F) heatmap.
• Other ranks: flattened single-row heatmap (best-effort fallback).

Parameters:

Name	Type	Description	Default
model	Union[Module, NNModel]	An NNModel (unwrapped to .net) or any torch.nn.Module.	required
x	Tensor	Input tensor (or any object) accepted by model.__call__. Moved to the same device as the model's first parameter when possible.	required
layer_name	str	Dotted name from model.named_modules() — e.g. "layers.2" for a Sequential, "conv1" for a class attribute. Pass an empty string "" to hook the top-level module itself.	required
max_channels	int	Cap on conv-channel subplots (4D case). Defaults to 16 — enough to spot patterns without crushing the layout for 256-channel feature maps.	16
cols	int	Subplot columns in the 4D grid.	4
fig_width	int	Total figure width in pixels.	900
cell_size	int	Per-subplot square cell size (px). Total height scales with the row count.	180

Returns:

Type	Description
Figure	A Plotly Figure containing one or more Heatmap traces.

Raises:

Type	Description
ValueError	If layer_name doesn't resolve to a submodule of model.
RuntimeError	If the forward hook on layer_name never fires (the layer is not reached by this input's forward path).

Source code in nnx/viz/activation.py

def activation_map(
    model: Union[nn.Module, NNModel],
    x: torch.Tensor,
    layer_name: str,
    *,
    max_channels: int = 16,
    cols: int = 4,
    fig_width: int = 900,
    cell_size: int = 180,
) -> go.Figure:
    """Capture the activation of `layer_name` for input `x` and render it.

    Registers a forward hook on the named submodule, runs `model(x)` under
    `torch.no_grad()`, then removes the hook and turns the captured tensor
    into a Plotly heatmap layout:

    - 4D ``(N, C, H, W)`` activations: grid of up to `max_channels` per-channel
      heatmaps from the first sample (``N=0``).
    - 2D ``(N, F)`` activations: single ``(N, F)`` heatmap.
    - Other ranks: flattened single-row heatmap (best-effort fallback).

    Args:
        model: An `NNModel` (unwrapped to `.net`) or any `torch.nn.Module`.
        x: Input tensor (or any object) accepted by `model.__call__`. Moved to
            the same device as the model's first parameter when possible.
        layer_name: Dotted name from `model.named_modules()` — e.g. `"layers.2"`
            for a Sequential, `"conv1"` for a class attribute. Pass an empty
            string `""` to hook the top-level module itself.
        max_channels: Cap on conv-channel subplots (4D case). Defaults to 16 —
            enough to spot patterns without crushing the layout for 256-channel
            feature maps.
        cols: Subplot columns in the 4D grid.
        fig_width: Total figure width in pixels.
        cell_size: Per-subplot square cell size (px). Total height scales
            with the row count.

    Returns:
        A Plotly `Figure` containing one or more `Heatmap` traces.

    Raises:
        ValueError: If `layer_name` doesn't resolve to a submodule of `model`.
        RuntimeError: If the forward hook on `layer_name` never fires
            (the layer is not reached by this input's forward path).
    """
    # Local import to avoid a circular import at package init time.
    from ..nn.nn_model import NNModel

    if isinstance(model, NNModel):
        model = model.net

    modules = dict(model.named_modules())
    if layer_name not in modules:
        # Surface the closest candidates so the caller can fix the typo
        # without re-grep'ing the model. Limit to a short list to keep
        # the error message readable on deep models.
        available = list(modules.keys())
        sample = ", ".join(repr(n) for n in available[:10])
        raise ValueError(
            f"activation_map: layer_name={layer_name!r} not found in model.named_modules(). "
            f"First {min(10, len(available))} available names: [{sample}]"
        )

    target = modules[layer_name]
    captured: dict[str, torch.Tensor] = {}

    def _hook(_mod: nn.Module, _inp: tuple, out: torch.Tensor) -> None:
        # Some layers return tuples / dicts (multi-head attention, etc.).
        # We grab the first tensor we can find so the helper degrades
        # gracefully on non-conv layers rather than raising.
        if isinstance(out, torch.Tensor):
            captured["v"] = out.detach().cpu()
        elif isinstance(out, (tuple, list)) and out and isinstance(out[0], torch.Tensor):
            captured["v"] = out[0].detach().cpu()
        else:
            # Last-ditch: try to coerce.
            captured["v"] = torch.as_tensor(out).detach().cpu()

    handle = target.register_forward_hook(_hook)
    try:
        # Best-effort device alignment — works for the common case where
        # the model has at least one parameter. Pure stateless modules
        # (e.g. nn.ReLU) keep `x` on whatever device the caller passed.
        try:
            device = next(model.parameters()).device
            if isinstance(x, torch.Tensor):
                x = x.to(device)
        except StopIteration:
            pass
        was_training = model.training
        model.eval()
        try:
            with torch.no_grad():
                model(x)
        finally:
            if was_training:
                model.train()
    finally:
        handle.remove()

    if "v" not in captured:
        # The hook *should* always fire (PyTorch invokes it as part of
        # the forward pass). If it didn't, the forward call must have
        # short-circuited before reaching the layer — surface that
        # rather than KeyError'ing on the next line.
        raise RuntimeError(
            f"activation_map: forward hook on {layer_name!r} did not fire. "
            "The layer may not be reached by this input — check the model's forward path."
        )
    return _activation_to_figure(
        captured["v"],
        layer_name=layer_name,
        max_channels=max_channels,
        cols=cols,
        fig_width=fig_width,
        cell_size=cell_size,
    )

nnx.viz.attribute.attribute(model: Union[nn.Module, NNModel], x: torch.Tensor, *, method: str = 'integrated_gradients', target: Any = None, **method_kwargs: Any) -> tuple[torch.Tensor, go.Figure]

Compute input attributions via Captum and render a Plotly heatmap.

Parameters:

Name	Type	Description	Default
model	Union[Module, NNModel]	An NNModel (unwrapped to .net) or any torch.nn.Module. The model is set to eval() for the duration of the attribution call; the original mode is restored on return.	required
x	Tensor	Input tensor to attribute. Shape (B, ...). Gradient-based methods will set requires_grad_(True) internally as needed.	required
method	str	One of "integrated_gradients", "gradient_shap", "deep_lift", "saliency", "input_x_gradient", "occlusion".	'integrated_gradients'
target	Any	Target class index (or per-batch indices) for classification attributors. Forwarded verbatim to Captum's .attribute(target=).	None
**method_kwargs	Any	Extra kwargs forwarded to the per-method .attribute(...) call. Overrides any defaults supplied for gradient_shap (baselines) or occlusion (sliding_window_shapes).	{}

Returns:

Type	Description
tuple[Tensor, Figure]	A tuple (attribution_tensor, figure) where attribution_tensor is a torch.Tensor with the same shape as x (per Captum's standard return contract for these six methods) and figure is a Plotly Heatmap visualizing the attribution. Image-shaped inputs (3-D / 4-D) are mean-pooled over channels before rendering.

Raises:

Type	Description
ImportError	If captum is not installed. Install via pip install thekaveh-nnx[viz] or pip install captum>=0.7.0.
ValueError	If method is not one of the supported keys.

Source code in nnx/viz/attribute.py

def attribute(
    model: Union[nn.Module, NNModel],
    x: torch.Tensor,
    *,
    method: str = "integrated_gradients",
    target: Any = None,
    **method_kwargs: Any,
) -> tuple[torch.Tensor, go.Figure]:
    """Compute input attributions via Captum and render a Plotly heatmap.

    Args:
        model: An `NNModel` (unwrapped to `.net`) or any `torch.nn.Module`.
            The model is set to `eval()` for the duration of the attribution
            call; the original mode is restored on return.
        x: Input tensor to attribute. Shape `(B, ...)`. Gradient-based
            methods will set `requires_grad_(True)` internally as needed.
        method: One of `"integrated_gradients"`, `"gradient_shap"`,
            `"deep_lift"`, `"saliency"`, `"input_x_gradient"`, `"occlusion"`.
        target: Target class index (or per-batch indices) for classification
            attributors. Forwarded verbatim to Captum's `.attribute(target=)`.
        **method_kwargs: Extra kwargs forwarded to the per-method
            `.attribute(...)` call. Overrides any defaults supplied for
            `gradient_shap` (`baselines`) or `occlusion` (`sliding_window_shapes`).

    Returns:
        A tuple `(attribution_tensor, figure)` where `attribution_tensor` is a
            `torch.Tensor` with the same shape as `x` (per Captum's standard
            return contract for these six methods) and `figure` is a Plotly
            `Heatmap` visualizing the attribution. Image-shaped inputs (3-D /
            4-D) are mean-pooled over channels before rendering.

    Raises:
        ImportError: If `captum` is not installed. Install via
            `pip install thekaveh-nnx[viz]` or `pip install captum>=0.7.0`.
        ValueError: If `method` is not one of the supported keys.
    """
    if method not in SUPPORTED_METHODS:
        raise ValueError(f"unknown method '{method}'; choose from {list(SUPPORTED_METHODS)}")

    # Local import to avoid a circular import at package init time
    # (NNModel pulls in nnx.viz indirectly through some training paths).
    from ..nn.nn_model import NNModel

    if isinstance(model, NNModel):
        model = model.net

    was_training = model.training
    model.eval()
    try:
        attributor = _build_attributor(method, model)
        kwargs = {**_default_kwargs(method, x), **method_kwargs}
        attr_tensor = attributor.attribute(x, target=target, **kwargs)
    finally:
        if was_training:
            model.train()

    figure = _attribution_to_figure(attr_tensor, x)
    return attr_tensor, figure

nnx.viz.gradient_flow.gradient_flow(model: Union[nn.Module, NNModel]) -> go.Figure

Return a Plotly bar chart of per-parameter L2 gradient norms.

Call AFTER loss.backward() and BEFORE optimizer.zero_grad(). Each bar is one trainable nn.Parameter of the model whose .grad has been populated by the backward pass; bar height is the L2 norm of that gradient.

Frozen parameters (requires_grad=False) are skipped. Parameters whose gradient is None (typically because they weren't reached during the forward pass) are also skipped.

Parameters:

Name	Type	Description	Default
model	Union[Module, NNModel]	an NNModel (unwrapped to its .net) or any nn.Module whose gradients have just been populated by loss.backward().	required

Returns:

Type	Description
Figure	A Plotly Figure with one bar per trainable parameter,
Figure	labeled by named_parameters() dotted name.

Raises:

Type	Description
ValueError	if no parameter has a populated gradient — most often because loss.backward() wasn't called before this function.

Source code in nnx/viz/gradient_flow.py

def gradient_flow(model: Union[nn.Module, NNModel]) -> go.Figure:
    """Return a Plotly bar chart of per-parameter L2 gradient norms.

    Call AFTER ``loss.backward()`` and BEFORE ``optimizer.zero_grad()``.
    Each bar is one trainable ``nn.Parameter`` of the model whose
    ``.grad`` has been populated by the backward pass; bar height is
    the L2 norm of that gradient.

    Frozen parameters (``requires_grad=False``) are skipped. Parameters
    whose gradient is ``None`` (typically because they weren't reached
    during the forward pass) are also skipped.

    Args:
        model: an ``NNModel`` (unwrapped to its ``.net``) or any
            ``nn.Module`` whose gradients have just been populated by
            ``loss.backward()``.

    Returns:
        A Plotly ``Figure`` with one bar per trainable parameter,
        labeled by ``named_parameters()`` dotted name.

    Raises:
        ValueError: if no parameter has a populated gradient — most
            often because ``loss.backward()`` wasn't called before
            this function.
    """
    from ..nn.nn_model import NNModel

    if isinstance(model, NNModel):
        model = model.net

    names: list[str] = []
    norms: list[float] = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.grad is None:
            continue
        names.append(name)
        norms.append(param.grad.detach().norm().item())

    if not names:
        raise ValueError(
            "gradient_flow: no parameters have populated gradients. "
            "Did you forget to call loss.backward() before this call?"
        )

    fig = go.Figure(data=[go.Bar(x=names, y=norms)])
    fig.update_layout(
        title="Per-layer gradient norms",
        xaxis_title="Parameter",
        yaxis_title="L2 norm of gradient",
        xaxis_tickangle=-45,
        margin=dict(b=120),
    )
    return fig

nnx.viz.netron.netron_export(model: Union[nn.Module, NNModel], path: str, example_input: Union[torch.Tensor, tuple, np.ndarray], *, launch: bool = False, opset_version: int = 17, dynamic_batch: bool = True) -> str

Export model to an ONNX file at path (optionally open Netron).

Parameters:

Name	Type	Description	Default
model	Union[Module, NNModel]	An NNModel (unwrapped to .net) or any torch.nn.Module.	required
path	str	Output filename, e.g. "model.onnx".	required
example_input	Union[Tensor, tuple, ndarray]	A tensor (or tuple of tensors) with realistic shape / dtype used to trace the network.	required
launch	bool	When True, call netron.start(path) to open the model in Netron's browser viewer. Requires pip install thekaveh-nnx[viz-interactive] (or pip install netron). Defaults to False so CI / tests can exercise export without spawning a long-lived process.	False
opset_version	int	ONNX opset to target. 17 is broadly supported by current runtimes.	17
dynamic_batch	bool	When True (default), marks dim 0 as dynamic so the exported graph accepts any batch size at inference.	True

Returns:

Type	Description
str	The path written (matches path — handy when chaining).

Raises:

Type	Description
ImportError	When launch=True and the netron package isn't installed. The ONNX export itself uses torch.onnx, which is part of core PyTorch.

Source code in nnx/viz/netron.py

def netron_export(
    model: Union[nn.Module, NNModel],
    path: str,
    example_input: Union[torch.Tensor, tuple, np.ndarray],
    *,
    launch: bool = False,
    opset_version: int = 17,
    dynamic_batch: bool = True,
) -> str:
    """Export `model` to an ONNX file at `path` (optionally open Netron).

    Args:
        model: An `NNModel` (unwrapped to `.net`) or any `torch.nn.Module`.
        path: Output filename, e.g. ``"model.onnx"``.
        example_input: A tensor (or tuple of tensors) with realistic
            shape / dtype used to trace the network.
        launch: When True, call `netron.start(path)` to open the model
            in Netron's browser viewer. Requires `pip install thekaveh-nnx[viz-interactive]`
            (or `pip install netron`). Defaults to False so CI / tests
            can exercise export without spawning a long-lived process.
        opset_version: ONNX opset to target. 17 is broadly supported
            by current runtimes.
        dynamic_batch: When True (default), marks dim 0 as dynamic so
            the exported graph accepts any batch size at inference.

    Returns:
        The path written (matches `path` — handy when chaining).

    Raises:
        ImportError: When `launch=True` and the `netron` package isn't
            installed. The ONNX export itself uses `torch.onnx`, which
            is part of core PyTorch.
    """
    # Local import to avoid a circular import at package init time.
    from ..nn.nn_model import NNModel

    if isinstance(model, NNModel):
        net = model.net
        device = model.device
    else:
        net = model
        # Best-effort device probe — fall back to CPU for stateless
        # modules with no parameters.
        try:
            device = next(net.parameters()).device
        except StopIteration:
            device = torch.device("cpu")

    if isinstance(example_input, torch.Tensor):
        example_input = (example_input.to(device),)
    elif isinstance(example_input, np.ndarray):
        example_input = (torch.from_numpy(example_input).to(device),)
    else:
        example_input = tuple(
            (e.to(device) if isinstance(e, torch.Tensor) else torch.from_numpy(np.asarray(e)).to(device))
            for e in example_input
        )

    in_names = [f"input_{i}" for i in range(len(example_input))]
    out_names = ["output"]
    dynamic_axes = {n: {0: "batch"} for n in in_names + out_names} if dynamic_batch else None

    was_training = net.training
    net.eval()
    try:
        # Mirror NNModel.to_onnx — pin to the legacy TorchScript path so
        # plain `pip install onnx` is enough; the dynamo path needs
        # `onnxscript` which we keep off the viz dep set.
        export_accepts_dynamo = "dynamo" in inspect.signature(torch.onnx.export).parameters
        if export_accepts_dynamo:
            torch.onnx.export(
                net,
                example_input,
                path,
                input_names=in_names,
                output_names=out_names,
                dynamic_axes=dynamic_axes,
                opset_version=opset_version,
                dynamo=False,
            )
        else:
            torch.onnx.export(
                net,
                example_input,
                path,
                input_names=in_names,
                output_names=out_names,
                dynamic_axes=dynamic_axes,
                opset_version=opset_version,
            )
    finally:
        if was_training:
            net.train()

    if launch:
        # Lazy import — keeps `netron` out of required deps. Only users
        # who actually want the interactive viewer pay the install cost.
        try:
            import netron  # type: ignore[import-not-found]
        except ImportError as e:
            raise ImportError(
                "netron_export(launch=True) requires the `netron` package. "
                "Install via `pip install thekaveh-nnx[viz-interactive]` (or `pip install netron`)."
            ) from e
        netron.start(path)

    return path

nnx.viz.summary.summary(model: Union[nn.Module, NNModel], *, input_size: tuple[int, ...] | None = None, input_data: Union[torch.Tensor, tuple, list, None] = None, depth: int = 4, col_names: tuple[str, ...] = ('output_size', 'num_params', 'mult_adds')) -> ModelStatistics

Return a torchinfo.ModelStatistics summary for model.

Parameters:

Name	Type	Description	Default
model	Union[Module, NNModel]	An NNModel (unwrapped to .net) or any torch.nn.Module.	required
input_size	tuple[int, ...] | None	Shape tuple for a synthetic dummy input, e.g. (1, 3, 224, 224). Mutually exclusive with input_data.	None
input_data	Union[Tensor, tuple, list, None]	An actual tensor / tuple / list to forward through the model. Useful when the model takes multiple positional arguments or a non-tensor input (graphs, dicts) that input_size can't describe.	None
depth	int	Maximum module-nesting depth to expand in the table.	4
col_names	tuple[str, ...]	Which torchinfo columns to include. Defaults to the three most useful ones for spotting parameter / FLOP regressions across runs.	('output_size', 'num_params', 'mult_adds')

Returns:

Type	Description
ModelStatistics	The torchinfo.ModelStatistics instance — print it for the Keras-style
ModelStatistics	table, or access .total_params / .trainable_params / .total_mult_adds
ModelStatistics	for programmatic regression assertions.

Raises:

Type	Description
ImportError	If torchinfo isn't installed. Install with pip install thekaveh-nnx[viz].

Source code in nnx/viz/summary.py

def summary(
    model: Union[nn.Module, NNModel],
    *,
    input_size: tuple[int, ...] | None = None,
    input_data: Union[torch.Tensor, tuple, list, None] = None,
    depth: int = 4,
    col_names: tuple[str, ...] = ("output_size", "num_params", "mult_adds"),
) -> ModelStatistics:
    """Return a `torchinfo.ModelStatistics` summary for `model`.

    Args:
        model: An `NNModel` (unwrapped to `.net`) or any `torch.nn.Module`.
        input_size: Shape tuple for a synthetic dummy input, e.g. `(1, 3, 224, 224)`.
            Mutually exclusive with `input_data`.
        input_data: An actual tensor / tuple / list to forward through the model.
            Useful when the model takes multiple positional arguments or a non-tensor
            input (graphs, dicts) that `input_size` can't describe.
        depth: Maximum module-nesting depth to expand in the table.
        col_names: Which torchinfo columns to include. Defaults to the three most
            useful ones for spotting parameter / FLOP regressions across runs.

    Returns:
        The `torchinfo.ModelStatistics` instance — print it for the Keras-style
        table, or access `.total_params` / `.trainable_params` / `.total_mult_adds`
        for programmatic regression assertions.

    Raises:
        ImportError: If `torchinfo` isn't installed. Install with `pip install thekaveh-nnx[viz]`.
    """
    try:
        from torchinfo import summary as _ti_summary
    except ImportError as e:
        raise ImportError(
            "nnx.viz.summary requires torchinfo — install via `pip install thekaveh-nnx[viz]` or `pip install torchinfo>=1.8.0`."
        ) from e
    # Accept NNModel directly — unwrap to .net so callers don't have to.
    # Local import to avoid a circular import at package init time.
    from ..nn.nn_model import NNModel

    if isinstance(model, NNModel):
        model = model.net

    if input_size is None:
        return _ti_summary(
            model,
            input_size=input_size,
            input_data=input_data,
            depth=depth,
            col_names=col_names,
            verbose=0,
        )

    # torchinfo synthesizes the input_size= dummy via torch.rand — an
    # incidental draw whose values never affect the statistics, but
    # which silently shifts any seeded pipeline probing a summary
    # mid-run (the concepts.md idiom). Snapshot/restore RNG around it.
    # Device mirrors torchinfo's own inference: the model's parameter
    # device, else CUDA when available.
    param = next(model.parameters(), None)
    if param is not None:
        dev = param.device
    elif torch.cuda.is_available():
        dev = torch.device("cuda")
    else:
        dev = torch.device("cpu")
    cpu_rng_state = torch.get_rng_state()
    device_rng_state = None
    if dev.type == "cuda":
        device_rng_state = torch.cuda.get_rng_state(dev)
    elif dev.type == "mps":
        device_rng_state = torch.mps.get_rng_state()
    try:
        return _ti_summary(
            model,
            input_size=input_size,
            input_data=input_data,
            depth=depth,
            col_names=col_names,
            verbose=0,
        )
    finally:
        torch.set_rng_state(cpu_rng_state)
        if dev.type == "cuda":
            torch.cuda.set_rng_state(device_rng_state, dev)
        elif dev.type == "mps":
            torch.mps.set_rng_state(device_rng_state)

nnx.viz.weight_histogram.weight_histogram(model: Union[nn.Module, NNModel], *, bins: int = 64, cols: int = 3, fig_width: int = 1000, row_height: int = 200) -> go.Figure

Return a Plotly grid of per-parameter weight histograms.

Parameters:

Name	Type	Description	Default
model	Union[Module, NNModel]	An NNModel (unwrapped to .net) or any torch.nn.Module.	required
bins	int	Number of histogram bins per parameter tensor.	64
cols	int	Number of columns in the subplot grid. Rows are computed from the parameter count.	3
fig_width	int	Figure width in pixels.	1000
row_height	int	Per-row height in pixels; total height = row_height * rows.	200

Returns:

Type	Description
Figure	A Plotly Figure with one Histogram trace per named parameter tensor.
Figure	Each subplot title is the dotted parameter name (e.g. layers.0.weight).
Figure	Empty parameter tensors are skipped from the grid.

Raises:

Type	Description
ValueError	If model has no named parameters (nothing to plot).

Source code in nnx/viz/weight_histogram.py

def weight_histogram(
    model: Union[nn.Module, NNModel],
    *,
    bins: int = 64,
    cols: int = 3,
    fig_width: int = 1000,
    row_height: int = 200,
) -> go.Figure:
    """Return a Plotly grid of per-parameter weight histograms.

    Args:
        model: An `NNModel` (unwrapped to `.net`) or any `torch.nn.Module`.
        bins: Number of histogram bins per parameter tensor.
        cols: Number of columns in the subplot grid. Rows are computed from the
            parameter count.
        fig_width: Figure width in pixels.
        row_height: Per-row height in pixels; total height = `row_height * rows`.

    Returns:
        A Plotly `Figure` with one `Histogram` trace per named parameter tensor.
        Each subplot title is the dotted parameter name (e.g. `layers.0.weight`).
        Empty parameter tensors are skipped from the grid.

    Raises:
        ValueError: If `model` has no named parameters (nothing to plot).
    """
    # Local import to avoid a circular import at package init time.
    from ..nn.nn_model import NNModel

    if isinstance(model, NNModel):
        model = model.net
    # numel() > 0: the docstring promises empty tensors are skipped
    # (a zero-size histogram subplot renders as a confusing blank).
    params = [(n, p.detach().cpu().flatten().numpy()) for n, p in model.named_parameters() if p.numel() > 0]
    n = len(params)
    if n == 0:
        raise ValueError("weight_histogram: model has no named parameters to plot.")
    rows = math.ceil(n / cols)
    fig = make_subplots(
        rows=rows,
        cols=cols,
        subplot_titles=[name for name, _ in params],
        vertical_spacing=0.08,
        horizontal_spacing=0.05,
    )
    for idx, (_name, vals) in enumerate(params):
        r, c = idx // cols + 1, idx % cols + 1
        fig.add_trace(
            go.Histogram(x=vals, nbinsx=bins, showlegend=False),
            row=r,
            col=c,
        )
    fig.update_layout(
        width=fig_width,
        height=row_height * rows,
        title="Weight histograms (per parameter tensor)",
        margin=dict(l=20, r=20, t=60, b=20),
    )
    return fig

20. Utilities

nnx.utils

Pretty-printing helpers used throughout nnx.

Both module-level functions (print_tree, print_table, flatten_dict) and the legacy Utils class API are exported. New code should prefer the module functions; Utils.method(...) is kept as a thin back-compat shim so existing notebooks keep working.

print_tree(tree, level: int = 0, *, file=None) -> None

Pretty-print a nested dict as an indented tree.

Pass file= (any object with .write) to redirect output away from stdout — useful for capturing in tests or writing to a log. Defaults to sys.stdout.

Source code in nnx/utils.py

def print_tree(tree, level: int = 0, *, file=None) -> None:
    """Pretty-print a nested dict as an indented tree.

    Pass ``file=`` (any object with ``.write``) to redirect output away
    from stdout — useful for capturing in tests or writing to a log.
    Defaults to ``sys.stdout``.
    """
    out = file if file is not None else sys.stdout
    if not isinstance(tree, dict) or not tree:
        # Empty dict (top-level or nested value): nothing to print, and
        # the max() below would raise on an empty key sequence.
        return

    # str(...): keys may be non-string (int epochs, enums).
    max_key_len = max(len(str(key)) for key in tree.keys())

    for key, val in tree.items():
        if isinstance(val, dict):
            print(" " * level * 4 + f"[-] {key}: ", file=out)
            print_tree(val, level + 1, file=out)
        else:
            print(" " * level * 4 + f"[+] {str(key).ljust(max_key_len)} : {val}", file=out)

print_table(data: dict, header: bool = True, title: Optional[str] = None, *, file=None) -> None

Print data as a 2-column key/value table.

Pass file= to redirect output. Defaults to sys.stdout.

Source code in nnx/utils.py

def print_table(data: dict, header: bool = True, title: Optional[str] = None, *, file=None) -> None:
    """Print ``data`` as a 2-column key/value table.

    Pass ``file=`` to redirect output. Defaults to ``sys.stdout``.
    """
    out = file if file is not None else sys.stdout
    table = PrettyTable(["Key", "Value"])
    table.header = header
    if title is not None:
        table.title = title

    for key, val in data.items():
        table.add_row([key, val])

    print(table, file=out)

flatten_dict(data: dict, parent_key: str = '', sep: str = '.') -> dict

Flatten a nested dict so nested keys become parent.child style.

flatten_dict({"a": 1, "b": {"c": 2}})

Source code in nnx/utils.py

def flatten_dict(data: dict, parent_key: str = "", sep: str = ".") -> dict:
    """Flatten a nested dict so nested keys become ``parent.child`` style.

    >>> flatten_dict({"a": 1, "b": {"c": 2}})
    {'a': 1, 'b.c': 2}
    """
    flattened = []
    for key, val in data.items():
        flattened_key = f"{parent_key}{sep}{key}" if parent_key else key
        if isinstance(val, dict):
            flattened.extend(flatten_dict(val, flattened_key, sep=sep).items())
        else:
            flattened.append((flattened_key, val))
    return dict(flattened)

20.1. Utils back-compat facade

nnx.Utils is a thin staticmethod facade over the module-level functions above, kept so existing notebook code that calls Utils.print_tree(...) / Utils.print_table(...) / Utils.flatten_dict(...) continues to work. New code should prefer the module-level functions directly.