2. Quickstart

An end-to-end CPU example you can paste into a Python REPL. Trains a tiny feed-forward classifier on random data so you can verify the install in under five seconds.

Install NNx first with pip install thekaveh-nnx. This guide describes the current source, and the newest PyPI release can be older. To install the source or choose extras, see the README install section.

1. Minimal example

import torch
from torch.utils.data import DataLoader, TensorDataset

from nnx import (
    NNModel, NNParams, NNModelParams, NNTrainParams,
    NNOptimParams, NNSchedulerParams,
    Activations, Devices, Losses, Nets, Optims,
    EarlyStopping,
)

# 1. Data
X_train, y_train = torch.randn(256, 8), torch.randint(0, 3, (256,))
X_val,   y_val   = torch.randn(64,  8), torch.randint(0, 3, (64,))
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
val_loader   = DataLoader(TensorDataset(X_val,   y_val),   batch_size=32)

# 2. Model
net_params   = NNParams(input_dim=8, output_dim=3, hidden_dims=[32, 16],
                        dropout_prob=0.1, activation=Activations.RELU)
model_params = NNModelParams(net=Nets.FEED_FWD, device=Devices.CPU,
                             loss=Losses.CROSS_ENTROPY)
model = NNModel(net_params=net_params, params=model_params)

# 3. Train
train_params = NNTrainParams(
    n_epochs=10,
    seed=42,                       # reproducibility
    train_loader=train_loader,
    val_loader=val_loader,
    optim=NNOptimParams(name=Optims.ADAM, max_lr=1e-2,
                        momentum=(0.9, 0.999), weight_decay=5e-5,
                        grad_clip_norm=1.0),
    scheduler=NNSchedulerParams(min_lr=1e-7, factor=0.5,
                                patience=3, cooldown=1, threshold=1e-3),
)
run = model.train(params=train_params, callbacks=[EarlyStopping(patience=5)])

# 4. Use it
print(f"trained {len(run.idps)} iterations; saved under runs/{run.id}/")
result = model.predict(X=X_val)
print(f"predicted {len(result.classes)} samples")

Run ids are content-addressed, so running the same script again in the same directory raises FileExistsError. Delete runs/ first, or pass overwrite_existing=True to NNTrainParams.

2. Common variations

2.1. GPU / Apple Silicon

from nnx import Devices
NNModelParams(net=Nets.FEED_FWD, device=Devices.get(), loss=Losses.CROSS_ENTROPY)
# Devices.get() picks MPS > CUDA > CPU.

2.2. Precision: FP32, FP16 and BF16

from nnx import PrecisionPolicy
NNModelParams(..., precision=PrecisionPolicy("bf16"))                   # CPU, or CUDA with bf16 support
NNModelParams(..., precision=PrecisionPolicy("fp16"))                   # CUDA only
NNModelParams(..., precision=PrecisionPolicy("fp16", fallback="fp32"))  # fp16 where it can run, else fp32

NNModelParams.precision declares the precision training, evaluation and prediction run in. The policy is resolved against the model's device when the model is built and again, afresh, before a run is reserved:

mode	runs as	where
"fp32" (default)	full precision — no autocast, no scaler	every device
"fp16"	float16 autocast + a torch.amp.GradScaler	CUDA
"bf16"	bfloat16 autocast, no scaler	CPU; CUDA with native bf16 (compute capability 8.0+ — emulated bf16 does not count)

An explicit PrecisionPolicy("fp32") is the default. It normalizes to no policy, so it keeps the default's run id. A mode the device cannot run raises PrecisionUnsupportedError before any run directory exists. The exception is a policy that names fallback="fp32": the run then trains in full precision and records why.

Autocast wraps only the forward pass and loss. The backward runs outside it, and parameters stay float32 (never model.half()). The update keeps its unscale → normalize → finite-check → clip → step order:

• FP16 unscale_s first, so grad_clip_norm clips true gradients. It steps through the scaler, which skips a non-finite update.
• Scaler-free BF16 checks its gradients itself. It raises FloatingPointError rather than apply a non-finite update.

Evaluation and predict() / predict_proba() run under the same policy and return float32 outputs, so their schemas are unchanged.

run.precision records what happened, and NNRun.load reads it back from metadata.yaml. The record holds these fields:

• requested, effective and fallback_reason.
• TF32, reported separately from torch's fp32_precision setting when made. NNx never sets it.
• covers: the NNx surfaces that apply the policy. These are the run's training where NNx runs the default step or an objective, and evaluate() / predict() whenever they are called.
• not_covered: the surfaces this run's policy did not reach. These are the NNx surfaces outside this run, plus the surfaces that never apply the policy.

Custom train_step_fn / eval_step_fn code is outside both lists. A training step reads ctx.precision and applies it itself. The policy does not reach nnx.lr_finder, nnx.diffusion.sampling or generation.

nnx.precision_support() reports each device's cells with one of four values:

• "verified": CPU FP32 and BF16, by NNx's seeded fixtures within nnx.precision.REFERENCE_TOLERANCES.
• "supported": the hardware is present and the mode resolves, but no NNx reference fixture checks it on that device.
• "unverified": the hardware is absent here, so nothing is claimed.
• "unsupported": the device type cannot run the mode.

The scaler's state dictionary and the precision record travel with every checkpoint's training-state sidecar. A warm resume therefore continues from the saved scale factor and growth tracker. A stateful resume into a different effective precision is rejected (resume precision mismatch) before the model is touched; a weights-only warm start (resume_mode="weights_only") may switch:

from nnx import Checkpoints, NNCheckpoint

bf16 = NNModelParams(net=Nets.FEED_FWD, device=Devices.CPU, loss=Losses.CROSS_ENTROPY, precision=PrecisionPolicy("bf16"))
first = NNModel(net_params=net_params, params=bf16).train(params=NNTrainParams(n_epochs=1, train_loader=train_loader))
print(first.precision.record()["effective"])   # "bf16"
state = NNCheckpoint.load_training_state(run=first.id, type=Checkpoints.LAST)
print(state["precision"]["effective"], state["scaler"])   # "bf16", None (fp16 keeps the scaler's state)
resumed = NNModel(net_params=net_params, params=bf16).train(
    params=NNTrainParams(n_epochs=1, train_loader=train_loader, resume_from_run_id=first.id),
)

NNModel.from_checkpoint(..., device=...) and Hub loading re-resolve the policy on the destination device, never from saved metadata. Both take precision= to replace it (say, a CUDA FP16 model loaded on a CPU).

The built-in imperative paradigm steps (finalize_step) run in full precision. They refuse a reduced policy before any work is done. A wrapper keeps that refusal when it is written with functools.wraps, which carries the marker, or marked with nnx.precision.full_precision_only. Otherwise finalize_step refuses at the first step. A custom train_step_fn applies ctx.precision.autocast() and ctx.scaler itself. In Trainer, an objective runs in the policy, while step functions (which own every update) refuse a reduced one.

The legacy NNModelParams(mixed_precision=True) keeps its meaning and run id. It means FP16 autocast with a scaler in NNModel.train on CUDA, silently plain FP32 elsewhere, and no effect on evaluation, prediction or Trainer. Combining it with a policy other than fp16 is refused.

A stateful resume can load a checkpoint written before precision was recorded. It counts that checkpoint as FP16 exactly when the checkpoint saved a scaler. resume GradScaler presence mismatch still names a scaler that appears or disappears.

The scaler factory is why NNx declares the PyTorch floor it does. The support matrix lists the tested torch / torchvision / Python combinations and which of them carry real CUDA evidence. Two examples exercise the policy:

• examples/01_synthetic_classification.py's precision_workflow(mode) runs one bounded CPU workflow: FP32 by default, or BF16 (--precision bf16) where precision_support allows it.
• examples/02_resume_training.py's amp_resume_compatibility() runs the CPU no-scaler cycle everywhere. It runs the enabled-AMP cycle only on a CUDA host.

2.3. Warm-resume training

# Train round 1
run = model.train(params=NNTrainParams(n_epochs=10, ...))

# Train round 2 — pick up from where the last run's LAST checkpoint left off.
model.train(params=NNTrainParams(
    n_epochs=10,
    resume_from_run_id=run.id,
    resume_from_checkpoint="last",   # or "best"
    ...
))

Resume validates the checkpoint's optimizer topology and matching immutable training-state generation before applying it. It restores scheduler/scaler, completed epoch, loader generators, and Python/NumPy/PyTorch CPU/CUDA/MPS RNG state; use num_workers=0 when exact continuation matters. Any re-iterable batch source works: a DataLoader, a list of (X, Y) batches, or the graph full-batch list. The worker warning only appears for a real loader with num_workers > 0.

Callback and step state resumes too. NNx writes the following state into every checkpoint:

• EarlyStopping's best value and patience counter.
• The I-JEPA target encoder, and the jepa_objective / diffusion_objective state (EMA target and counter, generator).
• Any object registered through train(..., components=[...]).

This state is restored after the callbacks' reset hooks, so a split run stops at the same epoch as the uninterrupted one (Concepts §14.1).

resume_mode="stateful" fails before restoring anything when the checkpoint is weights-only. resume_mode="weights_only" warm-starts from the weights alone, with a fresh optimizer, scheduler and components. The returned run.resume_status says which happened. The multi-optimizer Trainer resumes the same way through NNTrainerParams.builder().resume_from(run_id) (Concepts §8.4).

Probing a run that was never written is side-effect free. NNCheckpoint.load_with_training_state(run=..., type=Checkpoints.LAST) returns (None, None) for it and does not create the run's directory, so a preflight check never blocks the first fit. On an existing run, a probe takes the checkpoint's lock file, created when missing.

2.4. Loading a finished run

from nnx import NNRun, NNCheckpoint, Checkpoints, NNModel

run  = NNRun.load(id="<md5>")                              # rehydrate idps + params
ckpt = NNCheckpoint.load(run=run.id, type=Checkpoints.BEST)
model = NNModel.from_checkpoint(checkpoint=ckpt)

2.5. Tabular regression targets

NNTabularDataset treats targets as class labels by default. For regression, request a floating target dtype and use a one-output network with a regression loss. Declare the task so the default step, evaluate() and predict() score continuous values instead of argmax classes:

import pandas as pd
import torch
from nnx import Devices, Losses, Nets, NNModelParams, NNParams, NNTabularDataset, TaskSpec

frame = pd.DataFrame({"rooms": [1, 2, 3], "price": [95.0, 150.0, 220.0]})
dataset = NNTabularDataset(
    df=frame,
    feature_cols=["rooms"],
    target_col="price",
    target_dtype=torch.float32,
    batch_sizes=(2, 1, 1),
)
net_params = NNParams(
    input_dim=dataset.input_dim,
    output_dim=dataset.output_dim,
    hidden_dims=[16],
    dropout_prob=0.1,
)
model_params = NNModelParams(
    net=Nets.FEED_FWD,
    device=Devices.CPU,
    loss=Losses.MEAN_SQUARED_ERROR,
    task=TaskSpec.regression(dataset.output_dim),
)

With task=TaskSpec.regression(...) every training and validation record carries mse / mae in edp.metrics and the count of valid targets. The records leave accuracy / f1 / recall / precision and error as None instead of fabricating them. model.predict(X).classes holds the continuous predictions. A NaN target is masked out of the loss and every metric alike.

TaskSpec.multilabel(n, threshold=...) does the same for independent binary labels (Losses.BINARY_CROSS_ENTROPY). See Concepts §6.3 and examples/regression_task.py. Without a task the model keeps the legacy classification behaviour.

Leave target_dtype=None for classification, where targets remain torch.long and output_dim is the number of classes.

Admission checks. Only the selected feature_cols and target_col are inspected; unrelated columns may hold anything. NaN in any of them is rejected, naming the columns and the requested dtypes. The checks on values are:

• Non-finite features are never admitted. ±inf in a selected feature is rejected in source precision before the feature_dtype cast. An int64 or bool cast therefore cannot quietly absorb it (into an INT64 extreme or True). A non-finite target was already rejected.
• For an integer feature_dtype, finite values outside its range are rejected in source precision too, since the cast would wrap them. Examples are 300 with torch.int8 and 1e30 with torch.int64.
• A finite value that overflows a narrower floating or complex dtype during conversion is rejected after conversion. An example is 1e5 with feature_dtype or target_dtype=torch.float16.

Errors name the columns and dtypes. Every admission check, including the classification label check, runs before the train/val/test split. The caller's DataFrame and the global RNG are therefore untouched. See examples/tabular_validation.py.

Batch sizes. Every dataset wrapper (NNDataset, NNTabularDataset, NNPreferenceDataset, NNGraphDataset) takes batch_sizes=(train, val, test). None is the default for every slot and means one batch holding the complete split. The default train loader therefore performs one optimizer step per epoch. Pass an explicit train size for stochastic mini-batches. NNDataset wraps torchvision datasets, so it needs the vision extra (pip install "thekaveh-nnx[vision]"):

from torchvision import datasets, transforms
from nnx import NNDataset

dataset = NNDataset(
    ds_class=datasets.MNIST,
    transform=transforms.ToTensor(),
    batch_sizes=(128, None, None),   # 128-sample train batches; val/test = one full batch each
    val_proportion=0.1,
)

A positive integer is used verbatim (larger than the split → one smaller batch); NumPy integers are normalized to int. Zero, False, negatives, fractions, strings and anything but a 3-tuple raise a ValueError. The error names the slot (batch_sizes[0] (train)) and is raised before the dataset factory, tokenizer or split runs. Zero never disables a split.

Empty optional splits come from the proportions. With val_proportion=0.0 (tabular / preference: also test_proportion=0.0) that loader is None. Its resolved batch_sizes slot is a placeholder 1 when you requested it as None; an explicit size is kept. NNGraphDataset(sampler="full") rejects every explicit size, since the whole graph is always one batch.

2.6. Custom metrics

Declared, registered metrics receive the input they declare: decoded labels, probabilities or continuous outputs. They are computed over the full sample of each epoch. A named monitor makes BEST selection, ReduceLROnPlateau and EarlyStopping agree on one decision (Concepts §6.4):

from nnx import EarlyStopping, MetricSpec, MonitorSpec

monitor = MonitorSpec(metric="nll", min_delta=0.01)   # validation NLL, lower is better
run = model.train(
    params=NNTrainParams(..., metrics=[MetricSpec("nll"), MetricSpec("accuracy")], monitor=monitor),
    callbacks=[EarlyStopping(monitor=monitor, patience=5)],
)
run.idps[-1].val_edp.metrics["nll"]      # the whole validation set
run.idps[-1].train_summary              # the whole training epoch (full-epoch denominators)
run.idps[-1].selection                  # MonitorRecord: value, status, improved

Plain callables still work beside them:

from sklearn.metrics import f1_score

NNTrainParams(
    ...,
    extra_metrics={
        # Called as fn(y_true, y_pred): truth first, decoded class labels second.
        # zero_division=0 keeps per-batch calls quiet when a batch lacks a class.
        "weighted_f1": lambda y_true, y_pred: float(f1_score(y_true, y_pred, average="weighted", zero_division=0)),
    },
)
# Recorded in `.extra["weighted_f1"]` wherever NNx computes metrics.

Each callable receives (y_true, y_pred): the truth first and the decoded class predictions second, never probabilities. Use label-based metrics; ranking metrics such as ROC-AUC need scores that extra_metrics does not see. The callables run as follows:

• The default classification training step calls them per batch. A batch whose targets are all ignore_index has no metrics at all.
• evaluate(), and so the default validation pass, calls them once on the aggregate predictions.
• A custom train_step_fn / eval_step_fn decides whether and how to call them.

The values persist in idp.*_edp.extra and survive NNRun.load; see examples/03_custom_metrics.py.

2.7. Silencing the progress bar (CI / non-TTY)

The training loop draws a tqdm progress bar by default. Set NNX_TQDM_DISABLE=1 in the environment to silence it — useful in CI, in non-TTY contexts, and in test suites:

NNX_TQDM_DISABLE=1 python your_train_script.py

NNX_TQDM_DISABLE is read by both NNModel.train() and Trainer.train(). Any value of 1 / true / yes (case-insensitive) disables the bar; anything else leaves it enabled.

2.8. TensorBoard

pip install "thekaveh-nnx[tensorboard]"

from nnx import TensorBoardCallback
model.train(params=..., callbacks=[TensorBoardCallback(log_dir="tb_logs")])

2.9. LR finder pre-flight

Before a long training run, sweep learning rates exponentially and let the Smith-2017 steepest-descent heuristic pick a defensible max_lr for the real run. Its output is a Plotly figure, so it needs the plots extra (pip install "thekaveh-nnx[plots]") and refuses before touching anything without it. The sweep is non-destructive: model state, mixed per-module modes, loader generators, and Python/NumPy/PyTorch RNG streams are restored on exit.

import torch.nn.functional as F
from nnx import lr_finder

result = lr_finder(
    model.net, train_loader,
    loss_fn=F.cross_entropy,
    start_lr=1e-7, end_lr=10.0, num_iter=100,
)
print(f"Suggested max_lr: {result.suggested_lr:.2e}")
result.figure.show()  # Plotly: loss vs log(LR) with the suggestion marked

# Plug into the real training run:
NNTrainParams(..., optim=NNOptimParams(name=Optims.ADAM, max_lr=result.suggested_lr, ...))

See Concepts → LR finder for the algorithm details and divergence early-exit behavior.

2.10. Native NLL instead of cross-entropy

NNModelParams(net=Nets.FEED_FWD, device=Devices.CPU, loss=Losses.NEGATIVE_LOG_LIKELIHOOD)

The network still emits raw logits and predict().logits stays raw. NNx applies log_softmax internally before the native torch.nn.NLLLoss during training and evaluation. The reported loss therefore equals cross-entropy from the same weights. Checkpoints saved with this descriptor evaluate under the same rule when reloaded. A custom loss module or train_step_fn receives the raw logits.

2.11. Non-finite metrics, BEST and plateau scheduling

BEST checkpoints, runs/best and ReduceLROnPlateau all track the first finite value in a fixed order (lower is better). The order is validation error → validation loss → training error → training loss. The training values come from the epoch's last batch record; declare a named monitor (Concepts §6.4) to select on a whole-epoch training summary instead.

A custom eval_step_fn may report error=nan with a finite loss. It therefore still drives scheduling and BEST through its loss. The rejection is reported once per epoch as a RuntimeWarning.

A validation task record can be unavailable: every target is masked, or it has neither an error nor a loss. Such a record ends the walk, so that epoch gets no signal and the training values do not stand in (Concepts §6.2).

An epoch with no finite signal at all skips the plateau step instead of feeding it NaN. Its checkpoint counts as an unavailable baseline that the next finite epoch (or run) replaces. Raw NaN/inf observations stay in the run history as diagnostics.

2.12. Graph datasets with an empty validation split

NNGraphDataset reads the split sizes from the graph's train_mask / val_mask / test_mask. An empty optional mask is an absent split: that loader is None (resolved batch_sizes entry 0, "0" in state()). Passing it to train() skips validation instead of scoring zero seed rows. NNTabularDataset and NNPreferenceDataset also return a None loader for an empty split. Their resolved batch_sizes keep a placeholder 1 for a slot requested as None, and an explicit size is kept (§2.5).

An empty train_mask raises at construction. The example below runs fully offline and needs no neighbor-sampling backend:

import torch
from torch_geometric.data import Data
from nnx import (
    Activations, Devices, EarlyStopping, Losses, Nets,
    NNGraphDataset, NNModel, NNModelParams, NNParams, NNTrainParams,
)

class TinySplitGraph:                      # (root, transform) ctor + dataset[0] surface
    num_features, num_classes = 4, 2
    def __init__(self, root, transform=None):
        train = torch.tensor([True, True, True, False, False, False])
        self.data = Data(
            x=torch.randn(6, 4), y=torch.tensor([0, 1, 0, 1, 0, 1]),
            edge_index=torch.tensor([[0, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 0]]),
            train_mask=train, val_mask=torch.zeros(6, dtype=torch.bool), test_mask=~train,
        )
    def __getitem__(self, idx):
        return self.data

dataset = NNGraphDataset(ds_class=TinySplitGraph, sampler="full")
assert dataset.val_loader is None and dataset.batch_sizes == (3, 0, 3)

model = NNModel(
    net_params=NNParams(input_dim=dataset.input_dim, output_dim=dataset.output_dim,
                        hidden_dims=[8], dropout_prob=0.0, activation=Activations.RELU),
    params=NNModelParams(net=Nets.GRAPH_CONV, device=Devices.CPU, loss=Losses.CROSS_ENTROPY),
)
run = model.train(
    params=NNTrainParams(n_epochs=5, train_loader=dataset.train_loader, val_loader=dataset.val_loader),
    callbacks=[EarlyStopping(monitor="train_edp.loss", patience=2)],   # no validation signal to monitor
)
assert all(idp.val_edp is None for idp in run.idps)                   # persisted as absence, not zeros
print(model.evaluate(loader=dataset.test_loader).accuracy)            # test split scored on its seed rows

EarlyStopping's default monitor reads only validation data (val_edp.error, or val_edp.loss when the evaluator reports no error). Here it would have nothing to read — it would warn once and never stop — so point it at a train metric. See examples/graph_optional_splits.py for the runnable version (it also proves, with an eval_step_fn spy, that no validation hook runs).

2.13. Probabilities, decoded labels and sample ids

predict() returns raw logits and decoded classes. When you need probabilities — NLL, Brier score, calibration, a decision threshold — declare the task and call predict_proba instead; predict() itself is unchanged:

from nnx import ProbabilitySpec

spec = ProbabilitySpec(kind="categorical", class_axis=1, labels=("cat", "dog", "fox"))
result = model.predict_proba(X_val, spec)   # same inputs as predict(): arrays, tensors, tuples, loaders
result.logits                               # exactly predict(X_val).logits
result.probabilities                        # softmax over class_axis; each row sums to 1
result.decoded                              # argmax class indices == predict(X_val).classes
result.decoded_labels()                     # array(['dog', 'cat', ...], dtype=object)
result.sample_ids                           # identity of row i (input row index for arrays)

A model that declares a task (NNModelParams(task=TaskSpec...), §2.5) may omit the spec: model.predict_proba(X_val) derives it from the task. A regression task returns its continuous values in decoded with probabilities=None and spec=None.

kind="bernoulli" is for independent binary outputs: multi-label, or one binary logit trained with Losses.BINARY_CROSS_ENTROPY. It gives element-wise sigmoid probabilities that are not normalized across outputs. decoded holds 0/1 indicators (logit >= 0, the same threshold predict() uses). result.class_indices refuses to treat these indicators as class indices.

The kind is always declared, never inferred from the loss or output shape. class_axis defaults to -1, which fits (N, C) and class-last language-model logits. Class-first outputs such as (N, C, H, W) need class_axis=1. A label count that does not match the class axis, an invalid axis, or NaN / ±inf logits raise PredictionValidationError.

For an ordinary DataLoader, sample_ids are positions in iteration order. They are the dataset index only when the loader does not shuffle (a shuffling loader warns). For graph loaders only seed rows are returned, and sample_ids holds their global node indices.

A spec serializes with spec.state() / ProbabilitySpec.from_state(...), label order included. prediction_from_logits(logits, spec) applies the same rules to logits you already have.

2.14. Your own nn.Module

Any torch.nn.Module trains through the same loop. Register a factory to make it portable, or wrap an instance you already have:

from nnx import KeywordInputs, ModelSpec, NNModel, NNModelParams, register_model_factory

register_model_factory("acme.encoder", 1, lambda config: Encoder(**config))
model = NNModel(params=NNModelParams(net=ModelSpec("acme.encoder", 1, {"width": 64}), loss=Losses.CROSS_ENTROPY))
run = model.train(params=train_params)            # (x, y) batches reach forward(x)
rebuilt = NNModel.from_checkpoint(checkpoint)     # through the registry — no pickled code

wrapped = NNModel(module=encoder, params=NNModelParams(loss=Losses.CROSS_ENTROPY),
                  batch_adapter=KeywordInputs(("input_ids", "attention_mask"), target="labels"))

A registered module is built under the spec's seed (your RNG is left alone). It rebuilds from checkpoints, safetensors and the Hub. A wrapped instance is trained in place (wrapped.net is encoder) but is runtime-only. Its weights reload only into a module you pass again (NNModel.from_checkpoint(ckpt, module=...)). See Concepts §3.2 and examples/custom_module.py.

2.15. Recording what an experiment was

run.id names the configuration, and metadata.yaml the environment. To also record the declared intent and each execution, pass a manifest; this is opt-in and never changes the run id:

from nnx.provenance import ExperimentManifest, compare, hash_file

plan = ExperimentManifest.for_model(model, train=train_params, data={"train": hash_file("train.csv")})
run = model.train(params=train_params, provenance=plan)
run.provenance.fingerprint           # the plan: same manifest, same fingerprint
run.provenance.attempt.attempt_id    # this fit: fresh on every train() call
compare(run.id, other_run.id).differences

A plain id string (for example data={"train": "animals-v2"}) is recorded as declared and never reported as verified; only hash_file / hash_bytes digests are. A run without a manifest has run.provenance is None: absent, not equal. See Concepts §4.4.

2.16. Group, time and stratified splits

NNTabularDataset splits rows with a seeded random_split (val_proportion, test_proportion, seed). When rows are related (one patient, one user) or ordered in time, plan explicit membership by sample id instead and replay it:

from nnx.data_splits import SplitManifest, plan_split

plan = plan_split(df["row_id"], strategy="chronological", times=df["day"], cutoffs=(300, 330), gap=7)
plan.save("split.json")                         # plain JSON; plan.digest() is the split identity
dataset = NNTabularDataset(
    df=df, feature_cols=["age", "income"], target_col="label",
    split=SplitManifest.load("split.json"), id_col="row_id",
)

strategy="group" (with groups= and proportions=) keeps every group in one split, and strategy="stratified" (with labels=) keeps every class in every split. Rows are matched by id, so a reordered DataFrame replays into the same splits. Duplicate, missing or unexpected ids raise before any loader is built. An empty validation or test split gives a None loader. Without split= the random split is unchanged. See Concepts §13.3.

2.17. Train-only preprocessing and inference reload

Fit feature statistics on the training rows only, then serve the model with the same frozen statistics:

from nnx.preprocessing import Standardizer

dataset = NNTabularDataset(df=df, feature_cols=["age", "income"], target_col="label",
                           seed=0, standardize=True)       # fits on the training split only
run = model.train(params=NNTrainParams(n_epochs=5, train_loader=dataset.train_loader, ...))
dataset.standardizer.save("standardizer.json")          # float64 statistics + schema, as JSON

standardizer = Standardizer.load("standardizer.json")    # later: reload, never refit
prediction = model.predict(standardizer.transform(new_rows[["age", "income"]]))

NNModel.predict takes prepared features, the same space the model was trained on. Apply the reloaded standardizer exactly once to raw rows; a second pass would standardize twice. transform raises before inference on missing or reordered columns, a width mismatch and non-finite values. For image data, NNDataset(train_transform=..., eval_transform=...) gives training and evaluation their own transforms. See Concepts §13.4.

2.18. Calibrated probabilities

Fit one temperature on a calibration split, then report on held-out data:

from nnx.calibration import TemperatureCalibrator, fit_temperature, model_fingerprint

spec = ProbabilitySpec(kind="categorical", class_axis=1, labels=("cat", "dog", "fox"))
model_id = model_fingerprint(model)                       # changes whenever the weights change
fit = fit_temperature(model.predict_proba(X_calib, spec), y_calib, model_id=model_id,
                      split_id="calibration", train_split_id="train", test_split_id="test")
calibrator = fit.require()                                # CalibrationFitError with the reason if it failed
calibrator.save("calibrator.json")

calibrator = TemperatureCalibrator.load("calibrator.json")
result = calibrator.transform(model.predict_proba(X_test, spec), model_id=model_id)
result.probabilities, result.calibrated_probabilities     # raw vs calibrated, same label order
report = calibrator.report(model.predict_proba(X_test, spec), y_test, model_id=model_id, split_id="test")
print(report.summary())                                   # "improved" only if held-out NLL and Brier did not get worse

Use a calibration split that is neither the training nor the test split. fit_temperature rejects equal ids. The calibrator records that ids cannot prove the rows are disjoint. The fitted temperature is unrelated to the sampling temperature= of GenerativeNNModel.generate; it never changes the argmax. Logits with a different label order or from another model are refused unless you pass a named override=..., which is recorded. See Concepts §18.

2.19. Declining uncertain predictions

Choose an abstention threshold on validation, then apply it to new rows:

from nnx.abstention import select_threshold

selection = select_threshold(val_probabilities, y_val, kind="max_probability", risk_ceiling=0.05,
                             split_id="val", test_split_id="test", labels=labels, model_id=model_id)
policy = selection.require()                 # AbstentionError when no threshold meets the ceiling
policy.save("policy.json")

result = policy.apply(test_probabilities, labels=labels, model_id=model_id, sample_ids=test_ids)
result.accepted_ids, result.abstained_ids    # every row is one or the other
report = result.report(y_test)
report.coverage, report.risk                 # accepted/total and incorrect/accepted

A row is accepted when its top probability (kind="margin": the top-two gap) is at or above the threshold. Abstained rows keep their prediction, distribution and reason. With nothing accepted, report.risk is None (status "no_accepted"), not 0. The risk ceiling is measured on the validation split, not guaranteed. Pass a CalibratedPrediction (§2.18) to apply a policy declared on calibrated probabilities. See Concepts §19.

2.20. Branch, validate, probe and fit a plan

An ExperimentPlan is an immutable configuration you branch instead of copy:

from nnx import ExperimentPlan

base = (ExperimentPlan()
        .with_net(net_params).with_model(model_params).with_train(train_params)
        .with_data(make_train_loader, val=make_val_loader, identity="toy")  # factories: fresh loaders per fit
        .with_seed(7))
short, long = base.with_epochs(2), base.with_epochs(10)   # base, short and long stay unchanged

base.validate().raise_for_errors()   # every problem at once; reads no loader, builds no model
base.probe(next(iter(make_train_loader())))  # temporary model: eval no-grad forward, plus a train-mode loss pass if the batch has a target
result = base.fit()                  # FitResult(model, run, metrics, attempt_id)
result.metrics["val"].values["loss"] # from the run's final epoch; no second loader pass
child = base.with_epochs(20).resuming(result.run.id).fit()  # a new run; the parent stays intact

fit() seeds, calls the factories once, builds the model and calls NNModel.train, the same loop as the equally seeded imperative script. Every fit is a distinct attempt with its own run directory. Loader and callback instances you pass are borrowed and shared by every fit; pass factories for fresh ones. See Concepts §20.

2.21. Stream predictions and merge metrics

Stream a large loader instead of holding every batch's outputs:

from nnx import MetricSpec, StreamingMetrics, streaming_eval_step

with model.iter_predict(loader) as stream:          # one PredictionBatch per loader batch, in order
    for batch in stream:
        save(batch.sample_ids, batch.classes)        # keep what you need, not the batches

metrics = StreamingMetrics([MetricSpec("nll"), MetricSpec("brier")], "categorical")
with model.iter_predict(loader, rich=True) as stream: # PredictionResult per batch (model with a task)
    for batch in stream:
        metrics.update(targets[batch.sample_ids], probabilities=batch.probabilities)
metrics.finalize().values                             # {"nll": ..., "brier": ...}
total = shard_a.merge(shard_b).finalize()             # order-independent; refuses different declarations

model.train(params, eval_step_fn=streaming_eval_step) # the validation record from counts and sums

Concatenated, the streamed batches equal predict() / predict_proba(). Each batch restores the network's training mode, closing the stream drops its references, and the loader stays yours. See Concepts §21.

2.22. Share a run as a data-only bundle

A run bundle carries a checkpoint's weights, training state and calibrators as safetensors plus JSON, so it can be checked and rebuilt without unpickling anything:

from nnx.bundles import export_bundle, inspect_bundle, reconstruct_bundle, validate_bundle

export_bundle(run.id, "bundle", calibrators=[calibrator])  # reads your own run's LAST checkpoint
inspect_bundle("bundle")                                    # BundleInfo: capability, epoch, model, calibrators
validate_bundle("bundle")                                   # every payload's SHA-256, before any tensor is read
rebuilt = reconstruct_bundle("bundle", factories={("my.encoder", 1): my_factory})  # a registered module
rebuilt.model.predict(X)                                    # bit-for-bit the source model
rebuilt.resume(train_params)                                # continues the run from the bundle's training state

A weights-only checkpoint (a ModelCheckpoint snapshot) exports an "inference" bundle, which rebuilds the model but refuses resume(). See Concepts §22.

3. Beyond supervised classification

For tasks where loss isn't loss_fn(net(X), Y) — autoencoder reconstruction, VAE composite loss, recommendation pairwise loss, diffusion noise prediction — pass train_step_fn= to train(). Link prediction with negative sampling is built in (§3.9). See Concepts → Custom training paradigms.

The sections below point to the specialization paradigms built on these hooks (train_step_fn, objective=, eval_step_fn=) and on distributed=.

3.1. Fine-tuning (transfer learning)

Load external pretrained weights, freeze layers by glob pattern, and (optionally) train them at different learning rates. Parameter-group rules are matched in order and the first match wins (rules never merge), so specific rules go first:

# For a net with `encoder` / `head` submodules (FeedFwdNN uses `layers.0`, `layers.1`, ...;
# check `named_parameters()` — a pattern that matches nothing is silently empty).
param_groups=[
    NNParamGroupSpec(name_pattern="encoder.*bias", lr_multiplier=0.01, weight_decay=0.0),
    NNParamGroupSpec(name_pattern="encoder.*",     lr_multiplier=0.01),
    NNParamGroupSpec(name_pattern="*.bias",        weight_decay=0.0),
]

See Concepts → Fine-tuning and examples/06_finetune_with_layer_freezing.py.

3.2. Multi-optimizer training (GANs, actor-critic)

When per-batch updates need multiple optimizers, use nnx.trainer.Trainer. Examples are G + D for GANs and policy + value for actor-critic. Trainer accepts one NNModel and a dict of NNOptimParams, scoped via NNParamGroupSpec globs. See Concepts → Multi-optimizer training and examples/09_gan_with_trainer.py.

3.3. Diffusion (DDPM)

For DDPM-style diffusion: nnx.diffusion.{NoiseSchedulers, DiffusionMLP, diffusion_train_step_factory, sample}. The training step is a train_step_fn on NNModel.train() — no Trainer, no new params dataclass.

diffusion_objective(schedule) is the same loss as an objective (model.train(..., objective=...)). The shared update engine adds gradient accumulation (exact for uneven microbatches), mixed precision and clipping. The timesteps and noise come from the objective's own checkpointed generator. A sample(..., generator=...) preview never perturbs that generator. jepa_objective does the same for I-JEPA, advancing its EMA target once per committed update (I-JEPA §5). See Concepts → Diffusion and examples/08_diffusion_2d_mixture.py.

3.4. Training paradigms (KD, SimCLR, Mixup, CutMix)

nnx.paradigms.{kd, simclr, mixup, cutmix}_train_step_factory return train_step_fns for NNModel.train(). Knowledge distillation freezes the teacher and mixes soft/hard losses; SimCLR runs NT-Xent on paired-view batches; Mixup / CutMix interpolate samples within a batch. See Concepts → Training paradigms and examples/10_knowledge_distillation.py.

3.5. Parameter-efficient fine-tuning (LoRA, DoRA, IA3, Prefix, Prompt, Adapters)

nnx.peft.{LoRALinear, apply_lora_to, save_lora_weights, load_lora_weights, AdapterLayer} plus DoRA / IA3 / PrefixTuner / PromptTuner. LoRA wraps nn.Linear submodules with a frozen base + trainable low-rank residual. DoRA layers in a per-output magnitude vector, and IA3 is a per-output scaling. PrefixTuner / PromptTuner attach learned prefixes to a frozen TransformerNN. AdapterLayer is a bottleneck residual the user inserts manually. See Concepts → Parameter-efficient fine-tuning and examples/07_lora_finetuning.py.

3.6. Causal language models

nnx.lm_tasks.CausalLMTask(vocab_size=..., alignment="shift_inputs" | "pre_shifted", ignore_id=-100, pad_id=None, smoothing=0.0) is the built-in next-token path for a TransformerNN: pass objective=task.objective() and eval_step_fn=task.eval_step() to train(). The ignore id, padding id and an optional loss mask pick the valid tokens. The objective divides by the window's valid-token count, so accumulation is token-weighted. Validation reports the unsmoothed token NLL as its loss, with perplexity = exp(NLL) and token accuracy. It reports no placeholder classification metrics. See docs/lm.md §6 and examples/11_tinystories_lm.py.

3.7. Ranking and retrieval evaluation

nnx.ranking.RankingTask(k=(10,), max_relevance=2, weighting="query", candidate_sets="sampled") trains and evaluates a scorer per query:

• Pass objective=task.objective(), a pairwise logistic loss over unequal-relevance pairs within a query.
• Pass eval_step_fn=task.eval_step(). It reports MRR@k, Recall@k and NDCG@k per complete query, with ties broken by candidate id. Queries without a relevant candidate are excluded and counted.
• Select BEST with metrics=task.metric_specs(), monitor=MonitorSpec("ndcg_at_10").

Batches are (features, query_ids, candidate_ids, relevance). See Concepts → Query-grouped ranking and examples/ranking_offline.py.

3.8. Graph-level classification

nnx.graph_tasks (the graph extra, pip install "thekaveh-nnx[graph]") classifies whole graphs:

1. Build a GraphCollection(graphs, ids, targets=...). Flag graphs without a target with unlabeled=[...].
2. Split it by graph id with subset(ids).
3. Train NNModelParams(net=graph_classifier_spec(input_dim=..., num_classes=..., pool="mean"), loss=Losses.CROSS_ENTROPY, task=TaskSpec.categorical(n, ignore_index=IGNORE)) on collection.loader(batch_size).

Loss and metrics are averaged over labeled graphs. predict_proba returns one row per graph, with its graph id as the sample id. A reload rebuilds the encoder, pool and head from the recipe. See Concepts → Graph-level classification and examples/graph_classification_offline.py.

3.9. Link and edge prediction

nnx.link_tasks (the graph extra) splits edges, not nodes. split = split_links(edge_index, num_nodes, val=0.1, test=0.1, seed=0) gives a replayable manifest: canonical edges, and fixed held-out negatives from the complement. task = LinkTask(split) builds candidate batches over the training topology only, through task.loader("train" | "val" | "test", x, batch_size).

1. Train NNModelParams(net=link_predictor_spec(input_dim=...)) with objective=task.objective() and eval_step_fn=task.eval_step().
2. Select BEST with metrics=task.metric_specs(), monitor=MonitorSpec("auroc").
3. Score candidates with prediction = task.predict(model, loader) and link_metrics(prediction.logits, prediction.targets, from_logits=True).

The metrics are computed from the float64 logits, as the evaluator's are, so a confident model's AUROC / AP / BCE stay exact. A held-out positive in a message graph fails before any update. The default train / predict paths refuse link batches; use the task's objective, evaluator and predict. See Concepts → Link and edge prediction and examples/link_prediction_offline.py.

3.10. Data-parallel training on one machine

Launch one process per device with torchrun. NNx partitions the data and keeps every update equal to the union batch's. One rank writes the run:

# train.py — torchrun --standalone --nproc_per_node=2 train.py
from nnx import distributed

distributed.init_process_group()  # env:// from torchrun; gloo on CPU, nccl with CUDA
params = NNTrainParams(
    n_epochs=3,
    train_loader=distributed.train_loader(train_set, batch_size=32, policy="pad", seed=0),
    val_loader=distributed.validation_loader(val_set, batch_size=64),
    optim=NNOptimParams.builder().adamw(max_lr=1e-3).build(),
)
run = model.train(params=params, distributed=distributed.DDP())  # the same run on every rank
distributed.shutdown()  # leave together: a bounded barrier, then the group is destroyed

DDP supports FP32, the default steps and one optimizer on one node. Callbacks declare their rank behaviour (Callback.distributed) and TensorBoardCallback / WandbCallback go through distributed.writer_only(...). See Architecture §3 and examples/ddp_supervised.py.