12. API Reference¶
Generated from runtime signatures and source docstrings as portable Markdown so the same reference is readable in the repository, on the documentation site, and in the wiki. Sections are ordered from foundational APIs to specialized modules.
1. Top-level package¶
nnx¶
module nnx
nnx — lightweight PyTorch training / eval / visualization toolkit.
Details
The core lives under `nnx.nn` (model, params, datasets, enums, nets,
callbacks); feature subpackages and modules such as `nnx.paradigms`,
`nnx.peft`, `nnx.calibration` and `nnx.distributed` sit beside it (see
`__all__`). The curated re-exports below give a flat surface for the most
common imports without forbidding the deep paths existing notebook code
relies on.
nnx.__version__¶
nnx.__version__ = '0.4.1' # x-release-please-version
Exported value.
nnx.set_seed¶
nnx.set_seed(seed: 'int', strict: 'bool' = False) -> 'None'
Pin every RNG that affects training and toggle cuDNN deterministic.
Details
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.
nnx.dataloader_worker_init_fn¶
nnx.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.
Details
Pass as: `DataLoader(..., worker_init_fn=dataloader_worker_init_fn)`.
nnx.env_snapshot¶
nnx.env_snapshot(force_refresh: 'bool' = False) -> 'dict'
Capture a snapshot of the runtime environment for reproducibility.
Details
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.
nnx.LRFinderResult¶
class nnx.LRFinderResult(*, lrs: 'list[float]', losses: 'list[float]', suggested_lr: 'float', figure: 'go.Figure') -> 'None'
Result of an lr_finder sweep.
Details
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.
nnx.lr_finder¶
nnx.lr_finder(model: 'nn.Module', train_loader: 'DataLoader', *, loss_fn: 'Callable[[torch.Tensor, torch.Tensor], torch.Tensor]', optimizer_cls: 'type[torch.optim.Optimizer]' = torch.optim.adam.Adam, start_lr: 'float' = 1e-07, end_lr: 'float' = 10.0, num_iter: 'int' = 100, diverge_threshold: 'float' = 4.0, device: 'Optional[torch.device]' = None, ema_alpha: 'float' = 0.5) -> 'LRFinderResult'
Sweep LRs exponentially from start_lr to end_lr and suggest a one-cycle max_lr.
Details
Args:
model: the network to sweep against. ``model.train()`` is
called internally; the original training-mode state, the
weights, AND the RNG state are all restored on exit.
train_loader: a DataLoader yielding ``(X, Y)`` batches the
model can forward and ``loss_fn`` can score against. The
loader is iterated, and if the sweep exceeds one epoch
the loader is re-iterated from the start.
loss_fn: callable ``(y_hat, Y) -> scalar Tensor`` for the
per-batch loss. Same shape contract as torch loss
functions. It receives the model's raw output: unlike
``NNModel.train`` / ``evaluate``, no log-softmax is inserted
for ``torch.nn.NLLLoss``, so pass
``lambda y_hat, Y: F.nll_loss(F.log_softmax(y_hat, dim=1), Y)``
(or simply ``F.cross_entropy``) when sweeping an NLL model.
optimizer_cls: optimizer class. Adam by default; SGD also
works for the sweep.
start_lr: low end of the sweep range. Must be > 0.
end_lr: high end of the sweep range. Must be > start_lr.
num_iter: number of training iterations to run. Must be >= 2.
diverge_threshold: stop the sweep early when the EMA-smoothed
loss rises more than ``(diverge_threshold - 1) * |smoothed_min|``
above ``smoothed_min`` (the minimum EMA-smoothed loss observed so
far; for a positive loss, past ``diverge_threshold * smoothed_min``). Default 4. The
smoothed check matches fastai's lr_find heuristic — using
the raw ``min(losses)`` would let a single anomalous low
first-batch loss pull the threshold too tight and abort
the sweep prematurely.
device: device to move batches to. If None, inferred from
the first model parameter.
ema_alpha: smoothing coefficient for the loss curve before
the steepest-descent search. Default 0.5.
Returns:
:class:`LRFinderResult` with the raw sweep data, the suggested
max_lr, and a Plotly figure of loss vs log(LR).
Raises:
ImportError: Plotly (the ``plots`` extra, ``pip install
"thekaveh-nnx[plots]"``) is not installed — raised before any
data, model or RNG state is touched.
ValueError: on invalid arguments (``num_iter`` not an integer >= 2, a non-finite
``start_lr`` / ``end_lr`` / ``diverge_threshold``,
``start_lr <= 0``, ``end_lr <= start_lr``,
``diverge_threshold <= 1``, ``ema_alpha`` outside ``(0, 1)``), a
``train_loader`` with a live persistent-worker iterator, or a
loader that yields no batch.
2. Orchestrators¶
2.1. NNModel — supervised orchestrator¶
nnx.nn.nn_model.NNModel¶
class nnx.nn.nn_model.NNModel(net_params: 'Optional[NNParams]' = None, params: 'Optional[NNModelParams]' = None, *, module: 'Optional[torch.nn.Module]' = None, batch_adapter: 'Optional[BatchAdapter]' = None)
Top-level training/eval/predict wrapper around an nn.Module.
Details
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.
nnx.nn.nn_model.NNModel.task_adapter¶
property nnx.nn.nn_model.NNModel.task_adapter
The adapter for params.task, or None for a legacy classification model. Custom steps can call task_adapter.record(output, target, loss=...) to write the same task record the default step writes.
nnx.nn.nn_model.NNModel.resolved_precision¶
property nnx.nn.nn_model.NNModel.resolved_precision
The precision this model runs in on its current device: params.precision resolved against device — re- resolved, never read from saved metadata, whenever the device or the policy changes (a loaded model resolves on its destination device).
nnx.nn.nn_model.NNModel.to_onnx¶
nnx.nn.nn_model.NNModel.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.
Details
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.
A network that declares ``onnx_export_unsupported`` (a graph
classifier) is refused before anything is written, and so is a
mapping example: the export passes positional inputs only, so a
keyword-input module (``KeywordInputs``) needs a positional wrapper.
nnx.nn.nn_model.NNModel.from_checkpoint¶
nnx.nn.nn_model.NNModel.from_checkpoint(checkpoint: 'NNCheckpoint', device: 'Optional[Devices]' = None, *, module: 'Optional[torch.nn.Module]' = None, batch_adapter: 'Optional[BatchAdapter]' = None, precision: 'Optional[PrecisionPolicy]' = None, exclude_submodules: 'Sequence[str]' = (), **model_kwargs: 'Any') -> 'Self'
Rebuild a model, replay topology transforms, and load its weights.
Details
``exclude_submodules`` names top-level submodules the training run
attached to the net that are not part of the rebuilt architecture —
a JEPA predictor registered as ``model.net._jepa_predictor``, say —
whose weights are left out; every other weight still loads strictly.
A name the checkpoint does not hold raises ``ValueError``.
Ordinary and legacy FP32 checkpoints have no transforms. Converted
QAT checkpoints replay their persisted torchao recipe before state
loading; unsupported recipes fail explicitly rather than constructing
a model with the wrong topology.
A registered :class:`~nnx.models.ModelSpec` is rebuilt
through its factory — an unregistered one raises
:class:`~nnx.models.MissingModelFactoryError`, and a rebuilt topology
that differs from the saved weights raises ``ValueError``, both
before any weight is loaded. A runtime-only module
(``reconstructible=False``) needs ``module=`` — a module of the same
topology, into which the weights are loaded.
The precision policy is re-resolved on the destination
device (the constructor resolves it) — never read from the
checkpoint's recorded precision — so a policy that device cannot
run fails here; ``precision=`` replaces the saved policy (and the
legacy ``mixed_precision`` flag) for this model.
nnx.nn.nn_model.NNModel.save_pretrained¶
nnx.nn.nn_model.NNModel.save_pretrained(self, save_directory, *args: 'Any', **kwargs: 'Any') -> 'Any'
Hub save (huggingface_hub.PyTorchModelHubMixin.save_pretrained).
Details
A runtime-only module has no factory to rebuild it from,
so it is rejected with :class:`~nnx.models.MissingModelFactoryError`
before any file or directory is written.
nnx.nn.nn_model.NNModel.freeze¶
nnx.nn.nn_model.NNModel.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.
Details
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).
nnx.nn.nn_model.NNModel.unfreeze¶
nnx.nn.nn_model.NNModel.unfreeze(self, *patterns: 'str') -> 'int'
Mirror of freeze — set requires_grad=True on matching parameters.
nnx.nn.nn_model.NNModel.export_state_dict¶
nnx.nn.nn_model.NNModel.export_state_dict(self, path: 'str') -> 'str'
Save just self.net.state_dict() to path.
Details
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. It records no
transformation recipe (``nnx.transforms``): it cannot
rebuild a recipe's topology alone — materialize the recipe on a
fresh model before loading it, or keep a checkpoint or
``save_pretrained`` artifact, which carry the recipe.
Returns ``path`` so calls can be chained.
nnx.nn.nn_model.NNModel.train¶
nnx.nn.nn_model.NNModel.train(self, params: 'NNTrainParams', callbacks: 'Optional[list[CallbackLike]]' = None, train_step_fn: 'Optional[TrainStepFn]' = None, eval_step_fn: 'Optional[EvalStepFn]' = None, salt: 'Optional[str]' = None, components: 'Optional[list[Any]]' = None, objective: 'Optional[Callable[[Any], Any]]' = None, provenance: 'Optional[ExperimentManifest]' = None, history: 'Optional[HistoryJournal]' = None, compile: 'Optional[CompileSpec]' = None, distributed: 'Optional[Any]' = None) -> 'NNRun'
Train the model and return its persisted run history.
Details
Args:
params: Required loaders, optimizer/scheduler configuration,
epoch count, persistence controls, and optional resume source.
callbacks: Lifecycle callbacks invoked around training and epochs.
train_step_fn: Optional per-batch override; the default performs
supervised forward, loss, backward, and optimizer stepping.
eval_step_fn: Optional once-per-epoch validation override that
receives the complete validation loader.
salt: Optional string folded into the run.id hash so identical
(model, net, train) configs run as distinct experiments
without altering modeled params. ``None`` (the default)
preserves existing run.id hashes exactly.
components: Extra checkpointable components
(:class:`~nnx.StatefulComponent`) whose state is saved with
every checkpoint and restored on a stateful warm resume.
Callbacks and step functions that implement the protocol
(``EarlyStopping``, the JEPA step) are registered
automatically; names must be unique.
objective: Optional objective (``nnx.objectives``) —
a callable returning loss terms with explicit denominators.
NNx's shared update engine then owns backward,
accumulation (``accumulate_grad_batches``, exact for uneven
microbatches and masks), mixed precision, clipping and the
optimizer step, firing ``Callback.on_optimizer_update`` once
per committed update. Mutually exclusive with
``train_step_fn``.
provenance: Optional :class:`~nnx.provenance.ExperimentManifest`
— the declared intent. The fit records it
(``runs/<id>/provenance.json``, with its fingerprint) and a
fresh attempt (``attempt.json``: parent attempt and
checkpoint generation on resume; final status and last
committed checkpoint). Never part of the run id.
history: Optional :class:`~nnx.history.HistoryJournal`:
keep only its ``retention`` most recent records
in memory (``ctx.idps``, the returned ``NNRun.idps``) and
append every record once to ``runs/<id>/history/`` instead
of rewriting ``idps.csv`` each epoch. ``None`` (the default)
keeps the eager in-memory list and CSV. Never part of the
run id.
compile: Optional :class:`~nnx.compilation.CompileSpec`:
run the built-in step's FP32 forward (and the
fit's validation) through ``torch.compile``. The wrapper is
built per call beside ``model.net``, never assigned to it, so
checkpoints, export and optimizers see the eager module.
Refused, before any run is reserved, with a custom
``train_step_fn`` or ``objective``, a graph or custom net, a
reduced precision or a topology-changing callback (QAT).
``NNRun.compile`` records the request and what took effect.
``None`` (the default) trains eagerly. Never part of the run
id.
distributed: Optional :class:`nnx.distributed.DDP`:
train data-parallel over the process group ``torchrun``
started (``nnx.distributed.init_process_group()``), with
``nnx.distributed.train_loader`` / ``validation_loader``
partitions. Each update equals a single process training on
the ranks' union batch; records are global and identical on
every rank; only the writer rank holds the lease and writes
artifacts, and every rank returns the same run. FP32, the
default steps and one node only — refused, on every rank
together, before any run is reserved otherwise. Never part
of the run id.
Returns:
The completed :class:`NNRun`, persisted with run metadata,
iteration history, and configured checkpoints under
``<cwd>/runs/<run.id>/``. The printed completion line names
``runs/<id>`` relative to the working directory: a display path
with no absolute prefix, so captured notebook output stays
portable; it is not artifact provenance.
Raises:
ValueError: If required training inputs are missing or invalid,
the model is fully frozen, or resume state is incompatible.
FileExistsError: If the content-addressed run already exists and
``overwrite_existing`` is false.
FloatingPointError: If the default step encounters non-finite loss.
The run lease prevents another process using ``overwrite_existing``
from deleting or interleaving artifacts until final persistence ends.
nnx.nn.nn_model.NNModel.evaluate¶
nnx.nn.nn_model.NNModel.evaluate(self, loader: 'Iterable[Any]', extra_metrics=None, metrics: 'Sequence[MetricSpec]' = ()) -> '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.
Details
`extra_metrics` ({name -> callable(y_true, y_pred) -> float}) are
called once on the aggregate truth / decoded predictions.
`metrics` are declared :class:`~nnx.MetricSpec`s,
accumulated over every valid sample of the loader from the input
each declares (decoded labels, probabilities or continuous outputs)
and reported in the record's ``metrics`` under each spec's name. A
metric whose input this model cannot provide raises before any
batch is read.
Raises ValueError if the loader yields zero batches (there would be
nothing to average).
nnx.nn.nn_model.NNModel.predict¶
nnx.nn.nn_model.NNModel.predict(self, X) -> 'PredictResult'
Run the network in eval mode and return raw logits and decoded predictions (argmax classes; continuous values for a regression task; 0/1 indicators for a bernoulli / BCE output).
Details
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.
- a list or tuple of graph batches (``NNGraphDataset(sampler="full")``'s
loaders) — read like a loader, graph seed rows only.
Returns a ``PredictResult`` (a ``NamedTuple`` of (logits, classes))
that unpacks like the original 2-tuple. For probabilities, decoded
labels and per-row sample ids, use :meth:`predict_proba`.
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.
nnx.nn.nn_model.NNModel.iter_predict¶
nnx.nn.nn_model.NNModel.iter_predict(self, X: 'Iterable[Any]', spec: 'Optional[ProbabilitySpec]' = None, *, rich: 'bool' = False) -> 'PredictionStream'
Stream predictions one loader batch at a time.
Details
Returns a :class:`~nnx.streaming.PredictionStream` — use it as a
context manager — over ``X``, a ``DataLoader`` or another iterable of
batches (in-memory arrays and tensors go to :meth:`predict`). Each
item is a :class:`~nnx.streaming.PredictionBatch` of ``logits``,
``classes`` and ``sample_ids``, or, with a ``spec`` (or ``rich=True``
for a model with a task), a :class:`~nnx.prediction.PredictionResult`
as :meth:`predict_proba` builds it. The batches follow loader order,
and concatenated they are exactly the eager result for the same
``DataLoader`` (the eager calls read other iterables, except a list of
graph batches, as one in-memory input): ``predict(X)``'s logits and classes and ``predict_proba(X)``'s
sample ids, graph seed-row slicing included.
Each batch runs in eval mode under ``no_grad``, and every submodule's
training mode is restored before the batch is yielded or its error
raised. The stream holds only the batch in flight; closing it drops
its references to the loader's iterator and the model and ends the
iteration, and a closed or consumed stream cannot be iterated again.
An empty loader yields no batches (the eager calls raise instead).
Over a shuffling ``DataLoader``, the first batch whose sample ids are
iteration positions warns (graph seed rows carry global node indices,
graph-level rows their own graph ids).
nnx.nn.nn_model.NNModel.predict_proba¶
nnx.nn.nn_model.NNModel.predict_proba(self, X, spec: 'Optional[ProbabilitySpec]' = None) -> 'PredictionResult'
Probability-aware prediction declared by an explicit spec.
Details
``spec`` may be omitted for a model with a task: the task
supplies it — softmax over axis 1 for ``categorical``, sigmoid
decoded at the task's ``threshold`` for ``multilabel`` — and a
``regression`` task returns its continuous values with
``probabilities=None`` and ``spec=None`` (``decoded`` holds the
values). An explicit ``spec`` always wins.
Accepts the same inputs as :meth:`predict` (arrays, tensors, tuples
and ``DataLoader``s, including graph loaders whose rows are sliced
to seed nodes) with the same non-destructive eval-mode contract,
and returns a :class:`~nnx.prediction.PredictionResult`: the same
raw logits ``predict()`` returns, probabilities (softmax over
``spec.class_axis`` for ``"categorical"``, element-wise sigmoid for
``"bernoulli"``), decoded values and ``sample_ids``: the row index
for arrays, tensors and tuples; the position in iteration order for
an ordinary ``DataLoader`` (the dataset index only when the loader
does not shuffle — a shuffling loader warns); and the global node
index for graph seed rows. A spec that does not fit the logits is
rejected on the first loader batch.
The task kind is never inferred from the loss or the output shape.
For the default decoding rules, a categorical spec's ``decoded``
equals ``predict().classes`` when the class axes agree, and a
bernoulli spec's equals the ``BCEWithLogitsLoss`` threshold.
Raises :class:`~nnx.prediction.PredictionValidationError` for a
class axis or label count that does not fit the logits, or for
non-finite logits. Parameters and gradients are never touched.
nnx.nn.nn_model.PredictResult¶
class nnx.nn.nn_model.PredictResult(logits: 'np.ndarray', classes: 'np.ndarray')
Structured result of NNModel.predict().
Details
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.
nnx.nn.nn_model.TrainStepContext¶
class nnx.nn.nn_model.TrainStepContext(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', is_last_batch: 'bool' = False, accumulation_state: 'Optional[GradientAccumulationState]' = None, epoch_summary: 'Optional[_TrainEpochSummary]' = None, precision: 'Optional[ResolvedPrecision]' = None, report_update: 'Callable[[], None]' = <no update listener>) -> 'None'
Frozen bundle of state passed into a training-step function.
Details
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.
nnx.nn.nn_model.TrainStepFn¶
type alias nnx.nn.nn_model.TrainStepFn
Public type alias.
nnx.nn.nn_model.default_train_step¶
nnx.nn.nn_model.default_train_step(ctx: 'TrainStepContext') -> 'NNEvaluationDataPoint'
Standard supervised training step: forward → loss → backward → step.
Details
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). A trailing partial cycle is stepped at the epoch boundary;
gradients use each loss's effective normalization weight.
- the run's precision: autocast around the forward only,
the backward outside it; fp16 unscales before grad clip and steps
through the scaler, scaler-free bf16 checks its gradients are
finite before clipping
- 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.
nnx.nn.nn_model.EvalStepContext¶
class nnx.nn.nn_model.EvalStepContext(model: 'NNModel', val_loader: 'Iterable[Any]', extra_metrics: 'Optional[Mapping[str, Callable]]', epoch_idx: 'int', metrics: 'tuple[MetricSpec, ...]' = ()) -> 'None'
Frozen bundle of state passed into a validation-step function.
Details
Mirrors :class:`TrainStepContext` for the per-epoch VALIDATION pass: users
can pass ``eval_step_fn: Callable[[EvalStepContext], NNEvaluationDataPoint]``
to ``NNModel.train()`` to replace the built-in classification ``evaluate()``
for non-classification paradigms (next-token LM perplexity, DPO margins,
regression MAE, ...). The step runs under ``torch.no_grad()`` with the
net in eval mode (every submodule's mode is restored after it), and its
returned EDP becomes ``val_edp`` — recorded on the epoch's last idp and
persisted through the incremental run save like any built-in val metric.
nnx.nn.nn_model.EvalStepFn¶
type alias nnx.nn.nn_model.EvalStepFn
Public type alias.
2.2. GenerativeNNModel — decoder-only LM orchestrator¶
nnx.nn.generative_nn_model.GenerativeNNModel¶
class nnx.nn.generative_nn_model.GenerativeNNModel(net_params: 'NNParams', params: 'NNModelParams', tokenizer: 'Optional[NNTokenizerParams]' = None)
Language model with an autoregressive generate() method.
Details
``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.
nnx.nn.generative_nn_model.GenerativeNNModel.generate¶
nnx.nn.generative_nn_model.GenerativeNNModel.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, on_token: 'Optional[Callable[[int], None]]' = None, logits_pipeline: 'Optional[OrderedLogitsPipeline]' = None) -> 'str'
Autoregressive decode from prompt.
Details
Args:
prompt: input text. Encoded via ``self.tokenizer``.
max_new_tokens: hard cap on new tokens emitted. Generation
also stops when a ``stop`` string is decoded. It does not
stop on an end-of-sequence token, and a prompt longer than
the context window (max_seq_len) slides the window rather
than stopping.
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). The continuation is decoded without
special tokens, so a special token such as ``"<eos>"``
never matches a stop string.
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.
logits_pipeline: optional :class:`~nnx.generation.OrderedLogitsPipeline`:
its stages run exactly in declared order,
compiled once per call into fresh processors; like
``logits_chain`` it replaces the inline kwargs. Passing
both ``logits_chain`` and ``logits_pipeline`` raises
``ValueError`` before the model is touched.
on_token: optional callback invoked with each newly
generated token id immediately after it is appended.
Lets callers stream partial output or drive progress
reporting without re-running decode over the public forward /
apply_chain / sample_next_token primitives. ``None``
(default) is a no-op so existing callers are
unaffected. Fires only for newly generated tokens
(not prompt tokens) and on both the cached and
no-cache decode paths. Fires before any ``stop``
string check so the callback observes every emitted
token including the one that triggers a stop.
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``.
2.3. Trainer — multi-optimizer orchestrator¶
nnx.trainer.trainer.Trainer¶
class nnx.trainer.trainer.Trainer(model: 'NNModel')
Multi-optimizer training orchestrator.
Details
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.
nnx.trainer.trainer.Trainer.train¶
nnx.trainer.trainer.Trainer.train(self, params: 'NNTrainerParams', trainer_step_fn: 'Optional[TrainerStepFn]' = None, callbacks: 'Optional[list[CallbackLike]]' = None, salt: 'Optional[str]' = None, components: 'Optional[list[Any]]' = None, objective: 'Optional[Callable[[Any], Any]]' = None, provenance: 'Optional[ExperimentManifest]' = None, history: 'Optional[HistoryJournal]' = None) -> 'NNRun'
Run the multi-optimizer training loop and return the resulting NNRun.
Details
Args:
params: NNTrainerParams — train_loader + n_epochs + optims dict +
(optional) schedulers dict + (optional) val_loader, seed,
save_phase_checkpoints, extra_metrics. Schedulers step once
per epoch by default (an ``optimizer_update``-clock scheduler
after each reported update); set auto_step_schedulers=False when the
custom step function owns scheduler timing.
trainer_step_fn: `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: pass it or
``objective``.
objective: an objective (``nnx.objectives``) in place of
``trainer_step_fn``: the shared update engine accumulates its
loss terms over each window (the optimizers'
``accumulate_grad_batches``, which must agree), clips each
optimizer's parameters with its own ``grad_clip_norm`` and
steps every named optimizer that holds gradients once per
committed update (an optimizer without gradients is not
stepped but still counts), announcing each to
``Callback.on_optimizer_update``.
provenance: an optional ``nnx.provenance.ExperimentManifest``:
the declared intent, recorded with a fresh
attempt exactly as ``NNModel.train`` records it.
history: an optional ``nnx.history.HistoryJournal``:
a bounded in-memory window and an append journal instead of
the eager list and ``idps.csv``, exactly as in
``NNModel.train``.
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.
salt: mirrors ``NNModel.train()``'s ``salt`` parameter — an
optional string folded into the run.id hash so identical
(model, net, train) configs run as distinct experiments
without altering modeled params. ``None`` (the default)
preserves existing run.id hashes exactly.
components: extra checkpointable components;
callbacks and a step function that implement
:class:`~nnx.StatefulComponent` register automatically.
Returns:
NNRun with per-iteration idps, persisted under runs/<run.id>/
alongside the standard FIRST/Q1/Q2/Q3/LAST/BEST checkpoints.
The printed completion line uses the same cwd-relative
``runs/<id>`` display path as ``NNModel.train``.
Raises:
ValueError: when params is None, params.train_loader is None,
neither or both of trainer_step_fn and objective are given,
any optim's NNOptimParams.is_valid() returns False, a
trainer_step_fn run sets grad_clip_norm or
accumulate_grad_batches, an objective run's optimizers share
a parameter, an objective run sets
auto_step_schedulers=False, or a callback registers update
listeners or sets ctx.stop_at_update.
RuntimeError: in a process group of more than one rank.
nnx.trainer.trainer.TrainerStepContext¶
class nnx.trainer.trainer.TrainerStepContext(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', report_update: 'Callable[[str], None]' = <no update listener>) -> 'None'
Per-batch state passed into a trainer_step_fn.
Details
Mirrors TrainStepContext from NNModel.train() but with `optimizer`
(singular) replaced by `optimizers` (name-keyed dict) and `schedulers`
threaded through alongside for inspection. Step functions should only
call schedulers directly when ``auto_step_schedulers=False``.
`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).
nnx.trainer.trainer.TrainerStepFn¶
type alias nnx.trainer.trainer.TrainerStepFn
Public type alias.
nnx.trainer.params.NNTrainerParams¶
class nnx.trainer.params.NNTrainerParams(*, n_epochs: 'int', optims: 'Mapping[str, Union[NNOptimParams, NNOptimFactoryParams]]', schedulers: 'Mapping[str, NNSchedulerParams]' = <factory>, seed: 'Optional[int]' = None, data_id: 'Optional[str]' = None, save_phase_checkpoints: 'bool' = True, auto_step_schedulers: 'bool' = True, overwrite_existing: 'bool' = False, train_loader: 'Optional[DataLoader]' = None, val_loader: 'Optional[DataLoader]' = None, extra_metrics: 'Optional[Mapping[str, Callable]]' = None, resume_from_run_id: 'Optional[str]' = None, resume_from_checkpoint: 'str' = 'last', parent_run_id: 'Optional[str]' = None, resume_mode: 'str' = 'auto', metrics: 'tuple[MetricSpec, ...]' = (), monitor: 'Optional[MonitorSpec]' = None) -> 'None'
Configuration for Trainer.train() — the multi-optimizer parallel to NNModel.train() / NNTrainParams.
Details
`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 (`extra_metrics` are
``callable(y_true, y_pred)``; Trainer's validation calls them on the
aggregate). The custom `trainer_step_fn` owns every optimizer update;
by default Trainer steps every scheduler once after each epoch (an
``optimizer_update``-clock scheduler after each reported update) — set
`auto_step_schedulers=False` when the custom step function owns
scheduler timing too.
`optims`, `schedulers` and `extra_metrics` are captured as read-only,
insertion-ordered snapshots at construction (also via `from_state`,
`dataclasses.replace` and the `with_*_loader` helpers): mutating the
mapping you passed in — or reusing the builder that produced this
configuration — never changes it, and item assignment / deletion on
the exposed mappings raises `TypeError`. Values are shared by
identity (immutable params dataclasses; runtime-only metric callables
and loaders are never copied).
nnx.trainer.params.NNTrainerParams.with_train_loader¶
nnx.trainer.params.NNTrainerParams.with_train_loader(self, value: 'DataLoader') -> 'NNTrainerParams'
No public description is currently available.
nnx.trainer.params.NNTrainerParams.with_val_loader¶
nnx.trainer.params.NNTrainerParams.with_val_loader(self, value: 'DataLoader') -> 'NNTrainerParams'
No public description is currently available.
nnx.trainer.params.NNTrainerParams.state¶
nnx.trainer.params.NNTrainerParams.state(self)
No public description is currently available.
nnx.trainer.params.NNTrainerParams.from_state¶
nnx.trainer.params.NNTrainerParams.from_state(state: 'dict') -> 'NNTrainerParams'
No public description is currently available.
nnx.trainer.params.NNTrainerParams.builder¶
nnx.trainer.params.NNTrainerParams.builder() -> 'NNTrainerParamsBuilder'
Return a composite multi-optim builder. See NNTrainerParamsBuilder. Composes NNOptimParams.builder() + NNSchedulerParams.builder().
nnx.trainer.params_builder.NNTrainerParamsBuilder¶
class nnx.trainer.params_builder.NNTrainerParamsBuilder() -> 'None'
Composite builder for NNTrainerParams.
Details
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.
Ownership (see `NNTrainerParams`): the custom ``trainer_step_fn``
owns every optimizer update; Trainer steps registered schedulers once
per epoch (an ``optimizer_update``-clock scheduler after each reported
update instead) unless ``.auto_step_schedulers(False)``.
`copy()` branches a (possibly partial) builder and `from_params()`
rebuilds one from an existing `NNTrainerParams`, so a shared base
(epochs, loaders, a common optimizer) is written once and each branch
adds its own optimizers, schedulers or `data_id`.
nnx.trainer.params_builder.NNTrainerParamsBuilder.copy¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.copy(self) -> 'NNTrainerParamsBuilder'
Return an independent branch of this builder, complete or partial.
Details
The `optims` / `schedulers` maps and the `extra_metrics` mapping
are copied, so registering, replacing or dropping an entry on one
builder never shows up on the other (nor in values already
built). Their values are shared by identity: optimizer and
scheduler params are immutable, and `DataLoader` s and metric
callables are runtime objects that are never copied or iterated.
Nothing is built or validated here — `build()` validates each
branch, including the `schedulers ⊆ optims` rule.
nnx.trainer.params_builder.NNTrainerParamsBuilder.from_params¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.from_params(params: 'NNTrainerParams') -> 'NNTrainerParamsBuilder'
Return a builder pre-loaded with every field of params.
Details
`from_params(params).build()` reproduces `params` field for field,
with the same `state()` (sorted keys, omitted defaults) and the
same runtime insertion order of `optims`. Built-in and
registered-factory optimizer params, loaders and metric callables
keep their identity; the builder's own maps are fresh copies, so
extending it never alters `params`.
Raises:
TypeError: if `params` is not exactly an `NNTrainerParams`
(the single-optimizer `NNTrainParams` has no builder).
ValueError: if `params` carries a field this builder cannot
reproduce.
nnx.trainer.params_builder.NNTrainerParamsBuilder.n_epochs¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.n_epochs(self, n: 'int') -> 'NNTrainerParamsBuilder'
Number of training epochs. Required.
nnx.trainer.params_builder.NNTrainerParamsBuilder.optimizer¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.optimizer(self, name: 'str', params: 'Union[NNOptimParams, NNOptimFactoryParams]') -> 'NNTrainerParamsBuilder'
Register one optimizer under name. Each name gets its own torch.optim.Optimizer at Trainer.train() time. Use NNOptimParams.builder() to construct a built-in params, or nnx.NNOptimFactoryParams for a registered factory.
nnx.trainer.params_builder.NNTrainerParamsBuilder.scheduler¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.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.
nnx.trainer.params_builder.NNTrainerParamsBuilder.seed¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.seed(self, value: 'int') -> 'NNTrainerParamsBuilder'
Seed for reproducibility. None at default (no seeding via params; the caller's set_seed() is the only path).
nnx.trainer.params_builder.NNTrainerParamsBuilder.data_id¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.data_id(self, value: 'str') -> 'NNTrainerParamsBuilder'
Name the dataset or split this configuration trains on.
Details
Part of `state()` and therefore of the run id, so two otherwise
identical configurations (e.g. sibling builder branches) over
different data get distinct run directories. None at default.
nnx.trainer.params_builder.NNTrainerParamsBuilder.metrics¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.metrics(self, *specs: 'MetricSpec') -> 'NNTrainerParamsBuilder'
Declare named, registered metrics — computed over the whole validation set by evaluate() and reported under each spec's name. Replaces any earlier declaration.
nnx.trainer.params_builder.NNTrainerParamsBuilder.monitor¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.monitor(self, spec: 'MonitorSpec') -> 'NNTrainerParamsBuilder'
What BEST selection and plateau schedulers track.
nnx.trainer.params_builder.NNTrainerParamsBuilder.resume_from¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.resume_from(self, run_id: 'str', checkpoint: 'str' = 'last', mode: 'str' = 'auto') -> 'NNTrainerParamsBuilder'
Warm-resume from run_id's checkpoint: the model, every named optimizer and scheduler, the RNG and the registered components (callbacks such as EarlyStopping) continue where that run stopped. mode is "auto" (stateful when the checkpoint has training state, else weights-only), "stateful" (fail unless it has) or "weights_only". Serialized only as parent lineage, so the resumed run gets its own id.
nnx.trainer.params_builder.NNTrainerParamsBuilder.overwrite_existing¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.overwrite_existing(self, value: 'bool') -> 'NNTrainerParamsBuilder'
Allow a run to replace an existing run directory with the same id. Default False; not part of state() or the run id.
nnx.trainer.params_builder.NNTrainerParamsBuilder.save_phase_checkpoints¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.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).
nnx.trainer.params_builder.NNTrainerParamsBuilder.auto_step_schedulers¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.auto_step_schedulers(self, value: 'bool') -> 'NNTrainerParamsBuilder'
Choose whether Trainer steps every scheduler after each epoch (an optimizer_update-clock scheduler after each reported update).
Details
Disable this when the custom step function owns scheduler timing.
nnx.trainer.params_builder.NNTrainerParamsBuilder.train_loader¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.train_loader(self, loader: 'DataLoader') -> 'NNTrainerParamsBuilder'
Training DataLoader. Optional at Builder time (can be wired later via NNTrainerParams.with_train_loader).
nnx.trainer.params_builder.NNTrainerParamsBuilder.val_loader¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.val_loader(self, loader: 'DataLoader') -> 'NNTrainerParamsBuilder'
Validation DataLoader. Optional at Builder time (can be wired later via NNTrainerParams.with_val_loader).
nnx.trainer.params_builder.NNTrainerParamsBuilder.extra_metrics¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.extra_metrics(self, metrics: 'Mapping[str, Callable]') -> 'NNTrainerParamsBuilder'
Extra metrics, name-keyed callable(y_true, y_pred) -> float (truth first, decoded predictions second). Trainer's built-in validation calls each once on the aggregate predictions (NNModel.evaluate); the custom trainer_step_fn decides whether and how to call them on training batches.
nnx.trainer.params_builder.NNTrainerParamsBuilder.build¶
nnx.trainer.params_builder.NNTrainerParamsBuilder.build(self) -> 'NNTrainerParams'
Validate the key-subset invariant, then construct the dataclass.
Details
`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.
2.4. Probability-aware prediction (nnx.prediction)¶
nnx.prediction.ProbabilitySpec¶
class nnx.prediction.ProbabilitySpec(kind: 'str', class_axis: 'int' = -1, labels: 'Optional[tuple[str, ...]]' = None) -> 'None'
Explicit declaration of how raw logits become probabilities.
Details
Args:
kind: ``"categorical"`` (softmax over ``class_axis``, argmax
decoding) or ``"bernoulli"`` (independent sigmoid per output,
``logit >= 0`` decoding).
class_axis: axis of the logits that indexes classes / outputs.
Negative values count from the end; ``-1`` (default) fits
``(N, C)`` classifier logits and class-last ``(B, T, V)``
language-model logits, while class-first outputs such as
``(N, C, H, W)`` need ``class_axis=1``. Axis 0 is always the
sample axis and cannot be the class axis.
labels: optional ordered class / output names, one per entry
along ``class_axis``; unique, non-empty strings.
Serializes with :meth:`state` / :meth:`from_state` (plain YAML-safe
types, label order preserved).
nnx.prediction.ProbabilitySpec.resolve_class_axis¶
nnx.prediction.ProbabilitySpec.resolve_class_axis(self, ndim: 'int') -> 'int'
Return class_axis as a non-negative axis of an ndim-D logits array, or raise if it is out of range or the sample axis.
nnx.prediction.ProbabilitySpec.state¶
nnx.prediction.ProbabilitySpec.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.prediction.ProbabilitySpec.from_state¶
nnx.prediction.ProbabilitySpec.from_state(state: 'Mapping[str, Any]') -> 'ProbabilitySpec'
No public description is currently available.
nnx.prediction.PredictionResult¶
class nnx.prediction.PredictionResult(logits: 'np.ndarray', probabilities: 'Optional[np.ndarray]', decoded: 'np.ndarray', sample_ids: 'np.ndarray', spec: 'Optional[ProbabilitySpec]') -> 'None'
Probability-aware prediction for N samples.
Details
Attributes:
logits: raw network output, exactly what ``predict().logits``
returns for the same input.
probabilities: same shape as ``logits``; softmax over the class
axis (categorical, rows sum to 1) or element-wise sigmoid
(bernoulli, not normalized). ``None`` for a continuous
(regression-task) result, which has no probabilities.
decoded: categorical — argmax class indices (``logits`` without
the class axis); bernoulli — 0/1 indicators shaped like
``logits`` (``logit >= 0``, or the task's threshold);
continuous — the predicted values themselves.
sample_ids: ``int64[N]`` identity of each row: the input row index
for arrays, tensors and ordinary loaders, and the global node
index (``input_id``) for seed rows of a graph loader.
spec: the :class:`ProbabilitySpec` that produced the result, or
``None`` for a continuous result (``NNModel.predict_proba``
on a regression task).
nnx.prediction.PredictionResult.kind¶
property nnx.prediction.PredictionResult.kind
"categorical", "bernoulli" or "continuous" (no spec).
nnx.prediction.PredictionResult.class_axis¶
property nnx.prediction.PredictionResult.class_axis
The class axis of logits / probabilities, non-negative.
nnx.prediction.PredictionResult.labels¶
property nnx.prediction.PredictionResult.labels
No public description is currently available.
nnx.prediction.PredictionResult.class_indices¶
property nnx.prediction.PredictionResult.class_indices
Decoded class indices — categorical results only.
Details
Bernoulli outputs are independent per-output indicators, not
mutually exclusive classes, so they are refused here (and must
never be fed to class-index consumers such as
``VisUtils.confusion_matrix``).
nnx.prediction.PredictionResult.decoded_labels¶
nnx.prediction.PredictionResult.decoded_labels(self) -> 'np.ndarray'
Decoded class names (categorical results with spec.labels).
nnx.prediction.prediction_from_logits¶
nnx.prediction.prediction_from_logits(logits: 'Any', spec: 'ProbabilitySpec', *, sample_ids: 'Any' = None) -> 'PredictionResult'
Build a PredictionResult from raw logits.
Details
``logits`` is an array or tensor whose axis 0 indexes samples and whose
``spec.class_axis`` indexes classes / outputs. Raises
:class:`PredictionValidationError` when the class axis is invalid, the
label count differs from the class-axis size, any logit is NaN or
±inf (masked logits must use a large finite negative value), or
``sample_ids`` is not one id per sample. Probabilities are computed
stably (max-shifted softmax, ``exp(-|x|)`` sigmoid) in the logits'
float precision and returned in their dtype — float16 is computed in
float32 and returned as float16, bfloat16 tensors are upcast to
float32, integer logits give float64.
nnx.prediction.PredictionValidationError¶
class nnx.prediction.PredictionValidationError
An invalid ProbabilitySpec, or logits / sample ids that do not satisfy one (wrong label count, invalid class axis, non-finite logits, misaligned sample ids).
nnx.prediction.NonFiniteLogitsError¶
class nnx.prediction.NonFiniteLogitsError
The logits hold NaN or +-inf: the model's output, not the request, is at fault (a decision provider reports it as a provider failure).
2.5. Optimizer construction and registered factories (nnx.optimizers)¶
nnx.optimizers.build_optimizer¶
nnx.optimizers.build_optimizer(net: 'nn.Module', params: 'AnyOptimParams', *, strict_param_groups: 'bool' = False) -> 'torch.optim.Optimizer'
Build the optimizer params describes for net.
Details
The one construction hook behind ``NNModel.train`` (non-strict: an
unmatched trainable parameter joins a default group) and
``Trainer.train`` (``strict_param_groups=True``: each optimizer owns
only what its ``param_groups`` specs select, so disjoint optimizers
never co-own a parameter).
Built-in :class:`NNOptimParams` dispatch through
:class:`~nnx.nn.enum.optims.Optims` exactly as before. For
:class:`NNOptimFactoryParams` the factory is resolved in the registry,
called once with the groups a built-in would receive (the shared
:func:`~nnx.nn.enum.optims.resolve_param_groups` rule, each group with
explicit ``lr`` / ``weight_decay``) and the read-only config, and its
result is checked: it must be a ``torch.optim.Optimizer`` holding
exactly those parameters.
nnx.optimizers.OptimizerFactory¶
type alias nnx.optimizers.OptimizerFactory
Public type alias.
nnx.optimizers.OptimizerFactorySpec¶
class nnx.optimizers.OptimizerFactorySpec(id: 'str', version: 'int', config: 'Optional[Mapping[str, Any]]' = None) -> 'None'
Stable, serializable reference to a registered optimizer factory.
Details
``id`` and ``version`` name the factory in the registry; ``config`` is
handed to it read-only and must be JSON-like — ``None``, ``bool``,
``int``, finite ``float``, ``str``, lists (exposed as tuples) and
string-keyed mappings. The spec never holds the callable, so it
round-trips through ``run.yaml`` and reloads without the factory being
registered or imported. Immutable; equal specs have equal ``state()``.
nnx.optimizers.OptimizerFactorySpec.state¶
nnx.optimizers.OptimizerFactorySpec.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.optimizers.OptimizerFactorySpec.from_state¶
nnx.optimizers.OptimizerFactorySpec.from_state(state: 'Mapping[str, Any]') -> 'OptimizerFactorySpec'
No public description is currently available.
nnx.optimizers.register_optimizer_factory¶
nnx.optimizers.register_optimizer_factory(id: 'str', version: 'int', factory: 'OptimizerFactory', *, replace: 'bool' = False) -> 'None'
Register factory under (id, version) for this process.
Details
``factory(param_groups, config)`` receives the resolved parameter
groups and the spec's read-only config and must return a
``torch.optim.Optimizer`` over exactly those parameters. Registering an
existing ``(id, version)`` raises unless ``replace=True``; bump
``version`` instead when the factory's behaviour changes, so resumed
runs can tell the difference.
nnx.optimizers.unregister_optimizer_factory¶
nnx.optimizers.unregister_optimizer_factory(id: 'str', version: 'int') -> 'bool'
Remove (id, version) from the registry; returns whether it was registered. Run metadata that references it still loads — only training with it needs the registration.
nnx.optimizers.registered_optimizer_factories¶
nnx.optimizers.registered_optimizer_factories() -> 'tuple[tuple[str, int], ...]'
Sorted (id, version) pairs currently registered.
nnx.optimizers.resolve_optimizer_factory¶
nnx.optimizers.resolve_optimizer_factory(spec: 'OptimizerFactorySpec') -> 'OptimizerFactory'
Look spec up in the registry. Never imports anything: an unregistered id or version raises ValueError naming what is registered.
nnx.optimizers.optim_params_from_state¶
nnx.optimizers.optim_params_from_state(state: 'Mapping[str, Any]') -> 'AnyOptimParams'
Decode a serialized optimizer config into the right variant.
Details
A ``factory`` entry is the discriminator for
:class:`NNOptimFactoryParams`; anything else is a built-in
:class:`NNOptimParams` (every config without a ``factory``
entry). Decoding never resolves or runs a factory.
2.6. Task adapters (nnx.tasks)¶
nnx.tasks¶
module nnx.tasks
Task adapters: declare what a supervised model's outputs and targets mean.
Details
By default ``NNModel`` assumes categorical classification: the default
training step, :meth:`NNModel.evaluate` and :meth:`NNModel.predict` decode
by argmax (or the ``BCEWithLogitsLoss`` threshold) and every evaluation
record carries accuracy / f1 / recall / precision. A :class:`TaskSpec` on
``NNModelParams(task=...)`` opts a model into an explicit task instead::
params = NNModelParams(net=Nets.FEED_FWD, loss=Losses.MEAN_SQUARED_ERROR, task=TaskSpec.regression(2))
The task's adapter then owns, for the default step, ``evaluate()``,
``predict()`` and ``predict_proba()``:
* **validation** — each batch's outputs and targets are checked for the
declared shape, dtype and value range *before* any backward pass or
optimizer update (:class:`TaskValidationError`);
* **masking** — categorical targets equal to ``ignore_index``, and NaN
multilabel / regression targets, are excluded from the loss and every
metric alike;
* **loss units** — the loss is averaged over the valid targets, with the
same denominators the default step already uses for gradient
accumulation, so uneven batches and accumulated windows reduce exactly
like one full batch;
* **decoding and metrics** — argmax and accuracy / macro f1 / recall /
precision (``categorical``), independent sigmoid probabilities with a
declared threshold plus subset and element accuracy (``multilabel``),
continuous values with MSE / MAE and no classification fields
(``regression``).
Three kinds exist; the kind is always declared, never inferred:
* ``"categorical"`` — mutually exclusive classes along axis 1 of the
output (``(N, C)`` or ``(N, C, ...)``); integer targets shaped like the
output without the class axis. Needs ``Losses.CROSS_ENTROPY`` or
``Losses.NEGATIVE_LOG_LIKELIHOOD``.
* ``"multilabel"`` — independent binary labels: ``(N, L)`` logits and
0/1 targets of the same shape (NaN = masked). Needs
``Losses.BINARY_CROSS_ENTROPY`` (``BCEWithLogitsLoss``).
* ``"regression"`` — continuous targets shaped like the output (a
``(N,)`` target is accepted for an ``(N, 1)`` output; NaN = masked).
Needs a regression loss (``Losses.MEAN_SQUARED_ERROR``, or any
non-classification loss module assigned to ``model.loss_fn``).
Language-model, ranking, link-prediction and structured-prediction tasks
are out of scope; ``Nets.TRANSFORMER`` and ``Nets.VIT`` models are rejected. Models without
a task keep the legacy classification behaviour and serialization.
nnx.tasks.TaskSpec¶
class nnx.tasks.TaskSpec(kind: 'str', num_outputs: 'Optional[int]' = None, ignore_index: 'Optional[int]' = None, threshold: 'float' = 0.5, labels: 'Optional[tuple[str, ...]]' = None) -> 'None'
Declaration of a supervised task (see the module docstring).
Details
Prefer the constructors :meth:`categorical`, :meth:`multilabel` and
:meth:`regression`.
Args:
kind: ``"categorical"``, ``"multilabel"`` or ``"regression"``.
num_outputs: classes / labels / targets along output axis 1, or
``None`` to accept the net's width. When set, a model whose
``output_dim`` differs is rejected at construction.
ignore_index: categorical only — the target value excluded from
the loss and every metric (e.g. ``-100`` for padding).
threshold: multilabel only — probability at or above which a label
is decoded as positive (default ``0.5``, i.e. ``logit >= 0``).
labels: optional ordered names, one per output.
Serializes with :meth:`state` / :meth:`from_state` as a versioned
mapping (``version`` = :data:`TASK_SPEC_VERSION`).
nnx.tasks.TaskSpec.categorical¶
nnx.tasks.TaskSpec.categorical(num_classes: 'Optional[int]' = None, *, ignore_index: 'Optional[int]' = None, labels: 'Optional[Sequence[str]]' = None) -> 'TaskSpec'
Mutually exclusive classes along output axis 1.
nnx.tasks.TaskSpec.multilabel¶
nnx.tasks.TaskSpec.multilabel(num_labels: 'Optional[int]' = None, *, threshold: 'float' = 0.5, labels: 'Optional[Sequence[str]]' = None) -> 'TaskSpec'
Independent binary labels, decoded at probability >= threshold.
nnx.tasks.TaskSpec.regression¶
nnx.tasks.TaskSpec.regression(num_targets: 'Optional[int]' = None, *, labels: 'Optional[Sequence[str]]' = None) -> 'TaskSpec'
Continuous targets shaped like the output.
nnx.tasks.TaskSpec.state¶
nnx.tasks.TaskSpec.state(self) -> 'dict[str, Any]'
Versioned, YAML-safe representation (defaults omitted).
nnx.tasks.TaskSpec.from_state¶
nnx.tasks.TaskSpec.from_state(state: 'Mapping[str, Any]') -> 'TaskSpec'
Rebuild a spec from state, rejecting an unknown version or key.
nnx.tasks.TaskSpec.probability_spec¶
nnx.tasks.TaskSpec.probability_spec(self) -> 'Optional[ProbabilitySpec]'
The ProbabilitySpec this task implies: categorical softmax or multilabel (bernoulli) sigmoid over axis 1, or None for regression (no probabilities).
nnx.tasks.TaskAdapter¶
class nnx.tasks.TaskAdapter(spec: 'TaskSpec') -> 'None'
Validation, masking, loss units, decoding and metrics for one TaskSpec kind. Obtain one with task_adapter.
nnx.tasks.TaskAdapter.check_model¶
nnx.tasks.TaskAdapter.check_model(self, *, net: 'Any', loss: 'Losses', output_dim: 'Optional[int]') -> 'None'
Reject a model configuration this task cannot drive — before any net is built or loader is iterated. net is a Nets member or a registered / runtime descriptor, whose output width is unknown until the first batch (output_dim=None).
nnx.tasks.TaskAdapter.check_loss_fn¶
nnx.tasks.TaskAdapter.check_loss_fn(self, loss_fn: 'torch.nn.Module') -> 'None'
Reject a runtime model.loss_fn that cannot score this task.
nnx.tasks.TaskAdapter.prepare¶
nnx.tasks.TaskAdapter.prepare(self, output: 'torch.Tensor', target: 'torch.Tensor') -> 'tuple[torch.Tensor, torch.Tensor, torch.Tensor]'
Validate one batch and return (output, target, valid): tensors shaped for the loss plus a boolean mask of the targets that count. Raises TaskValidationError before any backward pass for a shape, dtype or value the task rejects.
nnx.tasks.TaskAdapter.loss_terms¶
nnx.tasks.TaskAdapter.loss_terms(self, loss_fn: 'torch.nn.Module', output: 'torch.Tensor', target: 'torch.Tensor', valid: 'torch.Tensor') -> 'tuple[torch.Tensor, torch.Tensor, Optional[float]]'
(display loss, additive numerator, normalization weight) over the valid targets only — the same contract as the default step's loss terms, so accumulation windows normalize once by their valid count. An all-masked batch contributes a differentiable zero with weight 0 and a NaN display loss.
nnx.tasks.TaskAdapter.decode¶
nnx.tasks.TaskAdapter.decode(self, output: 'torch.Tensor') -> 'torch.Tensor'
Decoded predictions for a prepared output.
nnx.tasks.TaskAdapter.decode_array¶
nnx.tasks.TaskAdapter.decode_array(self, logits: 'np.ndarray') -> 'np.ndarray'
Decoded predictions for raw predict() logits (numpy).
nnx.tasks.TaskAdapter.accumulator¶
nnx.tasks.TaskAdapter.accumulator(self, *, keep_arrays: 'bool' = False, bounded: 'bool' = False) -> 'TaskMetricAccumulator'
A fresh TaskMetricAccumulator; pass keep_arrays=True when extra_metrics will be computed from it. bounded=True keeps only counts and sums, so memory does not grow with the number of samples; it cannot keep arrays.
nnx.tasks.TaskAdapter.record¶
nnx.tasks.TaskAdapter.record(self, output: 'torch.Tensor', target: 'torch.Tensor', *, loss: 'Optional[float]', extra_metrics: 'Optional[Mapping[str, Callable]]' = None) -> 'NNEvaluationDataPoint'
Validate one batch and return its task record — for custom training steps that want the same record the default step writes.
nnx.tasks.TaskAdapter.check_logits¶
nnx.tasks.TaskAdapter.check_logits(self, logits: 'np.ndarray') -> 'None'
Shape check for raw prediction logits (a loader's first batch), so a mismatched net fails before the rest of inference runs.
nnx.tasks.TaskAdapter.prediction¶
nnx.tasks.TaskAdapter.prediction(self, logits: 'np.ndarray', sample_ids: 'np.ndarray') -> 'PredictionResult'
The rich prediction this task implies (see predict_proba).
nnx.tasks.TaskMetricAccumulator¶
class nnx.tasks.TaskMetricAccumulator(adapter: 'TaskAdapter', *, keep_arrays: 'bool' = False, bounded: 'bool' = False) -> 'None'
Mergeable per-task statistics over any number of batches.
Details
``update`` takes one prepared batch; ``result`` builds the
:class:`~nnx.NNEvaluationDataPoint` for everything seen so far.
Metrics are computed over the whole accumulation (not averaged per
batch), so uneven batches give the same numbers as one full batch.
``keep_arrays`` retains the valid targets and predictions for
user ``extra_metrics``; it is off unless those are requested.
nnx.tasks.TaskMetricAccumulator.update¶
nnx.tasks.TaskMetricAccumulator.update(self, output: 'torch.Tensor', target: 'torch.Tensor', valid: 'torch.Tensor') -> 'None'
No public description is currently available.
nnx.tasks.TaskMetricAccumulator.result¶
nnx.tasks.TaskMetricAccumulator.result(self, *, loss: 'Optional[float]', extra_metrics: 'Optional[Mapping[str, Callable]]' = None) -> 'NNEvaluationDataPoint'
No public description is currently available.
nnx.tasks.TaskValidationError¶
class nnx.tasks.TaskValidationError
An invalid TaskSpec, a model the task cannot drive, or a batch whose outputs / targets do not satisfy the declared task.
nnx.tasks.task_adapter¶
nnx.tasks.task_adapter(spec: 'TaskSpec') -> 'TaskAdapter'
Return the adapter that drives spec.
2.7. Checkpointable component state (nnx.components)¶
nnx.components.ComponentSpec¶
class nnx.components.ComponentSpec(name: 'str', version: 'int' = 1, required: 'bool' = True, numbered: 'bool' = False) -> 'None'
Identity and schema version of a checkpointable component.
Details
Args:
name: unique within a training run — a filename-safe slug
(letters, digits, ``.``, ``_``, ``:``, ``-``).
version: schema version of ``component_state()`` (a positive
integer). A checkpoint written by a *newer* schema than the
component supports is rejected before anything is restored.
required: when ``True`` (default) a resume from a stateful
checkpoint that lacks this component's state fails; an optional
component simply keeps its fresh state.
numbered: ``name`` is a default the registry may number when several
components share it — ``name``, ``name.2``, ``name.3`` … in
registration order — instead of rejecting the duplicate
(built-in callbacks such as ``EarlyStopping`` use it, so two of
them in one run still work). A resume then maps state by that
order. Explicit names leave it ``False`` and must be unique.
nnx.components.StatefulComponent¶
class nnx.components.StatefulComponent(*args, **kwargs)
The contract a checkpointable component implements.
Details
A component may also define ``check_component_state(state, *, version)
-> Iterable[str]``: the registry calls it while validating a resume,
before anything is mutated, and reports every string it returns (for
example a configuration that no longer matches the saved state) in the
same :class:`ComponentRestoreError` as the other problems.
nnx.components.StatefulComponent.component_spec¶
nnx.components.StatefulComponent.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.components.StatefulComponent.component_state¶
nnx.components.StatefulComponent.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.components.StatefulComponent.load_component_state¶
nnx.components.StatefulComponent.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.components.ComponentRegistry¶
class nnx.components.ComponentRegistry(components: 'Iterable[Any]' = ()) -> 'None'
Ordered, uniquely named set of the components of one training run.
nnx.components.ComponentRegistry.discover¶
nnx.components.ComponentRegistry.discover(*sources: 'Any', explicit: 'Iterable[Any]' = ()) -> 'ComponentRegistry'
Register every source that implements StatefulComponent (callbacks, step functions), in order — None and lists are flattened, anything else is ignored — then every explicit component, which must implement the protocol (TypeError otherwise). An object that appears more than once (say, a callback also passed through components=[...]) is registered once.
nnx.components.ComponentRegistry.register¶
nnx.components.ComponentRegistry.register(self, component: 'Any') -> 'str'
Register component and return the name its state is saved under. Registering the same object again is a no-op; a different object with a taken name raises ValueError unless its spec is numbered.
nnx.components.ComponentRegistry.names¶
property nnx.components.ComponentRegistry.names
No public description is currently available.
nnx.components.ComponentRegistry.collect¶
nnx.components.ComponentRegistry.collect(self) -> 'dict[str, dict[str, Any]]'
Every component's versioned state, for the training-state sidecar.
nnx.components.ComponentRegistry.fresh_plan¶
nnx.components.ComponentRegistry.fresh_plan(self) -> '_Plan'
A plan that restores nothing: every component keeps its fresh state (weights-only resumes, checkpoints without component state).
nnx.components.ComponentRegistry.plan¶
nnx.components.ComponentRegistry.plan(self, saved: 'Optional[Mapping[str, Any]]') -> '_Plan'
Validate saved component metadata against the registered components without mutating anything; raise one ComponentRestoreError listing every problem.
nnx.components.ComponentRegistry.restore¶
nnx.components.ComponentRegistry.restore(self, plan: '_Plan') -> 'tuple[str, ...]'
Load a validated plan transactionally: on any failure every component attempted — the failing one included — gets a deep copy of its pre-call state back, then the original error propagates. A component whose own rollback fails does not stop the others from being rolled back, nor replace the original error; it is reported in a RuntimeWarning.
nnx.components.ComponentRestoreError¶
class nnx.components.ComponentRestoreError(problems: 'Iterable[str]') -> 'None'
Saved component state that cannot be restored: every problem found (missing, unknown or incompatible components) is listed in problems, and nothing has been mutated when it is raised by validation.
nnx.components.ResumeStatus¶
class nnx.components.ResumeStatus(mode: 'str' = 'fresh', source_run_id: 'Optional[str]' = None, source_checkpoint: 'Optional[str]' = None, restored_components: 'tuple[str, ...]' = (), fresh_components: 'tuple[str, ...]' = (), source_epoch: 'Optional[int]' = None, source_generation: 'Optional[int]' = None, source_checkpoint_id: 'Optional[str]' = None) -> 'None'
How a training session started.
Details
Attributes:
mode: ``"fresh"`` (no resume), ``"stateful"`` (model plus the
complete training-state bundle, components included) or
``"weights_only"`` (model weights only — the checkpoint had no
training state, or ``resume_mode="weights_only"`` was asked).
source_run_id / source_checkpoint: where the resume came from.
source_epoch: the epoch of that checkpoint when the session resumed
from it (the history journal's lineage reads it); ``None`` for a
fresh run and for status recorded without it.
restored_components: names of the components whose state was
restored, in registration order.
fresh_components: registered components that kept their fresh
state (optional components absent from the checkpoint, or every
component on a weights-only resume).
source_generation / source_checkpoint_id: the resume point's
generation ordinal and checkpoint id from its manifest;
``None`` for a checkpoint without a manifest.
nnx.components.ResumeStatus.state¶
nnx.components.ResumeStatus.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.components.ResumeStatus.from_state¶
nnx.components.ResumeStatus.from_state(state: 'Mapping[str, Any]') -> 'ResumeStatus'
No public description is currently available.
2.8. Named metrics and monitors (nnx.monitors)¶
nnx.monitors.MetricSpec¶
class nnx.monitors.MetricSpec(id: 'str', version: 'int' = 1, config: 'Optional[Mapping[str, Any]]' = None, name: 'Optional[str]' = None) -> 'None'
Serializable declaration of a registered metric.
Details
Args:
id / version: the registration to use.
config: JSON-like options handed to the metric read-only
(``None``, ``bool``, ``int``, finite ``float``, ``str``, lists
and string-keyed mappings).
name: the key the value is reported under (default: ``id``); a slug,
unique among a run's metrics, and never ``"loss"`` / ``"error"``.
A name equal to a task's own metric (``mse``, ``subset_accuracy``)
replaces that value in the record's ``metrics``; loggers report
one equal to a record field (``accuracy``, ``f1``) as
``metrics/<name>``.
The spec never holds code, so it round-trips through ``run.yaml`` and
reloads without the metric being registered. It is resolved against
the registry when training starts: an unknown ``(id, version)`` then
fails before anything is trained, and no metric code has run.
nnx.monitors.MetricSpec.label¶
property nnx.monitors.MetricSpec.label
The key the metric is reported under.
nnx.monitors.MetricSpec.definition¶
property nnx.monitors.MetricSpec.definition
The registration this spec names (ValueError when unknown).
nnx.monitors.MetricSpec.input¶
property nnx.monitors.MetricSpec.input
No public description is currently available.
nnx.monitors.MetricSpec.mode¶
property nnx.monitors.MetricSpec.mode
No public description is currently available.
nnx.monitors.MetricSpec.check¶
nnx.monitors.MetricSpec.check(self) -> 'None'
Resolve the registration and validate the config, without computing anything.
nnx.monitors.MetricSpec.accumulator¶
nnx.monitors.MetricSpec.accumulator(self) -> 'MetricAccumulator'
No public description is currently available.
nnx.monitors.MetricSpec.state¶
nnx.monitors.MetricSpec.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.monitors.MetricSpec.from_state¶
nnx.monitors.MetricSpec.from_state(state: 'Mapping[str, Any]') -> 'MetricSpec'
No public description is currently available.
nnx.monitors.MetricDefinition¶
class nnx.monitors.MetricDefinition(id: 'str', version: 'int', input: 'str', mode: 'str', factory: 'Callable[[Mapping[str, Any]], MetricAccumulator]', check_config: 'Optional[Callable[[Mapping[str, Any]], None]]' = None) -> 'None'
A registered metric: what it receives and how it is computed.
nnx.monitors.MonitorSpec¶
class nnx.monitors.MonitorSpec(metric: 'str' = 'loss', split: 'str' = 'val', mode: 'Optional[str]' = None, min_delta: 'float' = 0.0, on_missing: 'str' = 'skip', on_nonfinite: 'str' = 'skip') -> 'None'
What BEST selection, early stopping and plateau scheduling track.
Details
Args:
metric: ``"loss"``, ``"error"`` or the name of a declared
:class:`MetricSpec` (``NNTrainParams.metrics``).
split: ``"val"`` (default) — the whole-validation-set record — or
``"train"`` — the whole-epoch training summary (full-epoch
denominators, never the last batch).
mode: ``"min"`` / ``"max"``; ``None`` (default) takes the metric's
natural direction (``"min"`` for loss and error).
min_delta: an improvement must beat the best by more than this
(finite, ``>= 0``); ties never improve.
on_missing: ``"skip"`` (default) — an epoch without a value makes no
decision (not counted toward patience, no plateau step, no
BEST) — or ``"error"``: raise :class:`MonitorUnavailableError`.
on_nonfinite: ``"skip"`` (default) — a NaN / ±inf value is an epoch
without improvement — or ``"error"``.
nnx.monitors.MonitorSpec.key¶
property nnx.monitors.MonitorSpec.key
"<split>.<metric>", as shown by displays and loggers.
nnx.monitors.MonitorSpec.check_names¶
nnx.monitors.MonitorSpec.check_names(self, metrics: 'Sequence[MetricSpec]', owner: 'str' = 'monitor') -> 'None'
Fail when the metric is neither a record field nor declared.
nnx.monitors.MonitorSpec.resolve¶
nnx.monitors.MonitorSpec.resolve(self, metrics: 'Sequence[MetricSpec]', owner: 'str' = 'monitor') -> 'MonitorSpec'
Check the name against the declared metrics and fill in the metric's natural direction when mode is None.
nnx.monitors.MonitorSpec.improved¶
nnx.monitors.MonitorSpec.improved(self, current: 'float', best: 'Optional[float]') -> 'bool'
The shared rule: the first finite value improves; later values must beat best by more than min_delta (ties never improve); a non-finite value never improves.
nnx.monitors.MonitorSpec.value¶
nnx.monitors.MonitorSpec.value(self, *, train: 'Optional[NNEvaluationDataPoint]', val: 'Optional[NNEvaluationDataPoint]') -> 'Optional[float]'
The monitored value from an epoch's training summary / validation record (None when absent).
nnx.monitors.MonitorSpec.state¶
nnx.monitors.MonitorSpec.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.monitors.MonitorSpec.from_state¶
nnx.monitors.MonitorSpec.from_state(state: 'Mapping[str, Any]') -> 'MonitorSpec'
No public description is currently available.
nnx.monitors.MonitorRecord¶
class nnx.monitors.MonitorRecord(monitor: 'MonitorSpec', value: 'Optional[float]', status: 'str', improved: 'bool') -> 'None'
One epoch's monitor decision, persisted with the epoch's record.
Details
``status`` is ``"ok"``, ``"missing"`` (no value; no decision) or
``"nonfinite"`` (an epoch without improvement); ``improved`` says
whether this epoch became the best.
nnx.monitors.MonitorRecord.state¶
nnx.monitors.MonitorRecord.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.monitors.MonitorRecord.from_state¶
nnx.monitors.MonitorRecord.from_state(state: 'Mapping[str, Any]') -> 'MonitorRecord'
No public description is currently available.
nnx.monitors.MonitorTracker¶
class nnx.monitors.MonitorTracker(spec: 'MonitorSpec', *, warn_missing: 'bool' = False) -> 'None'
Applies one MonitorSpec epoch by epoch: returns each epoch's MonitorRecord and remembers the best value.
Details
The run's tracker is also an optional checkpointable component
(``nnx.monitor``): a stateful warm resume continues from the
source run's best, so BEST keeps agreeing with a restored
``EarlyStopping`` and plateau scheduler after the split — the resumed
run writes BEST only on an epoch that beats it, and none when no
resumed epoch does.
nnx.monitors.MonitorTracker.observe¶
nnx.monitors.MonitorTracker.observe(self, value: 'Optional[float]', *, epoch: 'int') -> 'MonitorRecord'
No public description is currently available.
nnx.monitors.MonitorTracker.component_spec¶
nnx.monitors.MonitorTracker.component_spec(self) -> 'Any'
No public description is currently available.
nnx.monitors.MonitorTracker.component_state¶
nnx.monitors.MonitorTracker.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.monitors.MonitorTracker.check_component_state¶
nnx.monitors.MonitorTracker.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.monitors.MonitorTracker.load_component_state¶
nnx.monitors.MonitorTracker.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.monitors.MonitorUnavailableError¶
class nnx.monitors.MonitorUnavailableError
A monitor with on_missing="error" / on_nonfinite="error" found no finite value to decide on.
nnx.monitors.MetricAccumulator¶
class nnx.monitors.MetricAccumulator(*args, **kwargs)
Accumulates one metric over an epoch's full sample.
Details
``update`` receives one batch's valid targets and the declared
prediction input as NumPy arrays (sample axis first); ``result``
returns the metric over everything seen, or ``None`` if nothing was,
without changing the accumulation.
Optional, for bounded streaming (``nnx.streaming``):
``merge(other)`` adds another accumulator's state in place, keeping
nothing of ``other`` that a later update could change, and returns
``None``; ``stores_scores = True`` marks an accumulator that keeps every
score (such as a rank metric), which bounded streaming refuses. The
built-in metrics implement ``merge``.
nnx.monitors.MetricAccumulator.update¶
nnx.monitors.MetricAccumulator.update(self, target: 'np.ndarray', prediction: 'np.ndarray') -> 'None'
No public description is currently available.
nnx.monitors.MetricAccumulator.result¶
nnx.monitors.MetricAccumulator.result(self) -> 'Optional[float]'
No public description is currently available.
nnx.monitors.register_metric¶
nnx.monitors.register_metric(id: 'str', version: 'int', factory: 'Callable[[Mapping[str, Any]], MetricAccumulator]', *, input: 'str', mode: 'str', check_config: 'Optional[Callable[[Mapping[str, Any]], None]]' = None, replace: 'bool' = False) -> 'None'
Register a metric under (id, version) for this process.
Details
``factory(config)`` returns a fresh :class:`MetricAccumulator` for one
split of one epoch (``config`` is the spec's read-only config);
``input`` is one of :data:`METRIC_INPUTS`; ``mode`` is the natural
direction (``"min"``: lower is better). ``check_config(config)``, when
given, validates a spec's config as the run starts, before any batch.
Registering an existing ``(id, version)`` raises unless
``replace=True``; bump the version when the metric's meaning changes.
nnx.monitors.unregister_metric¶
nnx.monitors.unregister_metric(id: 'str', version: 'int') -> 'bool'
Remove a registration; True when one existed.
nnx.monitors.registered_metrics¶
nnx.monitors.registered_metrics() -> 'tuple[tuple[str, int], ...]'
Every registered (id, version), sorted.
2.9. Objectives and the shared update engine (nnx.objectives)¶
nnx.objectives.LossTerm¶
class nnx.objectives.LossTerm(name: 'str', numerator: 'torch.Tensor', denominator: 'Optional[Union[int, float]]' = None, reduction: 'str' = 'mean', weight: 'float' = 1.0) -> 'None'
One named part of an objective's loss for one microbatch.
Details
Args:
name: identifies the term across the microbatches of a window.
numerator: differentiable scalar — the **sum** of the term over its
samples (or, for ``reduction="sum"``, the term's total).
denominator: for ``reduction="mean"``, the number (or total weight)
of samples ``numerator`` sums over — ``0`` when every sample is
masked; must be ``None`` for ``reduction="sum"``.
reduction: ``"mean"`` (divided by the window-total denominator) or
``"sum"`` (the window total). A term keeps its reduction for
the whole window; mixing them fails.
weight: finite factor applied to the term's window value.
nnx.objectives.LossTerm.value¶
property nnx.objectives.LossTerm.value
This microbatch's own value (None when fully masked).
nnx.objectives.ObjectiveContext¶
class nnx.objectives.ObjectiveContext(model: 'NNModel', batch: 'Any', epoch_idx: 'int', batch_idx: 'int', extra_metrics: 'Optional[Mapping[str, Callable]]' = None, precision: 'Optional[ResolvedPrecision]' = None) -> 'None'
What an objective sees for one microbatch.
Details
``precision`` is the run's resolved precision (``None`` when
called outside a run). The objective runs under its autocast, so a
reduced run's forward outputs are float16 / bfloat16:
:meth:`full_precision` returns them as float32 for the loss's
bookkeeping and the records (bf16 has no NumPy dtype).
nnx.objectives.ObjectiveContext.full_precision¶
nnx.objectives.ObjectiveContext.full_precision(self, tensor: 'torch.Tensor') -> 'torch.Tensor'
tensor in full precision: a reduced run's float16 / bfloat16 output as float32 (differentiable), anything else unchanged.
nnx.objectives.ObjectiveResult¶
class nnx.objectives.ObjectiveResult(terms: 'Sequence[LossTerm]', record: 'Optional[NNEvaluationDataPoint]' = None) -> 'None'
An objective's output: the loss terms and, optionally, the microbatch's metric record (its loss is filled from the terms when absent).
nnx.objectives.ObjectiveResult.loss¶
nnx.objectives.ObjectiveResult.loss(self) -> 'Optional[float]'
The microbatch's weighted loss over its unmasked terms.
nnx.objectives.Objective¶
class nnx.objectives.Objective(*, nonfinite: 'str' = 'fail') -> 'None'
Base class for objectives: override __call__.
Details
``nonfinite`` (``"fail"`` default, or ``"skip"``) is the engine's policy
for a non-finite loss term or gradient. Any callable
``(ObjectiveContext) -> ObjectiveResult`` works as an objective; a plain
function uses ``"fail"``.
nnx.objectives.Objective.check_run¶
nnx.objectives.Objective.check_run(self, model: 'NNModel', *, optimizers: 'Mapping[str, torch.optim.Optimizer]', callbacks: 'Sequence[Any]') -> 'None'
Refuse a run this objective cannot train — called by NNModel.train and Trainer.train with the built optimizers and the callbacks, before any run is reserved or parameter changes. Raise to refuse; the default accepts every run.
nnx.objectives.Objective.after_update¶
nnx.objectives.Objective.after_update(self, events: 'tuple[UpdateEvent, ...]') -> 'None'
Called once per committed update — after the window's step (an optimizer without gradients is not stepped), never for a microbatch or a skipped window — with its events (one per optimizer), before callbacks see them. The default does nothing.
nnx.objectives.SupervisedObjective¶
class nnx.objectives.SupervisedObjective(*, nonfinite: 'str' = 'fail') -> 'None'
The standard supervised loss as an objective: the model's loss_fn on its outputs, normalized by the loss's own denominator (non-ignored targets, class weights) — or summed for a reduction="sum" loss — and, with a TaskSpec, the task's masked loss over its valid targets.
nnx.objectives.KDObjective¶
class nnx.objectives.KDObjective(teacher: 'NNModel', *, alpha: 'float' = 0.5, temperature: 'float' = 4.0, nonfinite: 'str' = 'fail')
Hinton knowledge distillation as an objective.
Details
Two terms: ``"distillation"`` — the temperature-softened
``KL(teacher ‖ student)`` summed over the rows and normalized by their
count, scaled by ``T²`` (weight ``alpha``) — and ``"supervised"`` —
the student's loss on the labels with its own denominator (weight
``1 − alpha``). Over a window this is exactly
``alpha · KL_batchmean + (1 − alpha) · loss`` of the full batch, however
the rows are split into microbatches or targets are ignored. The teacher
is frozen and put in eval mode once, as ``kd_train_step_factory`` does,
and never updated. At ``alpha=0`` only the ``"supervised"`` term is
returned.
nnx.objectives.supervised_objective¶
nnx.objectives.supervised_objective(*, nonfinite: 'str' = 'fail') -> 'SupervisedObjective'
The standard supervised loss as an objective (see SupervisedObjective).
nnx.objectives.kd_objective¶
nnx.objectives.kd_objective(teacher: 'NNModel', *, alpha: 'float' = 0.5, temperature: 'float' = 4.0, nonfinite: 'str' = 'fail') -> 'KDObjective'
Knowledge distillation as an objective (see KDObjective) — the objective counterpart of kd_train_step_factory, which stays available (and unchanged) as the imperative step.
nnx.objectives.UpdateEvent¶
class nnx.objectives.UpdateEvent(optimizer: 'str', update_idx: 'int', epoch_idx: 'int', batch_idx: 'int', microbatches: 'int', losses: 'Mapping[str, float]', loss: 'float') -> 'None'
One committed optimizer update (detached values only).
Details
Attributes:
optimizer: the optimizer that stepped (``"default"`` for
``NNModel.train``; the name for ``Trainer``).
update_idx: this optimizer's committed updates so far (1-based).
epoch_idx / batch_idx: where the window closed.
microbatches: microbatches the update accumulated.
losses: each term's window value (normalized terms divided by their
window-total denominator).
loss: the weighted total.
2.10. Arbitrary modules and registered model factories (nnx.models)¶
nnx.models.ModelSpec¶
class nnx.models.ModelSpec(id: 'str', version: 'int' = 1, config: 'Optional[Mapping[str, Any]]' = None, *, seed: 'int' = 0)
Stable, serializable reference to a registered model factory.
Details
``id`` and ``version`` name the factory; ``config`` is handed to it
read-only and must be JSON-like (``None``, ``bool``, ``int``, finite
``float``, ``str``, lists and string-keyed mappings); ``seed`` is the
``torch.manual_seed`` construction runs under. The spec never holds the
callable, so it round-trips through ``run.yaml``, checkpoints and Hub
configs. Immutable; equal specs have equal ``state()``.
nnx.models.ModelSpec.state¶
nnx.models.ModelSpec.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.models.ModelSpec.from_state¶
nnx.models.ModelSpec.from_state(state: 'Mapping[str, Any]') -> 'ModelSpec'
No public description is currently available.
nnx.models.RuntimeModule¶
class nnx.models.RuntimeModule(module: 'str', topology: 'str') -> 'None'
Descriptor of a caller-owned module that no factory describes.
Details
``module`` is the class's qualified name and ``topology`` a fingerprint
of its parameter and buffer names, shapes and dtypes and of its
``repr`` (layer hyperparameters such as dropout rates) — never of
weight values. Like a built-in net, the run id therefore reflects the
architecture, not the weights: two differently-initialized instances of
one architecture share a run id, so give them distinct ``data_id`` /
``salt`` values. ``reconstructible`` is always ``False``: the weights
can only be loaded back into a module the caller supplies again.
nnx.models.RuntimeModule.of¶
nnx.models.RuntimeModule.of(module: 'nn.Module') -> 'RuntimeModule'
The descriptor of module (its class and topology).
nnx.models.RuntimeModule.state¶
nnx.models.RuntimeModule.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.models.RuntimeModule.from_state¶
nnx.models.RuntimeModule.from_state(state: 'Mapping[str, Any]') -> 'RuntimeModule'
No public description is currently available.
nnx.models.ModelFactory¶
type alias nnx.models.ModelFactory
Public type alias.
nnx.models.register_model_factory¶
nnx.models.register_model_factory(id: 'str', version: 'int', factory: 'ModelFactory', *, replace: 'bool' = False) -> 'None'
Register factory under (id, version) for this process.
Details
``factory(config)`` receives a :class:`ModelSpec`'s read-only config and
must return a ``torch.nn.Module``. Registering an existing
``(id, version)`` raises unless ``replace=True``; bump ``version`` when
the factory's architecture changes, so saved weights are never loaded
into a different topology.
nnx.models.unregister_model_factory¶
nnx.models.unregister_model_factory(id: 'str', version: 'int') -> 'bool'
Remove (id, version); returns whether it was registered. Runs and checkpoints that reference it still load their metadata — only rebuilding the module needs the registration.
nnx.models.registered_model_factories¶
nnx.models.registered_model_factories() -> 'tuple[tuple[str, int], ...]'
Sorted (id, version) pairs currently registered.
nnx.models.resolve_model_factory¶
nnx.models.resolve_model_factory(spec: 'ModelSpec') -> 'ModelFactory'
Look spec up in the registry. Never imports anything: an unregistered id or version raises MissingModelFactoryError naming what is registered.
nnx.models.build_module¶
nnx.models.build_module(spec: 'ModelSpec') -> 'nn.Module'
Build spec's module: resolve the factory (failing before anything runs when it is unknown), then call it with Python's random, NumPy and torch all seeded from spec.seed and the ambient RNG states restored afterwards.
nnx.models.MissingModelFactoryError¶
class nnx.models.MissingModelFactoryError
A module cannot be rebuilt: its factory is not registered, or it is a runtime-only module (reconstructible=False) that no factory describes.
nnx.models.BatchAdapter¶
class nnx.models.BatchAdapter()
How a non-built-in module sees a batch.
Details
:meth:`split` returns ``(args, kwargs, target)``: the module is called as
``module(*args, **kwargs)`` and ``target`` (``None`` when the batch has
none, e.g. for prediction) is what the loss scores. :meth:`output`
turns the module's raw return value into the output tensor; the default
requires a tensor. Subclass either to support other layouts.
nnx.models.BatchAdapter.split¶
nnx.models.BatchAdapter.split(self, batch: 'Any') -> 'tuple[tuple[Any, ...], dict[str, Any], Any]'
No public description is currently available.
nnx.models.BatchAdapter.output¶
nnx.models.BatchAdapter.output(self, raw: 'Any') -> 'torch.Tensor'
No public description is currently available.
nnx.models.PositionalInputs¶
class nnx.models.PositionalInputs(n_inputs: 'int' = 1) -> 'None'
Positional inputs: a batch (x1, ..., xn, y) calls module(x1, ..., xn) and scores y; (x1, ..., xn) or a bare tensor has no target.
nnx.models.PositionalInputs.split¶
nnx.models.PositionalInputs.split(self, batch: 'Any') -> 'tuple[tuple[Any, ...], dict[str, Any], Any]'
No public description is currently available.
nnx.models.KeywordInputs¶
class nnx.models.KeywordInputs(inputs: 'Sequence[str]', target: 'Optional[str]' = 'labels') -> 'None'
Keyword inputs: a mapping batch calls module(**{name: batch[name] for name in inputs}) and scores batch[target] (None when the key is absent or target=None).
nnx.models.KeywordInputs.split¶
nnx.models.KeywordInputs.split(self, batch: 'Any') -> 'tuple[tuple[Any, ...], dict[str, Any], Any]'
No public description is currently available.
nnx.models.module_topology¶
nnx.models.module_topology(module: 'nn.Module') -> 'str'
Fingerprint of module's state-dict names, shapes and dtypes plus its repr (layer hyperparameters). Uninitialized lazy parameters (nn.LazyLinear before its first forward) count by name only.
2.11. Typed decisions (nnx.decisions)¶
nnx.decisions.Question¶
class nnx.decisions.Question()
Base of the decision primitives. kind is the discriminator; digest identifies the exact question (kind, prompt, and every option's id and description in order) so a result can never be matched to a reordered or reworded question.
nnx.decisions.Question.state¶
nnx.decisions.Question.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.Question.digest¶
nnx.decisions.Question.digest(self) -> 'str'
SHA-256 of the canonical question state (computed once; questions are immutable).
nnx.decisions.Question.model_view¶
nnx.decisions.Question.model_view(self) -> 'dict[str, Any]'
What a model may see: the prompt and the option texts in order — never the bookkeeping ids.
nnx.decisions.Question.option_ids¶
property nnx.decisions.Question.option_ids
No public description is currently available.
nnx.decisions.Option¶
class nnx.decisions.Option(id: 'str', description: 'str') -> 'None'
One answer of a Choice or one level of a Score.
Details
``id`` is bookkeeping — how results are keyed, never shown to a model;
``description`` is the model-facing text.
nnx.decisions.Choice¶
class nnx.decisions.Choice(prompt: 'str', options: 'tuple[Option, ...]') -> 'None'
Pick one of 2+ options; answered by a distribution over the options.
nnx.decisions.Choice.option_ids¶
property nnx.decisions.Choice.option_ids
No public description is currently available.
nnx.decisions.Choice.state¶
nnx.decisions.Choice.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.Choice.model_view¶
nnx.decisions.Choice.model_view(self) -> 'dict[str, Any]'
What a model may see: the prompt and the option texts in order — never the bookkeeping ids.
nnx.decisions.Boolean¶
class nnx.decisions.Boolean(prompt: 'str') -> 'None'
A yes / no question; answered by one probability p_true.
nnx.decisions.Boolean.state¶
nnx.decisions.Boolean.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.Boolean.model_view¶
nnx.decisions.Boolean.model_view(self) -> 'dict[str, Any]'
What a model may see: the prompt and the option texts in order — never the bookkeeping ids.
nnx.decisions.Score¶
class nnx.decisions.Score(prompt: 'str', levels: 'tuple[Option, ...]') -> 'None'
Place the input on 2+ ordered levels (lowest first); answered by a distribution over the levels. The levels are ordinal: their order is meaningful, their spacing is not.
nnx.decisions.Score.option_ids¶
property nnx.decisions.Score.option_ids
No public description is currently available.
nnx.decisions.Score.state¶
nnx.decisions.Score.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.Score.model_view¶
nnx.decisions.Score.model_view(self) -> 'dict[str, Any]'
What a model may see: the prompt and the option texts in order — never the bookkeeping ids.
nnx.decisions.question_from_state¶
nnx.decisions.question_from_state(state: 'Mapping[str, Any]') -> 'Union[Choice, Boolean, Score]'
Rebuild a question from Question.state, by its kind; malformed state raises InvalidDecisionRequest.
nnx.decisions.DecisionResult¶
type alias nnx.decisions.DecisionResult
Public type alias.
nnx.decisions.ChoiceResult¶
class nnx.decisions.ChoiceResult(question_digest: 'str', distribution: 'tuple[tuple[str, float], ...]', provider: 'Optional[str]' = None, raw: 'Any' = None) -> 'None'
A validated answer to a Choice: distribution in the question's option order. raw is the provider's own output, kept apart from the normalized fields (and from equality).
nnx.decisions.ChoiceResult.probabilities¶
nnx.decisions.ChoiceResult.probabilities(self) -> 'dict[str, float]'
No public description is currently available.
nnx.decisions.ChoiceResult.top¶
property nnx.decisions.ChoiceResult.top
The most probable option id (the first one on ties).
nnx.decisions.BooleanResult¶
class nnx.decisions.BooleanResult(question_digest: 'str', p_true: 'float', provider: 'Optional[str]' = None, raw: 'Any' = None) -> 'None'
A validated answer to a Boolean: one p_true in [0, 1].
nnx.decisions.BooleanResult.p_false¶
property nnx.decisions.BooleanResult.p_false
No public description is currently available.
nnx.decisions.ScoreResult¶
class nnx.decisions.ScoreResult(question_digest: 'str', distribution: 'tuple[tuple[str, float], ...]', provider: 'Optional[str]' = None, vendor_score: 'Optional[float]' = None, raw: 'Any' = None) -> 'None'
A validated answer to a Score: distribution over the levels in the question's order (lowest first).
Details
:attr:`expected_index` is ``sum(i * p_i)`` over zero-based levels — an
**ordinal** summary (a position between levels), not an interval-scale
score: the levels' spacing is undefined. A provider's own numeric
score, if it reports one, stays in ``vendor_score`` and is never mixed
into the distribution.
nnx.decisions.ScoreResult.probabilities¶
nnx.decisions.ScoreResult.probabilities(self) -> 'dict[str, float]'
No public description is currently available.
nnx.decisions.ScoreResult.expected_index¶
property nnx.decisions.ScoreResult.expected_index
No public description is currently available.
nnx.decisions.validate_response¶
nnx.decisions.validate_response(question: 'Union[Choice, Boolean, Score]', response: 'Any', *, provider: 'Optional[str]' = None, raw: 'Any' = None, vendor_score: 'Optional[float]' = None) -> 'DecisionResult'
Validate a provider's keyed output against question and return the normalized result.
Details
For a :class:`Choice` or :class:`Score`, ``response`` maps option ids to
probabilities (a mapping, or ``(id, p)`` pairs in any order). The result
is reordered into the question's order; a missing, duplicate, unknown or
unlabeled (empty / non-string) id, a probability outside ``[0, 1]``
(one too large for a float included, a ``vendor_score`` too) and a
distribution not summing to 1 within :data:`PROBABILITY_TOLERANCE`
raise :class:`InvalidDecisionResponse` — a malformed distribution is
never renormalized. For a :class:`Boolean`, ``response`` is ``p_true``
or ``{"true": p, "false": 1 - p}``. ``raw`` (the provider's own output)
and ``vendor_score`` (a :class:`Score` provider's own number) are kept
apart from the normalized fields; a ``vendor_score`` for another kind
of question raises :class:`InvalidDecisionResponse`.
nnx.decisions.Capabilities¶
class nnx.decisions.Capabilities(primitives: 'frozenset[str]', modalities: 'frozenset[str]', dynamic_labels: 'bool' = False, labels: 'Optional[tuple[str, ...]]' = None, max_batch: 'Optional[int]' = None, inference: 'bool' = True, training: 'bool' = False, export: 'bool' = False) -> 'None'
What a provider declares it can do. check rejects anything else with UnsupportedCapability before any model call or network I/O.
Details
Attributes:
primitives: the decision kinds it answers (``"choice"``,
``"boolean"``, ``"score"``).
modalities: the input kinds it reads (e.g. ``"tensor"``,
``"text"``).
dynamic_labels: whether a request may bring any option set; ``False``
means the options must map exactly onto ``labels``.
labels: the fixed label space when ``dynamic_labels`` is ``False``;
``None`` declares none, so option sets are not checked.
max_batch: the most inputs per call (``None``: unbounded).
inference / training / export: whether the backend can answer, be
trained and be exported. A hosted provider declares inference
only.
nnx.decisions.Capabilities.check¶
nnx.decisions.Capabilities.check(self, question: 'Question', *, modality: 'str', batch_size: 'int', label_ids: 'Optional[Sequence[str]]' = None) -> 'None'
Raise UnsupportedCapability unless this provider serves question for a batch of batch_size modality inputs.
Details
Without dynamic labels, a Choice's or Score's ids — or ``label_ids``,
the provider's own mapping of them — must be exactly the declared
``labels``, in any order.
nnx.decisions.DecisionProvider¶
class nnx.decisions.DecisionProvider(*args, **kwargs)
Anything that answers typed decision questions: it declares Capabilities and answers one question for a batch of inputs with validated results, one per input.
nnx.decisions.DecisionProvider.capabilities¶
nnx.decisions.DecisionProvider.capabilities(self) -> 'Capabilities'
No public description is currently available.
nnx.decisions.DecisionProvider.decide¶
nnx.decisions.DecisionProvider.decide(self, question: 'Question', inputs: 'Any') -> 'list[DecisionResult]'
No public description is currently available.
nnx.decisions.FixedHeadProvider¶
class nnx.decisions.FixedHeadProvider(model: 'NNModel', labels: 'Optional[Sequence[str]]' = None, option_map: 'Optional[Mapping[str, str]]' = None, ordinal: 'bool' = False, max_batch: 'Optional[int]' = None, name: 'str' = 'nnx.fixed_head', question: 'Optional[Boolean]' = None) -> 'None'
A trained NNx classifier as a decision provider.
Details
The head fixes the label space and the primitives it justifies: a
categorical head (softmax over its classes) answers :class:`Choice` —
and :class:`Score` when constructed with ``ordinal=True``, whose levels
must then follow the head's class order — and a single-logit Bernoulli
head (``BCEWithLogitsLoss``, or a one-output multilabel task) answers
one :class:`Boolean`: ``question``, the one it was trained for. Its
single probability is bound to no prompt, so any other Boolean (another
digest) raises :class:`UnsupportedCapability` before :attr:`model_calls`
advances. Probabilities come from the model's ``predict_proba``,
which evaluates in eval mode and restores every submodule's
training mode; they are computed in float64 from the returned logits,
so half-precision models meet the ``1e-6`` tolerance.
A request's option ids must be exactly the head's ``labels`` (in any
order), or ``option_map`` must map them one-to-one onto the labels; an
unseen or missing label raises :class:`UnsupportedCapability` before
:attr:`model_calls` advances. A label count that does not fit the head
is rejected at construction when the head's width is known (a task's
count or a built-in net's ``output_dim``), otherwise with
:class:`InvalidDecisionRequest` on the first call. Results come back in
the request's option order with each row's logits (a copy) as ``raw``;
an empty batch returns ``[]`` without calling the model.
Args:
model: the trained :class:`~nnx.NNModel`.
labels: one name per output class, in column order; defaults to the
model's ``TaskSpec`` labels, and must equal them when both exist.
option_map: request option id → head label (a bijection);
categorical heads only.
ordinal: whether the class order is meaningful (enables Score);
categorical heads only.
max_batch: the most inputs per call.
name: the provider name recorded on results.
question: the :class:`Boolean` a Bernoulli head was trained to
answer — required for one, refused for a categorical head
(which answers by its labels).
nnx.decisions.FixedHeadProvider.capabilities¶
nnx.decisions.FixedHeadProvider.capabilities(self) -> 'Capabilities'
No public description is currently available.
nnx.decisions.FixedHeadProvider.check¶
nnx.decisions.FixedHeadProvider.check(self, question: 'Question', inputs: 'Any') -> 'None'
Every check decide makes before calling the model — the question, the inputs' modality and batch size, and the label space (through option_map) or a Boolean's identity (its digest) — so a caller (a decision job) can validate a request without calling the model.
nnx.decisions.FixedHeadProvider.decide¶
nnx.decisions.FixedHeadProvider.decide(self, question: 'Question', inputs: 'Any') -> 'list[DecisionResult]'
Answer question for every row of inputs (a tensor, an array or a tuple of them).
nnx.decisions.nli¶
module nnx.decisions.nli
A local label-conditioned baseline: decisions from a caller-supplied natural-language-inference (NLI) model.
Details
A fixed-head classifier knows only the labels it was trained on. An NLI
cross-encoder scores whether a *premise* (the input text) entails a
*hypothesis* (a sentence built from a candidate's description), so the
candidates can be supplied at inference time — none of them needs to exist
when the provider is built::
from nnx.decisions import Choice, NLIProvider
provider = NLIProvider(model, tokenizer, entailment_id="entailment", contradiction_id="contradiction",
hypothesis_template="This text is about {}.", revision="abc123")
provider.decide(Choice("Topic?", (("t-sport", "sports"), ("t-econ", "the economy"))), texts)
**Scoring (explicit, never calibrated).**
- A :class:`Choice` scores every (input, candidate) pair and normalizes the
**entailment logits across the candidates** with a softmax
(``choice_scoring="entailment_softmax"``): logits ``log(3)`` and ``0``
give ``(0.75, 0.25)``. Results are keyed by the request's option ids, in
its order, whatever order the candidates arrive in. Only the candidates'
descriptions are scored: the Choice's prompt is not part of any
hypothesis (write it into ``hypothesis_template`` when it matters).
- A :class:`Boolean` scores one pair per input — the prompt rendered by
``boolean_template`` — and takes the softmax over that pair's
**contradiction and entailment** logits alone
(``boolean_scoring="entailment_vs_contradiction"``): contradiction ``0``,
entailment ``log(4)`` give ``p_true = 0.8``. Inputs are never normalized
together.
- :class:`Score` is unsupported (an NLI model has no ordinal semantics) and
is rejected before any model call.
These are zero-shot scores from a model trained for another task: they are
**not calibrated** probabilities, and how well they transfer to a decision
task stays an empirical question — measure it on labelled records (see
``examples/decision_nli.py``).
**The provider contract.** NNx imports no NLI library and downloads
nothing: the caller passes a tokenizer and a ``torch.nn.Module`` already
loaded (from a local path, a pinned ``revision`` or a test stub). The
tokenizer is called HuggingFace-style — ``tokenizer(premises, hypotheses,
truncation=..., max_length=..., padding=True, return_tensors="pt")``
returning a mapping of tensors (and, to measure lengths first, with
``truncation=False, padding=False`` and no ``return_tensors``, returning
per-pair ``input_ids`` — or an ``attention_mask``, whose sums are used) —
and the model as ``model(**encoded)``,
returning the logits ``(pairs, classes)`` or an object with ``.logits``.
Pairs are sent in chunks of ``pair_batch_size`` (the last chunk may be
shorter), moved to the model's device; the model runs in eval mode under
``no_grad`` and every submodule's training flag is restored afterwards, on
success and on failure.
**Truncation (explicit, reported).** Every pair is measured untruncated
first. ``truncation="only_first"`` cuts the premise to fit ``max_length``
and reports which pairs were cut in each result's ``raw["truncated"]``; a
hypothesis that leaves no room for the premise (measured with an empty
premise) rejects the request before any model call. ``truncation="error"`` rejects a request with an
over-long pair before any model call.
Each result's ``raw`` records the pair logits, the truncation report and
the provider's :meth:`NLIProvider.record` — the NLI template, label ids,
scoring methods, truncation policy and model ``revision`` — with
``"calibrated": False``, so a replayed or benchmarked result keeps how it
was scored.
nnx.decisions.NLIProvider¶
class nnx.decisions.NLIProvider(model: 'Any', tokenizer: 'Any', entailment_id: 'Union[int, str]', contradiction_id: 'Union[int, str]', hypothesis_template: 'str' = 'This example is {}.', boolean_template: 'str' = '{}', choice_scoring: 'str' = 'entailment_softmax', boolean_scoring: 'str' = 'entailment_vs_contradiction', max_length: 'int' = 512, truncation: 'str' = 'only_first', pair_batch_size: 'int' = 16, max_batch: 'Optional[int]' = None, revision: 'Optional[str]' = None, name: 'str' = 'nnx.nli') -> 'None'
A caller-supplied NLI model as a label-conditioned decision provider — see the module docstring for the scoring, the provider contract and the truncation policy.
Details
Args:
model: the NLI cross-encoder (a ``torch.nn.Module``; ``model(**encoded)``
returns logits ``(pairs, classes)`` or an object with ``.logits``).
Its device is its parameters' device.
tokenizer: the matching tokenizer, called HuggingFace-style.
entailment_id / contradiction_id: the model's class ids for
entailment and contradiction — ints, or label names resolved
through ``model.config.label2id``. Distinct, and within the
model's class count when its config declares one.
hypothesis_template: renders a Choice candidate's description
(exactly one ``{}``), e.g. ``"This example is {}."``.
boolean_template: renders a Boolean's prompt (exactly one ``{}``).
choice_scoring / boolean_scoring: the normalization methods (see the
module docstring); explicit so a record says how it scored.
max_length: the most tokens per pair.
truncation: ``"only_first"`` (cut the premise, report it) or
``"error"`` (reject over-long pairs before any model call).
pair_batch_size: pairs per model call.
max_batch: the most inputs per request (``None``: unbounded).
revision: the model revision the caller loaded (recorded, never
fetched).
name: the provider name recorded on results.
nnx.decisions.NLIProvider.capabilities¶
nnx.decisions.NLIProvider.capabilities(self) -> 'Capabilities'
No public description is currently available.
nnx.decisions.NLIProvider.record¶
nnx.decisions.NLIProvider.record(self) -> 'dict[str, Any]'
How this provider scores: what a replay record or a benchmark keeps beside each result. The scores are never calibrated.
nnx.decisions.NLIProvider.hypotheses¶
nnx.decisions.NLIProvider.hypotheses(self, question: 'Question') -> 'list[str]'
The hypothesis per candidate (a Choice, in its option order) or the one hypothesis of a Boolean.
nnx.decisions.NLIProvider.check¶
nnx.decisions.NLIProvider.check(self, question: 'Question', inputs: 'Any') -> 'None'
Every check decide makes before calling the model — the question, the texts (non-empty strings), the capabilities and the tokenizer-only length check of each pair — so a caller (a decision job, the benchmark) can refuse a request without spending a call.
nnx.decisions.NLIProvider.decide¶
nnx.decisions.NLIProvider.decide(self, question: 'Question', inputs: 'Any') -> 'list[DecisionResult]'
Answer question for every text in inputs.
nnx.decisions.jev¶
module nnx.decisions.jev
Typed decisions on Jev models through the TypeSafe Python SDK.
Details
:class:`JevProvider` answers :class:`Choice`, :class:`Boolean` and
:class:`Score` questions about text with one ``system_one`` call per input,
through the ``typesafe-sdk`` client — NNx adds no second transport stack::
from nnx.decisions import Boolean, Choice, JevProvider
with JevProvider(model="jev-1.13.0") as provider: # TYPESAFE_API_KEY from the environment
topic, spam = provider.decide_many(
[Choice("Topic?", (("t-sport", "sports"), ("t-econ", "the economy"))), Boolean("Is this spam?")],
texts,
)
Install the extra first: ``pip install "thekaveh-nnx[jev]"`` (tested against
``typesafe-sdk>=0.7,<0.8``). Importing :mod:`nnx` or :mod:`nnx.decisions`
never imports the SDK; building a provider without an injected client does,
and raises an :class:`ImportError` naming the extra when it is missing.
**Translation.** A :class:`Choice` becomes a Jev ``choice`` keyed by its
options' descriptions — the model-facing text, never the bookkeeping ids —
a :class:`Score` a ``score`` whose criteria are its levels' descriptions in
order (lowest first), a :class:`Boolean` a ``noul``; the prompt is the
``instructions``. Every question of one :meth:`JevProvider.decide_many` goes
into the same call, one call per input. Answers are realigned to each
question's own option order and validated
(:func:`~nnx.decisions.validate_response`); a Score's expected score stays
in ``vendor_score``.
**Confidence.** A Choice's or Score's ``confidence`` is Jev's own certainty
in its selection, recorded as ``raw["provider_confidence"]`` — it is not a
probability that the answer is correct. A Boolean has none.
**Metadata.** Each result's ``raw`` is plain JSON: ``provider``, the
``model`` the service resolved (even for an alias request), the
``requested_model`` when one was set, the ``request_id`` and the reported
``usage`` token counts. A field the service did not report is absent, and
clients, headers and credentials are never recorded.
**Retries and failures.** The SDK's :class:`~typesafe_sdk.RetryPolicy`
owns retries (pass ``retry=`` when the provider builds its client); the
adapter makes each request once. What still fails becomes a typed
:class:`JevError` — :class:`JevTimeout`, :class:`JevAuthenticationError`,
:class:`JevRateLimited` or :class:`JevMalformedResponse` — keeping the SDK
error as ``__cause__`` and, when the service answered, its ``request_id``
(a timeout or a lost connection has none).
**Cost.** Every text is its own ``system_one`` request, so one provider
call over ``n`` texts sends ``n`` requests: a benchmark ``Budget.max_calls``
or a job's ``Limits.max_requests`` counts provider calls, not requests — set
``max_batch`` to bound requests per call. A failure on one text stops the
call; answers already received for earlier texts are not returned.
**Lifecycle.** An injected ``client`` / ``async_client`` stays the
caller's: the provider never closes it, and never builds a second client
beside it — with only a sync client injected the provider is sync-only
(``adecide`` / ``adecide_many`` are ``None``, so a decision job's ``arun``
runs it in its own worker thread, one call at a time); with only an async
one, the sync methods are refused. With no client injected the provider builds its own (lazily, from
``api_key`` / ``base_url`` / ``timeout`` / ``retry`` or the SDK's
environment variables): :meth:`JevProvider.close` (``with``) closes a built
sync client, :meth:`JevProvider.aclose` (``async with``) both — after an
error or a cancellation too. A built async client belongs to one event
loop; a call (or :meth:`JevProvider.aclose`) from another loop drops it
with a ``ResourceWarning`` — close it with ``async with`` inside each
``asyncio.run``.
nnx.decisions.JevProvider¶
class nnx.decisions.JevProvider(client: 'Any' = None, *, async_client: 'Any' = None, model: 'Optional[str]' = None, api_key: 'Optional[str]' = None, base_url: 'Optional[str]' = None, timeout: 'Optional[float]' = None, retry: 'Any' = None, max_batch: 'Optional[int]' = None, name: 'str' = 'jev') -> 'None'
A DecisionProvider for Jev models through the TypeSafe SDK. See the module docstring for translation, confidence, metadata, retry and lifecycle semantics.
Details
Args:
client: a ``typesafe_sdk.TypeSafeClient`` (or anything with its
``system_one``) the caller owns; never closed by the provider.
async_client: an ``AsyncTypeSafeClient`` the caller owns, for
:meth:`adecide` / :meth:`adecide_many`. Injecting either means
the provider builds no client at all.
model: the model to request (pin a version such as
``"jev-1.13.0"``; ``None`` uses the client's default).
api_key / base_url / timeout / retry: settings for the clients the
provider builds itself (``api_key`` defaults to the
``TYPESAFE_API_KEY`` environment variable). Refused with an
injected client, whose settings are the caller's.
max_batch: the most inputs per call (``None``: unbounded) — each
input is one request.
name: the provider name recorded on results.
Raises:
ImportError: no client is injected and the ``jev`` extra is missing.
JevError: a client the provider builds cannot be configured (for
example no API key), at the first call.
nnx.decisions.JevProvider.capabilities¶
nnx.decisions.JevProvider.capabilities(self) -> 'Capabilities'
No public description is currently available.
nnx.decisions.JevProvider.record¶
nnx.decisions.JevProvider.record(self) -> 'dict[str, Any]'
The provider's identity for benchmarks and replay: no client, header or credential.
nnx.decisions.JevProvider.check¶
nnx.decisions.JevProvider.check(self, question: 'Question', inputs: 'Any') -> 'None'
Every check made before a request: the question, the inputs and the provider's capabilities — no network I/O.
nnx.decisions.JevProvider.close¶
nnx.decisions.JevProvider.close(self) -> 'None'
Close the sync client the provider built (never an injected one). A built async client needs aclose (async with).
nnx.decisions.JevProvider.aclose¶
nnx.decisions.JevProvider.aclose(self) -> 'None'
Close every client the provider built, async one included.
nnx.decisions.JevProvider.decide¶
nnx.decisions.JevProvider.decide(self, question: 'Question', inputs: 'Any') -> 'list[DecisionResult]'
Answer question for every text in inputs.
nnx.decisions.JevProvider.decide_many¶
nnx.decisions.JevProvider.decide_many(self, questions: 'Sequence[Question]', inputs: 'Any') -> 'list[list[DecisionResult]]'
Answer every question for every text: one system_one call per text holding all the questions. One result list per question, one result per text.
nnx.decisions.JevProvider.adecide¶
nnx.decisions.JevProvider.adecide(self, question: 'Question', inputs: 'Any') -> 'list[DecisionResult]'
decide through the async client.
nnx.decisions.JevProvider.adecide_many¶
nnx.decisions.JevProvider.adecide_many(self, questions: 'Sequence[Question]', inputs: 'Any') -> 'list[list[DecisionResult]]'
decide_many through the async client (absent — None — when only a sync client was injected).
nnx.decisions.DecisionError¶
class nnx.decisions.DecisionError
Base class of every decision-API error.
nnx.decisions.UnsupportedCapability¶
class nnx.decisions.UnsupportedCapability
A request the provider cannot serve — a primitive, modality, batch size or label space it does not declare. Raised before any model call or network I/O.
nnx.decisions.InvalidDecisionRequest¶
class nnx.decisions.InvalidDecisionRequest
A malformed question: too few options, duplicate or empty ids, empty text.
nnx.decisions.InvalidDecisionResponse¶
class nnx.decisions.InvalidDecisionResponse
A provider response that does not fit its question: missing, duplicate, unknown or unlabeled ids, or a probability outside [0, 1] or a distribution that does not sum to 1 (never renormalized).
nnx.decisions.ProviderFailure¶
class nnx.decisions.ProviderFailure
The provider's backend failed while answering a supported, valid request (the original error is the __cause__).
nnx.decisions.JevError¶
class nnx.decisions.JevError(message: 'str', *, request_id: 'Optional[str]' = None) -> 'None'
A Jev request failed (the SDK error is the __cause__). request_id is the service's request id when it sent one.
nnx.decisions.JevTimeout¶
class nnx.decisions.JevTimeout(message: 'str', *, request_id: 'Optional[str]' = None) -> 'None'
The request exceeded its timeout (after the SDK's retries).
nnx.decisions.JevAuthenticationError¶
class nnx.decisions.JevAuthenticationError(message: 'str', *, request_id: 'Optional[str]' = None) -> 'None'
The API key was rejected or lacks access (HTTP 401 / 403).
nnx.decisions.JevRateLimited¶
class nnx.decisions.JevRateLimited(message: 'str', *, request_id: 'Optional[str]' = None, retry_after_ms: 'Optional[float]' = None)
The service refused the request for its rate limit (HTTP 429, after the SDK's retries). retry_after_ms is the wait it asked for, if any.
nnx.decisions.JevMalformedResponse¶
class nnx.decisions.JevMalformedResponse(message: 'str', *, request_id: 'Optional[str]' = None) -> 'None'
The service answered with data that does not fit the request: a body the SDK cannot parse, a missing answer, or a distribution that does not fit the question.
nnx.decisions.jobs¶
module nnx.decisions.jobs
Deferred decision jobs: batch independent questions, chain dependent ones.
Details
A :class:`DecisionJob` is an **immutable description** of decision work —
building one calls no provider, runs no callback and draws no RNG::
from nnx.decisions import DecisionJob as Job
triage = Job.collect({
"topic": Job.ask(topic_question, id="topic"),
"urgent": Job.ask(urgent_question, id="urgent"),
}).map(lambda answers: answers["topic"])
routed = Job.ask(topic_question, id="topic").then(
lambda answer: Follow(Job.ask(team_question, id="team"), state=route_state(answer))
)
result = triage.run(provider, state=texts, limits=Limits(max_questions=2))
- :meth:`DecisionJob.ask` — one question over the run's (or a
continuation's) **state**: the inputs the provider reads. Its value is the
provider's results, one per input row (or, with an abstention ``policy``,
the selective decisions).
- :meth:`DecisionJob.collect` — independent jobs, keyed; the value is a
mapping in the given key order.
- :meth:`DecisionJob.map` — a pure function of a job's value.
- :meth:`DecisionJob.then` — a **dependent** job: once its prerequisite
succeeds, ``fn(value)`` returns a :class:`Follow` — the next job and the
state it reads, mapped explicitly from the answer. It runs once, never
feeds a sibling question, and is bounded by ``Limits.max_depth``.
:meth:`DecisionJob.run` and :meth:`DecisionJob.arun` are the only effect
boundaries. They **borrow** the provider (never close it): independent
questions over the same state are batched into the fewest calls —
``ceil(count / cap)`` per state, in a stable order — where ``cap`` is the
provider's ``max_questions`` (a provider answering one question per call,
such as :class:`~nnx.decisions.FixedHeadProvider`, has cap 1) bounded by
``Limits.max_questions``. Duplicate question ids, questions the provider
cannot serve and a token cap the provider cannot enforce are rejected
before any call.
**Answers are checked.** Each call's answers must hold one result list per
question, one result per input row, each a result of that question (its
digest); anything else is an ``InvalidDecisionResponse`` failure.
**Fail-fast.** A provider error stops scheduling: :class:`JobFailed` carries
every outcome known so far (``outcomes``; the answered ones are
``completed``), the failing questions and the skipped ones; nothing
completed is re-run (the job never retries — retries belong to the
provider). Request and depth limits raise :class:`JobLimitExceeded` (in
``arun``, after the calls already in flight finish), the ``timeout``
:class:`JobTimeout`, each with the outcomes so far; questions never sent
are ``skipped``.
**Async.** :meth:`DecisionJob.arun` runs up to ``Limits.max_concurrency``
provider calls at once through a provider's ``adecide_many`` / ``adecide``.
A provider with only synchronous methods is called in a worker thread,
never alongside another call (its methods need not be thread-safe), and the
run waits for a running thread before it returns — a cancellation that
arrives meanwhile is delivered once the thread is done. Setting the
``cancel`` event (an ``asyncio.Event``; set at any point, even if cleared
again) stops scheduling — no continuation runs after it — and returns a
``"cancelled"`` :class:`JobResult`; cancelling the task cleans up the job's own tasks and
re-raises. A request whose call began is reported as ``sent`` — never as
rolled back: whatever the provider did with it stays done.
**What "independent" means.** Computational, not statistical: questions
batched together share a call, nothing more. Their answers are separate
marginals; the job never multiplies them into a joint probability.
nnx.decisions.DecisionJob¶
class nnx.decisions.DecisionJob(node: '_Node') -> 'None'
An immutable, deferred description of decision work — see the module docstring. Build with ask and collect, transform with map, chain with then; run with run / arun.
nnx.decisions.DecisionJob.ask¶
nnx.decisions.DecisionJob.ask(question: 'Question', *, id: 'str', policy: 'Any' = None, model_id: 'Optional[str]' = None) -> 'DecisionJob'
One question over the state. id keys its outcome (unique in a run); an abstention policy (with the model_id it was tuned for) turns each row's result into a selective decision.
nnx.decisions.DecisionJob.collect¶
nnx.decisions.DecisionJob.collect(jobs: 'Mapping[str, DecisionJob]') -> 'DecisionJob'
Independent jobs; the value maps each key to its job's value, in jobs' order.
nnx.decisions.DecisionJob.map¶
nnx.decisions.DecisionJob.map(self, fn: 'Callable[[Any], Any]') -> 'DecisionJob'
This job with fn applied to its value (pure: no provider calls).
nnx.decisions.DecisionJob.then¶
nnx.decisions.DecisionJob.then(self, fn: 'Callable[[Any], Follow]') -> 'DecisionJob'
A dependent job: after this one succeeds, fn(value) returns a Follow (the next job and its state), run once.
nnx.decisions.DecisionJob.question_ids¶
nnx.decisions.DecisionJob.question_ids(self) -> 'tuple[str, ...]'
The ids of the questions known before running (not those a continuation will create), in scheduling order.
nnx.decisions.DecisionJob.state¶
nnx.decisions.DecisionJob.state(self) -> 'dict[str, Any]'
Plain data for a job of ask / collect only. A job holding a function (map, then) is a runtime value: refused.
nnx.decisions.DecisionJob.run¶
nnx.decisions.DecisionJob.run(self, provider: 'Any', *, state: 'Any', limits: 'Optional[Limits]' = None) -> 'JobResult'
Run the job synchronously with provider (borrowed, never closed) over state.
nnx.decisions.DecisionJob.arun¶
nnx.decisions.DecisionJob.arun(self, provider: 'Any', *, state: 'Any', limits: 'Optional[Limits]' = None, cancel: 'Optional[asyncio.Event]' = None) -> 'JobResult'
Run the job on the event loop: up to limits.max_concurrency calls at once; setting cancel stops scheduling and returns a cancelled result (see the module docstring).
nnx.decisions.Follow¶
class nnx.decisions.Follow(job: 'DecisionJob', state: 'Any') -> 'None'
What a continuation returns: the next job and the state it reads (the answer mapped explicitly into follow-up inputs).
nnx.decisions.Limits¶
class nnx.decisions.Limits(max_questions: 'Optional[int]' = None, max_tokens: 'Optional[int]' = None, max_depth: 'int' = 8, max_requests: 'Optional[int]' = None, timeout: 'Optional[float]' = None, max_concurrency: 'int' = 1) -> 'None'
What a run may do.
Details
Attributes:
max_questions: the most questions per provider call (``None``: the
provider's own cap, or any number when it declares none; a
provider without ``decide_many`` answers one per call).
max_tokens: the most tokens per provider call — enforceable only by
a provider with ``count_tokens(questions, state)``; any other is
refused before any call.
max_depth: the most nested continuations (``then``).
max_requests: the most provider calls in the run (``None``:
unbounded).
timeout: seconds for the whole run (``None``: none). ``run`` checks
it between calls and provider hooks; ``arun`` also cancels
in-flight asynchronous calls (a synchronous provider's call in
its worker thread is waited for).
max_concurrency: ``arun``'s concurrent provider calls.
nnx.decisions.JobResult¶
class nnx.decisions.JobResult(value: 'Any', outcomes: 'Mapping[str, QuestionOutcome]', calls: 'int', status: 'str' = 'completed') -> 'None'
A finished run: the job's value, every question's outcome in the order it was scheduled, the provider calls made and the status ("completed" or "cancelled"; a cancelled run has no value).
nnx.decisions.QuestionOutcome¶
class nnx.decisions.QuestionOutcome(id: 'str', kind: 'str', rows: 'tuple[RowOutcome, ...]' = (), error: 'Optional[BaseException]' = None, sent: 'bool' = False) -> 'None'
What happened to one question.
Details
``kind`` is ``"answered"`` (the provider answered; rows may still
abstain), ``"failed"`` (its provider call raised — ``error``),
``"skipped"`` (never sent: a failure, a limit or the timeout stopped
scheduling) or ``"cancelled"`` (the run was cancelled, a sibling call
failed first, or the timeout elapsed while its call was in flight in
``arun``; ``sent`` says whether the request had already gone out —
a sent request is not rolled back).
nnx.decisions.QuestionOutcome.abstained¶
property nnx.decisions.QuestionOutcome.abstained
No public description is currently available.
nnx.decisions.RowOutcome¶
class nnx.decisions.RowOutcome(kind: 'str', result: 'DecisionResult', decision: 'Any' = None) -> 'None'
One input row's answer: "answered", or "abstained" when the question's abstention policy declined it (the result is kept either way).
nnx.decisions.JobError¶
class nnx.decisions.JobError(message: 'str', *, outcomes: 'Optional[Mapping[str, QuestionOutcome]]' = None) -> 'None'
Base of the decision-job errors. outcomes maps question ids to every QuestionOutcome known when the error was raised (answered, failed, skipped or cancelled — a request already sent says so), in scheduling order; completed is the answered subset. Both are empty when the job was refused before any call.
Details
The error pickles (a process-pool worker's error reaches its parent): a
provider error that would not survive the round trip — on its own or as
a ``__cause__`` — is replaced by a :class:`ProviderFailure` naming it,
and an answer's provider ``raw`` output that would not is dropped.
nnx.decisions.InvalidJob¶
class nnx.decisions.InvalidJob(message: 'str', *, outcomes: 'Optional[Mapping[str, QuestionOutcome]]' = None) -> 'None'
A job that cannot run as described: duplicate question ids, an unsupported question, a token cap the provider cannot enforce, a continuation that returns something other than a Follow.
nnx.decisions.JobFailed¶
class nnx.decisions.JobFailed(message: 'str', *, outcomes: 'Mapping[str, QuestionOutcome]', failed: 'Sequence[str]', skipped: 'Sequence[str]') -> 'None'
A provider call failed (fail-fast): failed names the questions of that call, skipped the known questions that were never sent (none when the job was refused before any call); the provider's error is the __cause__.
nnx.decisions.JobLimitExceeded¶
class nnx.decisions.JobLimitExceeded(message: 'str', *, outcomes: 'Optional[Mapping[str, QuestionOutcome]]' = None) -> 'None'
A request or depth limit stopped the job before the next call.
nnx.decisions.JobTimeout¶
class nnx.decisions.JobTimeout(message: 'str', *, outcomes: 'Optional[Mapping[str, QuestionOutcome]]' = None) -> 'None'
The run's timeout elapsed; in-flight requests were sent and are not rolled back.
nnx.decisions.benchmark.Sample¶
class nnx.decisions.benchmark.Sample(id: 'str', question: 'Question', input: 'Any', label: 'Union[str, bool]', family: 'str' = 'default', heldout: 'bool' = False, group: 'Optional[str]' = None, perturbation: 'Optional[str]' = None) -> 'None'
One benchmark item.
Details
Attributes:
id: the sample id — how records join (never row position).
question: the typed question; its digest is the schema identity.
input: what the provider reads for this sample (a text, a tensor
row, an array row).
label: the true option id (Choice / Score) or ``True`` / ``False``
(Boolean).
family: the task family.
heldout: whether the family is held out (reported apart from
in-family rows).
group: the grouping unit intervals resample (default: the id).
perturbation: what produced this sample from another (``None``:
an original).
nnx.decisions.benchmark.Sample.digest¶
property nnx.decisions.benchmark.Sample.digest
No public description is currently available.
nnx.decisions.benchmark.Sample.unit¶
property nnx.decisions.benchmark.Sample.unit
The grouping unit (group, else the id).
nnx.decisions.benchmark.Sample.true_label¶
property nnx.decisions.benchmark.Sample.true_label
No public description is currently available.
nnx.decisions.benchmark.Record¶
class nnx.decisions.benchmark.Record(sample_id: 'str', question_digest: 'str', provider: 'str', status: 'str', distribution: 'Optional[tuple[tuple[str, float], ...]]' = None, p_true: 'Optional[float]' = None, reason: 'Optional[str]' = None, revision: 'Optional[str]' = None, prompt_identity: 'Optional[str]' = None, execution: 'Mapping[str, Any]' = <factory>) -> 'None'
One provider output for one sample.
Details
``status`` is ``"answered"`` (with ``distribution`` for a Choice or
Score, ``p_true`` for a Boolean), ``"unsupported"`` (the provider
declared it cannot serve the request) or ``"failed"`` (the call raised);
the latter two carry a ``reason``. ``revision`` is the model revision,
``prompt_identity`` how the provider rendered the question (a digest of
its settings), ``execution`` free JSON metadata (batch, timing, whether
the budget cut the batch short).
nnx.decisions.benchmark.Record.state¶
nnx.decisions.benchmark.Record.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.benchmark.Record.from_state¶
nnx.decisions.benchmark.Record.from_state(state: 'Mapping[str, Any]') -> 'Record'
No public description is currently available.
nnx.decisions.benchmark.write_records¶
nnx.decisions.benchmark.write_records(path: 'Union[str, os.PathLike[str]]', records: 'Iterable[Record]') -> 'None'
One JSON object per line (strict JSON, sorted keys).
nnx.decisions.benchmark.read_records¶
nnx.decisions.benchmark.read_records(path: 'Union[str, os.PathLike[str]]') -> 'list[Record]'
The records of a JSONL file written by write_records.
nnx.decisions.benchmark.Budget¶
class nnx.decisions.benchmark.Budget(max_calls: 'int', max_samples: 'Optional[int]' = None) -> 'None'
What a live collection may spend: max_calls provider calls and, optionally, max_samples samples. A batch the sample budget cuts short is sent shortened and recorded as partial_batch.
nnx.decisions.benchmark.collect¶
nnx.decisions.benchmark.collect(provider: 'Any', samples: 'Sequence[Sample]', *, provider_id: 'str', budget: 'Budget', batch_size: 'int' = 16, revision: 'Optional[str]' = None, prompt_identity: 'Optional[str]' = None) -> 'Collection'
Ask provider every sample's question, in batches of samples that share one question (at most batch_size inputs per call).
Details
The provider and its id and the :class:`Budget` are explicit. Each batch
is attempted once: a provider that raises gives ``"failed"`` records
(retries are the provider's own), a request the provider declares it
cannot serve (its ``check(question, inputs)``, or an
``UnsupportedCapability`` from ``decide``) gives ``"unsupported"``
records. Collection stops when the budget is spent; samples never
attempted have no record (``missing`` when evaluated).
``prompt_identity`` defaults to a digest of the provider's ``record()``.
nnx.decisions.benchmark.Collection¶
class nnx.decisions.benchmark.Collection(records: 'tuple[Record, ...]', calls: 'int', complete: 'bool', stopped: 'Optional[str]' = None) -> 'None'
A live collection: the records, the provider calls made, whether every sample was attempted (complete) and, if not, why.
nnx.decisions.benchmark.evaluate¶
nnx.decisions.benchmark.evaluate(samples: 'Sequence[Sample]', records: 'Sequence[Record]', *, split: 'str', provider: 'Optional[str]' = None, epsilon: 'Optional[float]' = None, n_bins: 'int' = 10, policy: 'Any' = None, model_id: 'Optional[str]' = None, resources: 'Optional[Resources]' = None) -> 'BenchmarkReport'
Replay: join records to samples by sample id and question digest and report every slice. Calls no provider and fits nothing.
Details
Args:
split: the benchmark split's identity (reports of different splits
are never compared).
provider: which provider's records to score (``None``: all records,
which must then come from one provider).
epsilon: the NLL floor; ``None`` (default) is exact — ``+inf`` when a
true label has probability 0.
n_bins: reliability bins.
policy / model_id: an ``nnx.abstention.AbstentionPolicy`` applied as
given (never tuned here) for selective coverage and risk.
resources: how cost was measured, if at all.
nnx.decisions.benchmark.BenchmarkReport¶
class nnx.decisions.benchmark.BenchmarkReport(split: 'str', provider: 'Optional[str]', metric_identity: 'Mapping[str, Any]', slices: 'Mapping[str, SliceReport]', resources: 'Resources' = <factory>, extra: 'int' = 0, sample_set: 'Optional[str]' = None) -> 'None'
A replayed benchmark: the split and metric identity, every slice's coverage, metrics and bins, and the declared Resources.
nnx.decisions.benchmark.BenchmarkReport.state¶
nnx.decisions.benchmark.BenchmarkReport.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.benchmark.BenchmarkReport.to_json¶
nnx.decisions.benchmark.BenchmarkReport.to_json(self) -> 'str'
No public description is currently available.
nnx.decisions.benchmark.BenchmarkReport.csv_rows¶
nnx.decisions.benchmark.BenchmarkReport.csv_rows(self) -> 'list[dict[str, Any]]'
One row per slice and metric: value, unit, denominator, the unavailable reason, an available value's note, every coverage count of the slice (samples, eligible, failed, unsupported, missing, duplicate, mismatched, invalid — they add up to the slice's samples) and the report's extra records. A non-finite value is written as in the JSON.
nnx.decisions.benchmark.BenchmarkReport.to_csv¶
nnx.decisions.benchmark.BenchmarkReport.to_csv(self) -> 'str'
No public description is currently available.
nnx.decisions.benchmark.BenchmarkReport.text¶
nnx.decisions.benchmark.BenchmarkReport.text(self) -> 'str'
No public description is currently available.
nnx.decisions.benchmark.BenchmarkReport.save¶
nnx.decisions.benchmark.BenchmarkReport.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write the JSON report atomically (serialized first: a report that cannot be written never replaces an existing file).
nnx.decisions.benchmark.BenchmarkReport.load_state¶
nnx.decisions.benchmark.BenchmarkReport.load_state(path: 'Union[str, os.PathLike[str]]') -> 'dict[str, Any]'
A saved report's JSON state (for compare_reports).
nnx.decisions.benchmark.SliceReport¶
class nnx.decisions.benchmark.SliceReport(name: 'str', coverage: 'Coverage', metrics: 'Mapping[str, MetricValue]', bins: 'tuple[Any, ...]') -> 'None'
One slice: its coverage, every metric and the reliability bins.
nnx.decisions.benchmark.SliceReport.state¶
nnx.decisions.benchmark.SliceReport.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.benchmark.Coverage¶
class nnx.decisions.benchmark.Coverage(samples: 'int', eligible: 'int', missing: 'int', duplicate: 'int', mismatched: 'int', invalid: 'int', unsupported: 'int', failed: 'int', reasons: 'Mapping[str, int]' = <factory>) -> 'None'
How the samples of a slice joined their records: eligible (one answered record with the sample's digest), missing (no record), duplicate (several), mismatched (a record for another question digest), invalid (an answered record that does not fit its question: options in another order or set, a Boolean answer for a Choice), unsupported and failed (with their reasons). Records for no sample are the report's extra count.
nnx.decisions.benchmark.Coverage.state¶
nnx.decisions.benchmark.Coverage.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.benchmark.MetricValue¶
class nnx.decisions.benchmark.MetricValue(name: 'str', value: 'Optional[float]', denominator: 'int', reason: 'Optional[str]' = None, note: 'Optional[str]' = None) -> 'None'
A metric's value over denominator rows, or None with the reason it is unavailable. note qualifies an available value (rows it does not count, say).
nnx.decisions.benchmark.MetricValue.unit¶
property nnx.decisions.benchmark.MetricValue.unit
No public description is currently available.
nnx.decisions.benchmark.MetricValue.state¶
nnx.decisions.benchmark.MetricValue.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.benchmark.Resources¶
class nnx.decisions.benchmark.Resources(warmup: 'Optional[int]' = None, hardware: 'Optional[str]' = None, timing_boundary: 'Optional[str]' = None, concurrency: 'Optional[int]' = None, batch_count: 'Optional[int]' = None, seconds: 'Optional[float]' = None, source: 'Optional[str]' = None) -> 'None'
How cost was obtained — every field None unless declared.
Details
``source`` is ``"measured"`` (timed here; ``hardware`` is then required)
or ``"supplied"`` (reported by someone else); ``timing_boundary`` says
what the time covers (for example ``"provider.decide calls only"``).
nnx.decisions.benchmark.Resources.state¶
nnx.decisions.benchmark.Resources.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.decisions.benchmark.bootstrap_interval¶
nnx.decisions.benchmark.bootstrap_interval(samples: 'Sequence[Sample]', records: 'Sequence[Record]', *, metric: 'str', seed: 'int', slice: 'str' = 'in_family', resamples: 'int' = 1000, level: 'float' = 0.95, epsilon: 'Optional[float]' = None, provider: 'Optional[str]' = None) -> 'Interval'
Bootstrap metric (accuracy, macro-F1, NLL or Brier) over the eligible rows of slice (of provider's records, as evaluate selects them), resampling grouping units — every row of a drawn unit comes along — with numpy.random.default_rng(seed). The bounds are resampled values (no interpolation); a non-finite bound (an exact NLL of +inf) flags the interval degenerate.
nnx.decisions.benchmark.Interval¶
class nnx.decisions.benchmark.Interval(metric: 'str', estimate: 'Optional[float]', low: 'Optional[float]', high: 'Optional[float]', level: 'float', seed: 'int', resamples: 'int', units: 'int', degenerate: 'bool', unit: 'str' = 'group', method: 'str' = 'percentile bootstrap, grouping units resampled with replacement', reason: 'Optional[str]' = None) -> 'None'
A percentile bootstrap interval of one metric, resampling the declared grouping units (unit="group") with seed. degenerate flags an interval that carries no spread information: fewer than two units, every resample giving the same value, bounds that coincide (low == high) while the resamples still vary, or a non-finite bound.
nnx.decisions.benchmark.Interval.state¶
nnx.decisions.benchmark.Interval.state(self) -> 'dict[str, Any]'
Strict JSON: a non-finite bound is written as MetricValue writes it.
nnx.decisions.benchmark.compare_reports¶
nnx.decisions.benchmark.compare_reports(a: 'Union[BenchmarkReport, Mapping[str, Any]]', b: 'Union[BenchmarkReport, Mapping[str, Any]]') -> 'dict[str, dict[str, Optional[float]]]'
b - a per slice and metric, for two reports of the same split and metric identity and — when both record it — the same sample set (a BenchmarkReport or a saved report's state); anything else raises BenchmarkError. A metric unavailable on either side has no delta.
nnx.decisions.benchmark.permute_options¶
nnx.decisions.benchmark.permute_options(sample: 'Sample', order: 'Sequence[str]', *, id: 'Optional[str]' = None) -> 'Sample'
The same options in another order (order: every option id once). A Score's levels are ordered, so it is refused.
nnx.decisions.benchmark.redescribe¶
nnx.decisions.benchmark.redescribe(sample: 'Sample', descriptions: 'Mapping[str, str]', *, id: 'Optional[str]' = None) -> 'Sample'
New descriptions for some options (id -> description): the label space is the same, the text the provider sees is not.
nnx.decisions.benchmark.add_distractors¶
nnx.decisions.benchmark.add_distractors(sample: 'Sample', distractors: 'Sequence[Union[Option, tuple[str, str]]]', *, id: 'Optional[str]' = None) -> 'Sample'
Extra wrong options appended to a Choice.
nnx.decisions.benchmark.add_none_of_the_above¶
nnx.decisions.benchmark.add_none_of_the_above(sample: 'Sample', option: 'Union[Option, tuple[str, str]]' = ('none', 'None of the above'), *, remove_label: 'bool' = True, id: 'Optional[str]' = None) -> 'Sample'
A none-of-the-above option; with remove_label the true option is removed, so none-of-the-above becomes the answer.
nnx.decisions.benchmark.add_context¶
nnx.decisions.benchmark.add_context(sample: 'Sample', context: 'str', *, kind: 'str' = 'long_context', id: 'Optional[str]' = None) -> 'Sample'
Irrelevant context appended to a text input.
nnx.decisions.benchmark.rewrite_input¶
nnx.decisions.benchmark.rewrite_input(sample: 'Sample', text: 'str', *, kind: 'str', id: 'Optional[str]' = None) -> 'Sample'
The same question over a rewritten text input — a translation (kind="multilingual") or an adversarial rewrite (kind="adversarial").
nnx.decisions.benchmark.BenchmarkError¶
class nnx.decisions.benchmark.BenchmarkError
A malformed sample, record, budget or report, or reports that cannot be compared.
2.12. Errors as values at fallible boundaries (nnx.result)¶
Opt-in; see Concepts §27.
nnx.result.Ok¶
class nnx.result.Ok(value: 'T') -> 'None'
A success holding value (any value: 0, None and False included).
nnx.result.Ok.is_ok¶
property nnx.result.Ok.is_ok
No public description is currently available.
nnx.result.Ok.is_err¶
property nnx.result.Ok.is_err
No public description is currently available.
nnx.result.Ok.map¶
nnx.result.Ok.map(self, fn: 'Callable[[T], U]') -> 'Ok[U]'
Ok(fn(value)).
nnx.result.Ok.bind¶
nnx.result.Ok.bind(self, fn: 'Callable[[T], Result[U, E]]') -> 'Result[U, E]'
fn(value), which must be an Ok or an Err.
nnx.result.Ok.map_error¶
nnx.result.Ok.map_error(self, fn: 'Callable[[Any], Any]') -> 'Ok[T]'
Unchanged: there is no error to map.
nnx.result.Ok.recover¶
nnx.result.Ok.recover(self, fn: 'Callable[[Any], Any]') -> 'Ok[T]'
Unchanged: there is nothing to recover from.
nnx.result.Ok.unwrap¶
nnx.result.Ok.unwrap(self) -> 'T'
No public description is currently available.
nnx.result.Err¶
class nnx.result.Err(error: 'E') -> 'None'
A failure holding a typed error.
nnx.result.Err.is_ok¶
property nnx.result.Err.is_ok
No public description is currently available.
nnx.result.Err.is_err¶
property nnx.result.Err.is_err
No public description is currently available.
nnx.result.Err.map¶
nnx.result.Err.map(self, fn: 'Callable[[Any], Any]') -> 'Err[E]'
Unchanged: there is no value to map.
nnx.result.Err.bind¶
nnx.result.Err.bind(self, fn: 'Callable[[Any], Any]') -> 'Err[E]'
Unchanged: there is no value to bind.
nnx.result.Err.map_error¶
nnx.result.Err.map_error(self, fn: 'Callable[[E], F]') -> 'Err[F]'
Err(fn(error)).
nnx.result.Err.recover¶
nnx.result.Err.recover(self, fn: 'Callable[[E], Result[T, F]]') -> 'Result[T, F]'
fn(error), which must be an Ok or an Err.
nnx.result.Err.unwrap¶
nnx.result.Err.unwrap(self) -> 'NoReturn'
Raise UnwrapError carrying the error (and chaining it when it is an exception).
nnx.result.Result¶
type alias nnx.result.Result
Public type alias.
nnx.result.UnwrapError¶
class nnx.result.UnwrapError(error: 'Any') -> 'None'
Err.unwrap was called; error is the typed error (and the __cause__ when it is an exception).
nnx.result.BoundaryError¶
class nnx.result.BoundaryError(code: 'str', where: 'str', context: 'Mapping[str, Any]' = <factory>, cause: 'Optional[BaseException]' = None) -> 'None'
A typed failure at a boundary: a stable code, where it happened (a field name or a path), a read-only context mapping and the original exception as cause.
Details
Equality compares ``code``, ``where`` and ``context`` (not ``cause``:
exceptions compare by identity); the hash uses ``code`` and ``where``.
It pickles and deep-copies (``context`` as a plain mapping); being
read-only, ``context`` is not supported by ``dataclasses.asdict``.
nnx.result.validate_decision_request_result¶
nnx.result.validate_decision_request_result(kind: 'str', prompt: 'str', options: 'Optional[Sequence[Any]]' = None) -> 'Result[Any, BoundaryError]'
Build a typed decision question from plain data.
Details
``kind`` is ``"choice"``, ``"boolean"`` or ``"score"``; ``options`` are
``(id, description)`` pairs (a Choice's options, a Score's levels, none
for a Boolean). The question objects are built inside the boundary, so a
malformed request — duplicate or empty ids, too few options, an empty
prompt, an unknown kind — is ``Err(BoundaryError(code=
"invalid_decision_request"))`` and nothing is ever sent to a provider.
Only :class:`~nnx.decisions.InvalidDecisionRequest` is converted.
nnx.result.decide_result¶
nnx.result.decide_result(provider: 'Any', question: 'Any', inputs: 'Any', *, policy: 'Any' = None, model_id: 'Optional[str]' = None) -> 'Result[tuple[Any, ...], BoundaryError]'
Ask provider question for inputs.
Details
``Ok`` holds the results, one per input (an answer of another count,
question, kind or option order is an ``"invalid_decision_response"``) — or, with an abstention
``policy`` (and the ``model_id`` it was tuned for), each row's
``SelectiveDecision``: an abstention is a success whose probabilities
and reason are kept. Only the provider's own failures are errors:
``InvalidDecisionResponse`` (an answer that does not fit the question)
becomes ``"invalid_decision_response"``, any other ``ProviderFailure``
(an outage, a timeout, a refused credential) ``"provider_failure"``, and
``UnsupportedCapability`` ``"unsupported"``; the first two keep the
error's ``request_id`` in ``context``. The codes follow the exception
type, not whether a retry would help. Anything else propagates. A
``policy`` that does not fit the question (a ``Boolean``, other labels,
another ``model_id``) raises before the provider is called.
nnx.result.inspect_bundle_result¶
nnx.result.inspect_bundle_result(path: 'Union[str, os.PathLike[str]]') -> 'Result[Any, BoundaryError]'
Summarize the run bundle at path with nnx.bundles.inspect_bundle — manifest and JSON records only: it never unpickles a checkpoint, loads a tensor, calls a factory or a registry.
Details
``Ok`` holds the ``BundleInfo``. A path that does not exist is
``Err(BoundaryError(code="artifact_missing"))``; any other
``BundleError`` (not a bundle, a malformed or tampered manifest) is
``"bundle_invalid"``, its cause kept. Other exceptions propagate —
a ``PermissionError`` on the path itself, or a ``TypeError`` for a
``path`` that is not a str or ``os.PathLike`` (``where`` is always a
str).
2.13. Opt-in compilation (nnx.compilation)¶
train(compile=...) compiles the built-in FP32 forward; see Concepts §28.
nnx.compilation.CompileSpec¶
class nnx.compilation.CompileSpec(*, backend: 'str' = 'inductor', mode: 'Optional[str]' = None, fullgraph: 'bool' = False, dynamic: 'Optional[bool]' = None, on_failure: 'str' = 'error') -> 'None'
What to compile with: a torch.compile backend (default "inductor"), its mode (one of COMPILE_MODES, or None), fullgraph (reject graph breaks), dynamic (None lets torch decide) and on_failure ("error" or "eager").
nnx.compilation.CompileSpec.options¶
nnx.compilation.CompileSpec.options(self) -> 'dict[str, Any]'
The keyword options passed to torch.compile.
nnx.compilation.CompileSpec.state¶
nnx.compilation.CompileSpec.state(self) -> 'dict[str, Any]'
Fields that differ from their defaults ({} for the defaults).
nnx.compilation.CompileSpec.from_state¶
nnx.compilation.CompileSpec.from_state(state: 'Mapping[str, Any]') -> 'CompileSpec'
No public description is currently available.
nnx.compilation.CompileRecord¶
class nnx.compilation.CompileRecord(*, backend: 'str', options: 'Mapping[str, Any]', on_failure: 'str', torch_version: 'str', effective: 'str' = 'pending', graph_breaks: 'Optional[int]' = None, restart: 'Optional[Mapping[str, Any]]' = None) -> 'None'
What a run asked to compile and what took effect.
Details
``backend`` / ``options`` / ``on_failure`` are the request and
``torch_version`` the framework that served it. ``effective`` is
``pending`` (no compiled forward has run yet), ``compiled``, ``eager``
(restarted eager under ``on_failure="eager"``) or ``failed``.
``graph_breaks`` counts the breaks of the first compiled forward
(``capture`` is ``"full"`` or ``"partial"``); ``restart`` holds where
and why an eager restart (or a failure) happened.
nnx.compilation.CompileRecord.requested¶
property nnx.compilation.CompileRecord.requested
Always True: a record exists only when compilation was requested.
nnx.compilation.CompileRecord.capture¶
property nnx.compilation.CompileRecord.capture
"full" / "partial" once a compiled forward ran, else None.
nnx.compilation.CompileRecord.record¶
nnx.compilation.CompileRecord.record(self) -> 'dict[str, Any]'
The plain mapping stored in metadata, checkpoints and attempts.
nnx.compilation.CompileRecord.from_record¶
nnx.compilation.CompileRecord.from_record(record: 'Mapping[str, Any]') -> 'CompileRecord'
No public description is currently available.
nnx.compilation.CompileFailed¶
class nnx.compilation.CompileFailed(message: 'str', record: 'CompileRecord') -> 'None'
A compiled forward failed to compile under on_failure="error"; __cause__ is torch's error and record the run's CompileRecord (effective state failed).
2.14. Forward benchmarks and profiling (nnx.benchmarking)¶
Bounded forward benchmarks and profiling.
nnx.benchmarking.benchmark_forward¶
nnx.benchmarking.benchmark_forward(net: 'torch.nn.Module', inputs: 'Inputs', *, compile: 'Optional[CompileSpec]' = None, warmup: 'int' = 3, repeats: 'int' = 10, device: 'Optional[Union[str, torch.device]]' = None, label: 'Optional[str]' = None) -> 'BenchmarkReport'
Time net's forward on inputs (eager, or compiled with compile): the first call, then repeats calls after warmup untimed ones, on a deep copy in eval mode under torch.no_grad. device defaults to the module's.
nnx.benchmarking.compare_compile¶
nnx.benchmarking.compare_compile(net: 'torch.nn.Module', inputs: 'Inputs', *, compile: 'Optional[CompileSpec]' = None, warmup: 'int' = 3, repeats: 'int' = 10, device: 'Optional[Union[str, torch.device]]' = None) -> 'CompileComparison'
Benchmark net eager and compiled (compile, default CompileSpec()) from identical weights.
nnx.benchmarking.profile_forward¶
nnx.benchmarking.profile_forward(net: 'torch.nn.Module', inputs: 'Inputs', *, output_dir: 'Union[str, os.PathLike[str]]', wait: 'int' = 1, warmup: 'int' = 1, active: 'int' = 3, repeat: 'int' = 1, compile: 'Optional[CompileSpec]' = None, device: 'Optional[Union[str, torch.device]]' = None, row_limit: 'int' = 10) -> 'ProfileReport'
Profile net's forward with torch.profiler for exactly (wait + warmup + active) * repeat steps, writing Chrome traces to output_dir (created if needed). Every count is finite (active and repeat at least 1). The profiler is stopped and its traces flushed whether the steps finish, raise or are cancelled. With compile and wait=warmup=0 the first active step includes the compilation itself.
nnx.benchmarking.BenchmarkReport¶
class nnx.benchmarking.BenchmarkReport(*, label: 'str', backend: 'Optional[str]', device: 'str', shapes: 'tuple[tuple[int, ...], ...]', dtypes: 'tuple[str, ...]', batch_size: 'int', warmup: 'int', repeats: 'int', first_call_seconds: 'float', latencies: 'tuple[float, ...]', peak_memory_bytes: 'Optional[int]', weights_digest: 'str', torch_version: 'str') -> 'None'
One benchmarked configuration. Times are seconds; latencies are the warmed per-call times (repeats of them).
nnx.benchmarking.BenchmarkReport.mean_seconds¶
property nnx.benchmarking.BenchmarkReport.mean_seconds
No public description is currently available.
nnx.benchmarking.BenchmarkReport.median_seconds¶
property nnx.benchmarking.BenchmarkReport.median_seconds
No public description is currently available.
nnx.benchmarking.BenchmarkReport.stdev_seconds¶
property nnx.benchmarking.BenchmarkReport.stdev_seconds
Sample standard deviation (0.0 for a single repeat).
nnx.benchmarking.BenchmarkReport.min_seconds¶
property nnx.benchmarking.BenchmarkReport.min_seconds
No public description is currently available.
nnx.benchmarking.BenchmarkReport.max_seconds¶
property nnx.benchmarking.BenchmarkReport.max_seconds
No public description is currently available.
nnx.benchmarking.BenchmarkReport.throughput¶
property nnx.benchmarking.BenchmarkReport.throughput
Warmed samples per second (batch_size / mean_seconds), where batch_size is the first input's leading dimension (batch-first inputs).
nnx.benchmarking.BenchmarkReport.state¶
nnx.benchmarking.BenchmarkReport.state(self) -> 'dict[str, Any]'
A plain, JSON-ready summary.
nnx.benchmarking.CompileComparison¶
class nnx.benchmarking.CompileComparison(*, eager: 'BenchmarkReport', compiled: 'BenchmarkReport', max_abs_diff: 'float') -> 'None'
An eager and a compiled benchmark of identical weights, and the largest absolute difference between their outputs.
nnx.benchmarking.CompileComparison.speedup¶
property nnx.benchmarking.CompileComparison.speedup
Warmed eager mean over compiled mean (> 1 means faster compiled; nothing is guaranteed).
nnx.benchmarking.CompileComparison.state¶
nnx.benchmarking.CompileComparison.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.benchmarking.ProfileReport¶
class nnx.benchmarking.ProfileReport(*, output_dir: 'str', trace_files: 'tuple[str, ...]', schedule: 'Mapping[str, int]', steps: 'int', table: 'str') -> 'None'
A finished profile: where its traces were written, the schedule and step count, and the top operators by self time.
2.15. Single-node data parallelism (nnx.distributed)¶
Under torchrun; see Architecture §3.
nnx.distributed.init_process_group¶
nnx.distributed.init_process_group(backend: 'Optional[str]' = None, *, timeout_seconds: 'float' = 120.0) -> 'tuple[int, int]'
Join the process group torchrun described in the environment (env://) and return (rank, world_size). backend defaults to nccl when CUDA is available, else gloo; timeout_seconds bounds every collective, so a rank that dies cannot hang the others forever. With CUDA, each process takes cuda:LOCAL_RANK.
nnx.distributed.shutdown¶
nnx.distributed.shutdown(*, timeout_seconds: 'float' = 60.0) -> 'None'
Leave the process group together, then destroy it: the counterpart of init_process_group, called once training is over.
Details
Every rank meets at a teardown barrier before the group is destroyed, so
no rank tears its connections down while a peer still uses them (with
Gloo, a peer can otherwise abort: "terminate called without an active
exception"). With Gloo the wait is bounded by ``timeout_seconds``: when a
rank does not arrive in time the others raise ``RuntimeError`` instead of
blocking in the call, and leave the group to process exit (destroying it
then could block on the missing rank). The timed-out barrier stays queued,
so the process itself may still wait at exit until the missing rank
arrives (its barrier then completes against the queued one) or the
group's own timeout (``init_process_group(timeout_seconds=...)``) expires.
With NCCL, older torch releases ignore ``timeout_seconds`` unless
``TORCH_NCCL_BLOCKING_WAIT=1`` is set, so the wait is bounded by the
group's timeout instead. Without a process group, or after a successful
call, it does nothing; after a failed call the group is still
initialized (``init_process_group()`` returns it as is), so exit.
nnx.distributed.train_loader¶
nnx.distributed.train_loader(dataset: 'Any', batch_size: 'int', *, policy: 'str' = 'pad', seed: 'int' = 0, shuffle: 'bool' = True, collate_fn: 'Optional[Callable[[list[Any]], Any]]' = None) -> 'ShardedLoader'
This rank's training loader: a per-epoch seeded partition of dataset evened out by policy ("pad" or "drop"), so every rank takes the same number of steps.
nnx.distributed.validation_loader¶
nnx.distributed.validation_loader(dataset: 'Any', batch_size: 'int', *, collate_fn: 'Optional[Callable[[list[Any]], Any]]' = None) -> 'ShardedLoader'
This rank's validation shard: rows rank::world_size, unpadded — every row is scored exactly once across the ranks.
nnx.distributed.DDP¶
class nnx.distributed.DDP(*, writer_rank: 'int' = 0, find_unused_parameters: 'bool' = False) -> 'None'
model.train(params, distributed=DDP()): data-parallel training over the initialized process group. writer_rank owns every artifact; find_unused_parameters is passed to DistributedDataParallel.
nnx.distributed.writer_only¶
nnx.distributed.writer_only(factory: 'Callable[[], Any]') -> 'Any'
A callback NNx builds on the writer rank only — for callbacks that write when constructed (TensorBoardCallback, WandbCallback): callbacks=[writer_only(lambda: TensorBoardCallback("tb"))]. Without distributed=, the factory is simply called.
nnx.distributed.RankPartition¶
class nnx.distributed.RankPartition(n: 'int', *, rank: 'int', world_size: 'int', kind: 'str' = 'train', policy: 'str' = 'pad', shuffle: 'bool' = True, seed: 'int' = 0) -> 'None'
This rank's dataset indices, per epoch.
Details
``kind="train"``: a seeded permutation of ``range(n)`` (``seed + epoch``;
identity order with ``shuffle=False``), evened out by ``policy`` —
``"pad"`` repeats leading indices, ``"drop"`` drops the tail — and
dealt ``rank::world_size``, so every rank serves the same count.
``kind="validation"``: rows ``rank::world_size`` in order, unpadded
(counts may differ; a rank may serve none).
nnx.distributed.RankPartition.set_epoch¶
nnx.distributed.RankPartition.set_epoch(self, epoch: 'int') -> 'None'
No public description is currently available.
nnx.distributed.RankPartition.global_order¶
nnx.distributed.RankPartition.global_order(self) -> 'list[int]'
Every rank's indices for this epoch, before dealing (train: after padding or dropping).
nnx.distributed.RankPartition.indices¶
nnx.distributed.RankPartition.indices(self) -> 'list[int]'
No public description is currently available.
nnx.distributed.RankPartition.descriptor¶
nnx.distributed.RankPartition.descriptor(self) -> 'dict[str, Any]'
What a resume must match: the dataset size, world, kind, policy, shuffle and seed (never the epoch or rank).
nnx.distributed.ShardedLoader¶
class nnx.distributed.ShardedLoader(dataset: 'Any', batch_size: 'int', partition: 'RankPartition', *, collate_fn: 'Optional[Callable[[list[Any]], Any]]' = None) -> 'None'
A re-iterable loader over one rank's RankPartition of a map-style dataset (DataLoader underneath, in-process, ordered). set_epoch reseeds a train partition; ids() names the dataset rows this rank serves this epoch.
nnx.distributed.ShardedLoader.set_epoch¶
nnx.distributed.ShardedLoader.set_epoch(self, epoch: 'int') -> 'None'
No public description is currently available.
nnx.distributed.ShardedLoader.ids¶
nnx.distributed.ShardedLoader.ids(self) -> 'list[int]'
No public description is currently available.
nnx.distributed.ShardedLoader.descriptor¶
nnx.distributed.ShardedLoader.descriptor(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.distributed.DistributedFailure¶
class nnx.distributed.DistributedFailure
Raised on every rank when another rank failed (preflight, a non-finite loss, a step-count mismatch, a writer-side commit): the failing rank raises its own error, the others this one naming it.
2.16. Budgeted experiment search (nnx.search)¶
Through Optuna (extra optuna); see Concepts §29.
nnx.search¶
module nnx.search
Budgeted experiment search through an optional Optuna adapter.
Details
``search(plan, space, apply=..., monitor=..., budget=..., study_name=...)``
asks an Optuna study for parameter sets one at a time; ``apply(plan,
params)`` turns each into a **fresh** :class:`~nnx.plans.ExperimentPlan`
(the base plan is never changed), which is fitted — a new model, optimizer,
loaders and callbacks per trial, its own attempt id and run directory — and
scored by one declared **validation** monitor. Requires the ``optuna``
extra (``pip install "thekaveh-nnx[optuna]"``); importing ``nnx`` never
imports Optuna or creates a study.
- **Validated first.** The space (bounds, steps, log scales, choices), the
budget, the monitor (a validation monitor with a direction, equal to the
plan's own ``NNTrainParams.monitor`` so the BEST checkpoint is chosen the
same way) and the plan (validation data, data and callbacks as factories)
are checked before any trial; reopening a stored study whose space,
monitor or plan identity differs is refused before any trial too.
- **The budget counts every trial.** Completed, failed, pruned and orphaned
(left running by a crashed process) trials all count against
``SearchBudget.trials``: a study holding two trials, searched with a
budget of 3, asks once.
- **Per-trial caps.** ``updates_per_trial`` stops a trial right after that
many committed optimizer updates (a partial accumulation window is not an
update). ``deadline_seconds`` (measured on ``clock``, from the search's
start) starts no new trial once expired and stops a running one at its
next update or epoch boundary, recording the overrun. Neither is
preemptive: a long batch finishes first. A trial stopped mid-epoch keeps
only its committed epochs — LAST, its tensors and the history are those
of the last completed epoch (no LAST before a first completed one).
- **Pruning** (Optuna's ``pruner``) sees only the monitor's validation
value of each completed epoch, reported at step = the epoch index
(strictly increasing). A missing or non-finite value fails the trial; it
is never a candidate. A trial's value is the run's own monitor decision:
the value of the epoch its BEST checkpoint holds (``min_delta`` included).
- **Errors.** A fit that raises fails its trial (``TrialOutcome.error``
keeps the message) and the search continues; a configuration error —
``apply`` raising or returning a plan the search cannot score — stops the
search. An exhausted sampler ends it (``stop_reason="exhausted"``).
- **Outcomes.** Every trial's :class:`TrialOutcome` keeps its parameters,
state and terminal reason, resource counters (epochs, committed updates,
seconds), best monitor value, attempt id, run id and BEST checkpoint path.
:attr:`SearchResult.best` is the completed trial with the best finite
monitor value under the monitor's direction.
- **Cancellation** (``KeyboardInterrupt``) runs the fit's callback cleanup,
releases its run lease, records the trial as failed and propagates.
Trials run sequentially in this process, which is the study's only
writer: inside a ``torch.distributed`` process group of more than one rank
(a ``torchrun`` launch) ``search`` raises ``RuntimeError`` before loading or
creating a study, rather than run one study per rank. NNx artifacts live
under ``runs/`` and Optuna's in ``storage`` (an in-memory study when
``None``). Samplers keep their own limits (a seeded
``TPESampler`` is reproducible only sequentially).
nnx.search.search¶
nnx.search.search(plan: 'Any', space: 'SearchSpace', *, apply: 'Callable[[Any, Mapping[str, Any]], Any]', monitor: 'Any', budget: 'SearchBudget', study_name: 'str', storage: 'Optional[str]' = None, sampler: 'Any' = None, pruner: 'Any' = None, clock: 'Callable[[], float]' = <built-in function monotonic>) -> 'SearchResult'
Run the budgeted search; see the module documentation.
nnx.search.SearchSpace¶
class nnx.search.SearchSpace(*params: 'SearchParam') -> 'None'
The parameters a trial draws, by unique name.
nnx.search.SearchSpace.state¶
nnx.search.SearchSpace.state(self) -> 'list[dict[str, Any]]'
No public description is currently available.
nnx.search.SearchSpace.digest¶
nnx.search.SearchSpace.digest(self) -> 'str'
No public description is currently available.
nnx.search.FloatParam¶
class nnx.search.FloatParam(*, name: 'str', low: 'float', high: 'float', log: 'bool' = False, step: 'Optional[float]' = None) -> 'None'
A float in [low, high]: log samples on a log scale (low > 0), step discretizes (not with log).
nnx.search.FloatParam.state¶
nnx.search.FloatParam.state(self) -> 'dict[str, Any]'
Fields that differ from their defaults; bounds as floats (so low=0 and low=0.0 are one space).
nnx.search.FloatParam.suggest¶
nnx.search.FloatParam.suggest(self, trial: 'Any') -> 'float'
No public description is currently available.
nnx.search.IntParam¶
class nnx.search.IntParam(*, name: 'str', low: 'int', high: 'int', log: 'bool' = False, step: 'int' = 1) -> 'None'
An int in [low, high] (inclusive): log (low >= 1) or a step.
nnx.search.IntParam.state¶
nnx.search.IntParam.state(self) -> 'dict[str, Any]'
Fields that differ from their defaults.
nnx.search.IntParam.suggest¶
nnx.search.IntParam.suggest(self, trial: 'Any') -> 'int'
No public description is currently available.
nnx.search.CategoricalParam¶
class nnx.search.CategoricalParam(*, name: 'str', choices: 'tuple[Any, ...]') -> 'None'
One of choices (distinct None / bool / int / float / str).
nnx.search.CategoricalParam.state¶
nnx.search.CategoricalParam.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.search.CategoricalParam.suggest¶
nnx.search.CategoricalParam.suggest(self, trial: 'Any') -> 'Any'
No public description is currently available.
nnx.search.SearchBudget¶
class nnx.search.SearchBudget(*, trials: 'int', updates_per_trial: 'Optional[int]' = None, deadline_seconds: 'Optional[float]' = None) -> 'None'
trials — the study's total (every state counts); updates_per_trial — committed optimizer updates per trial; deadline_seconds — wall time from the search's start.
nnx.search.TrialOutcome¶
class nnx.search.TrialOutcome(*, number: 'int', params: 'Mapping[str, Any]', state: 'str', reason: 'str', monitor_value: 'Optional[float]' = None, epochs: 'int' = 0, updates: 'int' = 0, seconds: 'float' = 0.0, overrun_seconds: 'Optional[float]' = None, observations: 'tuple[tuple[int, float], ...]' = (), attempt_id: 'Optional[str]' = None, run_id: 'Optional[str]' = None, checkpoint: 'Optional[str]' = None, error: 'Optional[str]' = None) -> 'None'
One trial: its Optuna number, params, state ("completed", "pruned", "failed" or "running" for an orphan) and terminal reason; the resources it used (epochs, committed updates, seconds, overrun_seconds past the deadline); its best finite monitor_value, the validation observations (epoch, value) pruning saw; and its attempt_id, run_id and BEST checkpoint path.
nnx.search.SearchResult¶
class nnx.search.SearchResult(*, study_name: 'str', outcomes: 'tuple[TrialOutcome, ...]', best: 'Optional[TrialOutcome]', asked: 'int', stop_reason: 'str', direction: 'str', monitor: 'str', deadline_overrun_seconds: 'Optional[float]' = None, trial_runs: 'Mapping[int, Any]' = <factory>) -> 'None'
Every trial of the study (priors included), the best completed one, how many this call asked, and why it stopped ("budget", "deadline" or "exhausted" — a sampler ran out of points).
2.17. Experiment provenance (nnx.provenance)¶
nnx.provenance.ExperimentManifest¶
class nnx.provenance.ExperimentManifest(task: 'Optional[Mapping[str, Any]]' = None, labels: 'Optional[Sequence[str]]' = None, model: 'Optional[Mapping[str, Any]]' = None, data: 'Mapping[str, Any]' = <factory>, splits: 'Mapping[str, Any]' = <factory>, objective: 'Optional[Mapping[str, Any]]' = None, config: 'Mapping[str, Any]' = <factory>) -> 'None'
The declared intent of an experiment. Everything is JSON-like and immutable; fingerprint is the SHA-256 of its canonical bytes (with FORMAT), so changing the task, the label order, a split's digest or the objective's version changes it.
Details
Values are validated and copied into immutable form at construction, so
later changes to the dicts you passed never change the manifest;
equality and hashing follow the canonical bytes (``1`` and ``1.0``
differ, as they do in the fingerprint).
Args:
task: the task declaration (e.g. ``TaskSpec.state()``).
labels: the ordered output labels.
model: the model descriptor (e.g. ``NNModelParams.state()``).
data: named data identities (:class:`IdentityRef`, a declared id
string, or ``None`` for unknown).
splits: named split identities, likewise; a
:class:`nnx.data_splits.SplitManifest` is recorded as its digest
(:meth:`~nnx.data_splits.SplitManifest.identity`).
objective: the objective's identity, e.g. ``{"id": "kd", "version": 1}``.
config: any other JSON-like configuration.
nnx.provenance.ExperimentManifest.state¶
nnx.provenance.ExperimentManifest.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.provenance.ExperimentManifest.from_state¶
nnx.provenance.ExperimentManifest.from_state(state: 'Mapping[str, Any]') -> 'ExperimentManifest'
No public description is currently available.
nnx.provenance.ExperimentManifest.document¶
nnx.provenance.ExperimentManifest.document(self) -> 'dict[str, Any]'
The hashed document: the format version plus the manifest.
nnx.provenance.ExperimentManifest.canonical_bytes¶
nnx.provenance.ExperimentManifest.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.provenance.ExperimentManifest.fingerprint¶
nnx.provenance.ExperimentManifest.fingerprint(self) -> 'str'
sha256:<hex> of the canonical bytes.
nnx.provenance.ExperimentManifest.for_model¶
nnx.provenance.ExperimentManifest.for_model(model: 'Any', *, train: 'Any' = None, data: 'Optional[Mapping[str, Any]]' = None, splits: 'Optional[Mapping[str, Any]]' = None, objective: 'Optional[Mapping[str, Any]]' = None, config: 'Optional[Mapping[str, Any]]' = None, preprocessing: 'Any' = None) -> 'ExperimentManifest'
A manifest from an existing model's declarations: its task and label order, its model descriptor and built-in net params, its recorded transformation recipe when it has one, and — when given — the training configuration (NNTrainParams.state() without n_epochs, resume_epochs and the resume lineage, which describe an attempt; loaders are never read) and a fitted nnx.preprocessing.Standardizer, recorded as config["preprocessing"]: its schema, frozen statistics and fit-membership identity, never recomputed. Builds no model and iterates no loader.
nnx.provenance.IdentityRef¶
class nnx.provenance.IdentityRef(kind: 'str', value: 'Optional[str]' = None) -> 'None'
How a data or split identity is known.
Details
- ``declared`` — a name you supplied (a dataset id, a split name): it is
recorded and compared, but **never reported as verified**;
- ``digest`` — a content digest NNx computed (:func:`hash_file`,
:func:`hash_bytes`): equal digests are verified equality;
- ``unknown`` — not known.
nnx.provenance.IdentityRef.declared¶
nnx.provenance.IdentityRef.declared(value: 'str') -> 'IdentityRef'
No public description is currently available.
nnx.provenance.IdentityRef.unknown¶
nnx.provenance.IdentityRef.unknown() -> 'IdentityRef'
No public description is currently available.
nnx.provenance.IdentityRef.verified¶
property nnx.provenance.IdentityRef.verified
Whether this identity is a computed content digest.
nnx.provenance.IdentityRef.state¶
nnx.provenance.IdentityRef.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.provenance.IdentityRef.from_state¶
nnx.provenance.IdentityRef.from_state(state: 'Mapping[str, Any]') -> 'IdentityRef'
No public description is currently available.
nnx.provenance.hash_file¶
nnx.provenance.hash_file(path: 'Union[str, os.PathLike[str]]', *, chunk_size: 'int' = 1048576) -> 'IdentityRef'
A verified sha256:<hex> identity of a file's bytes — the one call here that reads data, and only when you make it.
nnx.provenance.hash_bytes¶
nnx.provenance.hash_bytes(data: 'bytes') -> 'IdentityRef'
A verified sha256:<hex> identity of data.
nnx.provenance.canonical_bytes¶
nnx.provenance.canonical_bytes(value: 'Any') -> 'bytes'
Canonical UTF-8 JSON of a JSON-like value: sorted keys, arrays in order, compact separators, Unicode kept as UTF-8. Raises TypeError / ValueError (naming the path) for anything else.
nnx.provenance.Attempt¶
class nnx.provenance.Attempt(attempt_id: 'str', fingerprint: 'str', run_id: 'str', status: 'str', started_at: 'str', finished_at: 'Optional[str]' = None, parent: 'Optional[Mapping[str, Any]]' = None, last_committed: 'Optional[Mapping[str, Any]]' = None, error: 'Optional[Mapping[str, Any]]' = None, compile: 'Optional[Mapping[str, Any]]' = None) -> 'None'
One execution of a plan: a fresh attempt_id per train() call, the plan's fingerprint and run_id, its status (running / completed / failed / cancelled), the parent it resumed from (run, attempt, checkpoint tag and generation), the last_committed checkpoint (tag, epoch, generation) and, when it did not complete, the error (type and message). compile is the fit's compile record — requested and effective — when train(compile=...) was used; it is absent from state() otherwise.
nnx.provenance.Attempt.state¶
nnx.provenance.Attempt.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.provenance.Attempt.from_state¶
nnx.provenance.Attempt.from_state(state: 'Mapping[str, Any]') -> 'Attempt'
No public description is currently available.
nnx.provenance.ProvenanceRecord¶
class nnx.provenance.ProvenanceRecord(manifest: 'ExperimentManifest', fingerprint: 'str', attempt: 'Optional[Attempt]' = None) -> 'None'
A run's provenance: the declared manifest, its fingerprint and the attempt that produced the run (NNRun.provenance).
nnx.provenance.load_provenance¶
nnx.provenance.load_provenance(run_id: 'str', root: 'Optional[str]' = None) -> 'Optional[ProvenanceRecord]'
Read a run's provenance files; None when it has none (absent). Raises ValueError for an unsupported format, a fingerprint that does not match its manifest, or JSON nested too deeply to read; a malformed entry (a missing key, a non-object entry) raises KeyError or TypeError.
nnx.provenance.compare¶
nnx.provenance.compare(left: 'Any', right: 'Any', *, root: 'Optional[str]' = None) -> 'ProvenanceComparison'
Compare two plans field by field — manifests, provenance records, NNRuns or run ids (their provenance files are read; no model is loaded, no loader iterated, no file hashed).
nnx.provenance.ProvenanceComparison¶
class nnx.provenance.ProvenanceComparison(fields: 'tuple[FieldComparison, ...]') -> 'None'
Every compared field. verified_equal is true only when both sides have provenance and every field is equal or verified — declared or unknown identities and absent provenance never count as verified equality.
nnx.provenance.ProvenanceComparison.differences¶
property nnx.provenance.ProvenanceComparison.differences
No public description is currently available.
nnx.provenance.ProvenanceComparison.verified_equal¶
property nnx.provenance.ProvenanceComparison.verified_equal
No public description is currently available.
nnx.provenance.ProvenanceComparison.by_path¶
nnx.provenance.ProvenanceComparison.by_path(self) -> 'dict[str, FieldComparison]'
No public description is currently available.
nnx.provenance.FieldComparison¶
class nnx.provenance.FieldComparison(path: 'str', status: 'str', left: 'Any' = None, right: 'Any' = None) -> 'None'
One field of a comparison: a stable dotted path and a status — equal, different, missing (on one side only), unknown (an unknown identity), declared (equal declared identities: unverified), verified (equal digests) or absent (no provenance at all).
2.18. Reproducible splits (nnx.data_splits)¶
nnx.data_splits.plan_split¶
nnx.data_splits.plan_split(ids: 'Optional[Iterable[Any]]' = None, *, strategy: 'str', groups: 'Optional[Iterable[Any]]' = None, times: 'Optional[Iterable[Any]]' = None, labels: 'Optional[Iterable[Any]]' = None, proportions: 'Optional[Iterable[float]]' = None, cutoffs: 'Optional[Iterable[Any]]' = None, gap: 'Any' = None, seed: 'Optional[int]' = None, insufficient: 'Optional[str]' = None, source: 'Union[IdentityRef, str, None]' = None) -> 'SplitManifest'
Plan a group, chronological or stratified split.
Details
Args:
ids: stable sample ids (str or int), one per row. ``None`` plans on
row positions, which needs ``source`` (row positions alone cannot
detect a reordered source) and replays only against it.
strategy: ``"group"``, ``"chronological"`` or ``"stratified"``.
groups: per-row group keys (``"group"``).
times: per-row numbers or dates / datetimes (``"chronological"``).
labels: per-row class labels (``"stratified"``).
proportions: ``(train, validation, test)`` shares summing to 1, train
> 0; a zero share leaves that split empty (``"group"``,
``"stratified"``).
cutoffs: ``(validation_start, test_start)``; either may be ``None``
(``"chronological"``).
gap: rows within ``gap`` before a cutoff are excluded; a number, or a
``timedelta`` for dates (``"chronological"``, default 0).
seed: orders groups / ids (``"group"``, ``"stratified"``); ``None``
draws one from the OS and records it. Chronological plans take
no seed.
insufficient: ``"raise"`` (default) or ``"train"`` for a class with
fewer rows than active splits (``"stratified"``).
source: the source's identity — an ``IdentityRef`` (e.g.
``nnx.provenance.hash_file``), a declared id string, or ``None``.
Raises:
SplitError: duplicate ids, too few groups or class rows, or no
training rows. ``TypeError`` / ``ValueError`` for malformed
arguments, including ones the strategy does not use.
nnx.data_splits.SplitManifest¶
class nnx.data_splits.SplitManifest(strategy: 'str', train: 'tuple[SampleId, ...]', validation: 'tuple[SampleId, ...]', test: 'tuple[SampleId, ...]', excluded: 'tuple[SampleId, ...]' = (), ids: 'str' = 'sample', source: 'IdentityRef' = <factory>, parameters: 'Mapping[str, Any]' = <factory>, seed: 'Optional[int]' = None) -> 'None'
Who is in which split, and how that was decided.
Details
Memberships are sample ids (``ids="sample"``) or row positions
(``ids="position"``, which can be replayed only against a recorded source
identity). They are disjoint, stored in sorted order, and ``excluded``
holds the rows a chronological ``gap`` left out. Equality and hashing
follow :meth:`canonical_bytes`, computed once (the manifest is frozen).
nnx.data_splits.SplitManifest.n_rows¶
property nnx.data_splits.SplitManifest.n_rows
Rows the plan covers, excluded ones included.
nnx.data_splits.SplitManifest.state¶
nnx.data_splits.SplitManifest.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.data_splits.SplitManifest.from_state¶
nnx.data_splits.SplitManifest.from_state(state: 'Mapping[str, Any]') -> 'SplitManifest'
No public description is currently available.
nnx.data_splits.SplitManifest.canonical_bytes¶
nnx.data_splits.SplitManifest.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.data_splits.SplitManifest.digest¶
nnx.data_splits.SplitManifest.digest(self) -> 'str'
sha256:<hex> of the canonical manifest: the split identity.
nnx.data_splits.SplitManifest.identity¶
nnx.data_splits.SplitManifest.identity(self) -> 'IdentityRef'
The split identity as a verified provenance reference.
nnx.data_splits.SplitManifest.to_json¶
nnx.data_splits.SplitManifest.to_json(self) -> 'str'
No public description is currently available.
nnx.data_splits.SplitManifest.from_json¶
nnx.data_splits.SplitManifest.from_json(text: 'str') -> 'SplitManifest'
No public description is currently available.
nnx.data_splits.SplitManifest.save¶
nnx.data_splits.SplitManifest.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write the manifest as JSON, atomically.
nnx.data_splits.SplitManifest.load¶
nnx.data_splits.SplitManifest.load(path: 'Union[str, os.PathLike[str]]') -> 'SplitManifest'
No public description is currently available.
nnx.data_splits.SplitManifest.resolve¶
nnx.data_splits.SplitManifest.resolve(self, ids: 'Iterable[Any]', *, source: 'Union[IdentityRef, str, None]' = None) -> 'SplitIndices'
Map this manifest onto a source whose rows carry ids.
Details
``ids`` are the source's sample ids in row order (``range(n)`` for a
positional manifest). When the manifest records a source identity,
``source`` must match it; passing one to a manifest that records none
raises too, since nothing could be checked. Raises
:class:`SplitError` on a changed identity, a positional-only
manifest, ids of another type, and duplicate, missing or unexpected
ids, and on an empty train split. Excluded sample ids may be absent
(they are in no split); a positional manifest needs every position,
since a dropped row would shift the rest.
nnx.data_splits.SplitIndices¶
class nnx.data_splits.SplitIndices(train: 'tuple[int, ...]', validation: 'tuple[int, ...]', test: 'tuple[int, ...]', excluded: 'tuple[int, ...]' = ()) -> 'None'
Row positions of each split in the source a manifest was resolved on, in the manifest's membership order.
nnx.data_splits.SplitError¶
class nnx.data_splits.SplitError
A split that cannot be planned or replayed as asked.
2.19. Train-fitted preprocessing (nnx.preprocessing)¶
nnx.preprocessing.Standardizer¶
class nnx.preprocessing.Standardizer(mean: 'tuple[float, ...]', scale: 'tuple[float, ...]', columns: 'Optional[tuple[ColumnLabel, ...]]' = None, dtype: 'str' = 'float32', fit_rows: 'int' = 0, fit_membership: 'IdentityRef' = <factory>) -> 'None'
Frozen per-column (x - mean) / scale with a declared schema.
Details
Build one with :meth:`fit`; reload one with :meth:`load` /
:meth:`from_state`. ``columns`` names the fitted DataFrame columns in
order (``str`` or ``int`` labels; ``None`` for unnamed arrays);
``dtype`` is the output dtype of :meth:`transform`; ``fit_rows`` and
``fit_membership`` identify the training rows it was fitted on. Calling
the standardizer on one sample (a 1-D row) standardizes that row, so it
also works as a :class:`SplitView` transform. Equality and hashing
follow :meth:`canonical_bytes`.
nnx.preprocessing.Standardizer.fit¶
nnx.preprocessing.Standardizer.fit(source: 'Any', *, rows: 'Iterable[Any]', columns: 'Optional[Sequence[ColumnLabel]]' = None, dtype: 'Union[str, torch.dtype]' = 'float32', membership: 'Any' = None) -> 'Standardizer'
Fit on source's training rows (integer positions) only.
Details
``source`` is a DataFrame (``columns`` selects and orders the schema,
default all columns; only the training rows of those columns are
copied), a 2-D tensor or array, or any row-indexable sequence. No
other row is read. The population mean and standard deviation are
computed in float64; a constant column (all training values equal)
gets its value as the mean and scale 1. ``membership`` identifies the
rows (an ``IdentityRef``, a ``SplitManifest``, a declared string); by
default it is a digest of the positions. Raises
:class:`PreprocessingError` for an empty, repeated or out-of-range
membership and for non-numeric or non-finite training values.
nnx.preprocessing.Standardizer.n_features¶
property nnx.preprocessing.Standardizer.n_features
No public description is currently available.
nnx.preprocessing.Standardizer.transform¶
nnx.preprocessing.Standardizer.transform(self, data: 'Any') -> 'torch.Tensor'
Standardize data with the frozen statistics.
Details
A DataFrame must hold the fitted ``columns`` once each and in the
fitted order (other columns are ignored); a tensor, array or row
sequence must be 2-D with ``n_features`` columns. Values must be
numeric and finite, and so must the result in ``dtype``. Returns a
new tensor; the input is never modified. Raises
:class:`PreprocessingError` otherwise — before any model sees it.
nnx.preprocessing.Standardizer.state¶
nnx.preprocessing.Standardizer.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.preprocessing.Standardizer.from_state¶
nnx.preprocessing.Standardizer.from_state(state: 'Mapping[str, Any]') -> 'Standardizer'
Rebuild from state, rejecting anything malformed with PreprocessingError (bad lengths, non-finite values, non-positive scales, missing or mistyped entries).
nnx.preprocessing.Standardizer.canonical_bytes¶
nnx.preprocessing.Standardizer.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.preprocessing.Standardizer.digest¶
nnx.preprocessing.Standardizer.digest(self) -> 'str'
sha256:<hex> of the canonical state: the schema identity.
nnx.preprocessing.Standardizer.to_json¶
nnx.preprocessing.Standardizer.to_json(self) -> 'str'
No public description is currently available.
nnx.preprocessing.Standardizer.from_json¶
nnx.preprocessing.Standardizer.from_json(text: 'str') -> 'Standardizer'
No public description is currently available.
nnx.preprocessing.Standardizer.save¶
nnx.preprocessing.Standardizer.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write the statistics and schema as JSON, atomically.
nnx.preprocessing.Standardizer.load¶
nnx.preprocessing.Standardizer.load(path: 'Union[str, os.PathLike[str]]') -> 'Standardizer'
No public description is currently available.
nnx.preprocessing.SplitView¶
class nnx.preprocessing.SplitView(base: 'Any', indices: 'Iterable[int]', *, transform: 'Optional[Callable[..., Any]]' = None) -> 'None'
Rows indices of base, with transform applied to each sample's input.
Details
A sample is read from ``base`` (with whatever transform ``base`` already
applies) and only its first element is transformed; targets and any
other elements pass through unchanged, and tuples, named tuples and
lists keep their type. ``base`` itself — its ``transform``, labels and
order — is never modified, so a training view with a random
augmentation and an evaluation view with a fixed transform can share one
base and be read alternately or from DataLoader workers. ``indices``
are the view's sample ids in ``base``; ``dataset`` aliases ``base`` like
``torch.utils.data.Subset``, and ``classes`` forwards the base's class
names. Dict samples are rejected: which entry is the input is ambiguous.
nnx.preprocessing.SplitView.dataset¶
property nnx.preprocessing.SplitView.dataset
No public description is currently available.
nnx.preprocessing.SplitView.classes¶
property nnx.preprocessing.SplitView.classes
The base dataset's class names (dataset-wide, like its labels).
nnx.preprocessing.SplitView.reconstructible¶
property nnx.preprocessing.SplitView.reconstructible
Whether the transform can be rebuilt from state alone.
nnx.preprocessing.SplitView.state¶
nnx.preprocessing.SplitView.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.preprocessing.describe_transform¶
nnx.preprocessing.describe_transform(transform: 'Optional[Callable[..., Any]]') -> 'Optional[dict[str, Any]]'
How a view's transform can be rebuilt: a Standardizer from its state (by digest); any other callable is runtime-only and must be registered again — passed to the view anew — to reconstruct it.
nnx.preprocessing.PreprocessingError¶
class nnx.preprocessing.PreprocessingError
Preprocessing that cannot be fitted, loaded or applied as declared.
2.20. Fitted classifier calibration (nnx.calibration)¶
nnx.calibration.fit_temperature¶
nnx.calibration.fit_temperature(logits: 'Any', targets: 'Any', *, model_id: 'str', split_id: 'str', labels: 'Optional[Sequence[str]]' = None, train_split_id: 'Optional[str]' = None, test_split_id: 'Optional[str]' = None, min_temperature: 'float' = 0.001, max_temperature: 'float' = 1000.0, tolerance: 'float' = 1e-10, max_iterations: 'int' = 200) -> 'CalibrationFit'
Fit a scalar temperature on the calibration split split_id.
Details
``logits`` are categorical ``(N, C)`` raw scores (array, tensor, or a
categorical ``nnx.prediction.PredictionResult``, whose spec supplies the
``labels``) and ``targets`` their integer class indices. The NLL of
``softmax(logits / T)`` is minimized over ``T`` in
``[min_temperature, max_temperature]`` in float64 on a copy of the
logits, by bisection on its (monotone) slope in ``1 / T`` until the
log-temperature bracket is narrower than ``tolerance``.
Raises :class:`CalibrationError` for non-finite logits, ``C < 2``,
out-of-range targets, labels that do not name each class once, and a
``split_id`` equal to ``train_split_id`` or ``test_split_id``. Returns a
failed :class:`CalibrationFit` — never a calibrator with a NaN or
boundary temperature — when the optimum lies outside the search range (a
separable split, logits no better than uniform), the NLL does not depend
on ``T`` (constant logits), or the search does not converge.
nnx.calibration.CalibrationFit¶
class nnx.calibration.CalibrationFit(status: 'str', calibrator: 'Optional[TemperatureCalibrator]', reason: 'Optional[str]', config: 'Mapping[str, Any]', split: 'Mapping[str, Any]') -> 'None'
The outcome of fit_temperature: status == "ok" with a calibrator, or "failed" with calibrator=None and a reason. config and split are recorded either way (read-only). Compared and hashed by identity; compare the calibrators themselves.
nnx.calibration.CalibrationFit.ok¶
property nnx.calibration.CalibrationFit.ok
No public description is currently available.
nnx.calibration.CalibrationFit.require¶
nnx.calibration.CalibrationFit.require(self) -> 'TemperatureCalibrator'
The calibrator, or CalibrationFitError with the reason.
nnx.calibration.TemperatureCalibrator¶
class nnx.calibration.TemperatureCalibrator(temperature: 'float', labels: 'tuple[str, ...]', model_id: 'str', split_id: 'str', train_split_id: 'Optional[str]' = None, test_split_id: 'Optional[str]' = None, fit_config: 'Mapping[str, Any]' = <factory>, fit_result: 'Mapping[str, Any]' = <factory>) -> 'None'
A fitted scalar temperature bound to its label schema and model.
Details
Build one with :func:`fit_temperature`; reload one with :meth:`load` /
:meth:`from_state`. ``labels`` are the ordered class names the logits'
columns must carry, ``model_id`` the model they must come from,
``split_id`` the calibration split it was fitted on (and, when declared,
``train_split_id`` / ``test_split_id``, which must differ from it).
``fit_config`` / ``fit_result`` record how it was fitted (the result's
``nll_before`` / ``nll_after`` are the calibration split's exact,
unfloored NLL, computed in the log domain). Equality,
hashing and :attr:`id` follow the canonical state.
nnx.calibration.TemperatureCalibrator.split¶
property nnx.calibration.TemperatureCalibrator.split
The split ids, and that their disjointness is unverified.
nnx.calibration.TemperatureCalibrator.id¶
property nnx.calibration.TemperatureCalibrator.id
The calibrator id: digest.
nnx.calibration.TemperatureCalibrator.transform¶
nnx.calibration.TemperatureCalibrator.transform(self, logits: 'Any', *, model_id: 'str', labels: 'Optional[Sequence[str]]' = None, override: 'Optional[str]' = None, sample_ids: 'Any' = None) -> 'CalibratedPrediction'
Calibrate logits — an (N, C) array or tensor, or a categorical nnx.prediction.PredictionResult.
Details
``sample_ids`` names the rows of an array (``0..N-1`` by default); a
prediction keeps its own ids, and ``sample_ids=`` must then equal
them.
``model_id`` and the column ``labels`` (taken from the prediction's
spec when omitted) must equal the fitted ones; otherwise
:class:`CalibrationMismatchError` is raised before the logits are
read, unless ``override`` names the exception — then the result
records the name and each mismatch. An override exists only to accept
a mismatch: with none, the name is validated and nothing is recorded.
Non-finite logits or a width that differs from the labels raise
:class:`CalibrationError`. The input is never modified.
nnx.calibration.TemperatureCalibrator.report¶
nnx.calibration.TemperatureCalibrator.report(self, logits: 'Any', targets: 'Any', *, model_id: 'str', split_id: 'str', labels: 'Optional[Sequence[str]]' = None, override: 'Optional[str]' = None, n_bins: 'int' = 10, epsilon: 'Optional[float]' = 1e-12) -> 'CalibrationReport'
Held-out NLL, Brier, ECE and reliability bins before and after calibration.
Details
``split_id`` names the held-out split and must differ from the
calibration and training split ids. Identity checks match
:meth:`transform`. The report's :attr:`CalibrationReport.outcome` is
``"improved"`` only when neither NLL nor Brier got worse and one got
better; a worse result is reported as ``"worsened"`` / ``"mixed"``.
nnx.calibration.TemperatureCalibrator.state¶
nnx.calibration.TemperatureCalibrator.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.calibration.TemperatureCalibrator.from_state¶
nnx.calibration.TemperatureCalibrator.from_state(state: 'Mapping[str, Any]') -> 'TemperatureCalibrator'
Rebuild from state, rejecting anything malformed with CalibrationError.
nnx.calibration.TemperatureCalibrator.digest¶
nnx.calibration.TemperatureCalibrator.digest(self) -> 'str'
sha256:<hex> of the canonical state.
nnx.calibration.TemperatureCalibrator.from_json¶
nnx.calibration.TemperatureCalibrator.from_json(text: 'str') -> 'TemperatureCalibrator'
No public description is currently available.
nnx.calibration.TemperatureCalibrator.load¶
nnx.calibration.TemperatureCalibrator.load(path: 'Union[str, os.PathLike[str]]') -> 'TemperatureCalibrator'
No public description is currently available.
nnx.calibration.TemperatureCalibrator.canonical_bytes¶
nnx.calibration.TemperatureCalibrator.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.calibration.TemperatureCalibrator.to_json¶
nnx.calibration.TemperatureCalibrator.to_json(self) -> 'str'
No public description is currently available.
nnx.calibration.TemperatureCalibrator.save¶
nnx.calibration.TemperatureCalibrator.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write to_json to path atomically.
nnx.calibration.CalibratedPrediction¶
class nnx.calibration.CalibratedPrediction(logits: 'np.ndarray', probabilities: 'np.ndarray', calibrated_probabilities: 'np.ndarray', decoded: 'np.ndarray', labels: 'tuple[str, ...]', sample_ids: 'np.ndarray', calibrator_id: 'str', model_id: 'str', temperature: 'float', override: 'Optional[Mapping[str, Any]]' = None) -> 'None'
Raw and calibrated views of one categorical prediction, kept apart.
Details
Attributes:
logits: the raw logits as given (a copy; tensors become arrays).
probabilities: uncalibrated float64 ``softmax(logits)``.
calibrated_probabilities: float64 ``softmax(logits / temperature)``.
decoded: argmax class indices of the raw logits — calibration never
changes the argmax.
labels: the ordered class names of the columns.
sample_ids: ``int64[N]`` row identities (a ``PredictionResult``'s
own ids, else ``transform``'s ``sample_ids=``, else ``0..N-1``).
calibrator_id: :meth:`TemperatureCalibrator.digest` of the calibrator.
model_id: the model id the logits were declared to come from.
temperature: the calibrator's temperature.
override: ``None``, or the named override and the label / model-id
mismatches it accepted.
nnx.calibration.CalibratedPrediction.decoded_labels¶
nnx.calibration.CalibratedPrediction.decoded_labels(self) -> 'np.ndarray'
Decoded class names.
nnx.calibration.CalibrationReport¶
class nnx.calibration.CalibrationReport(calibrator_id: 'str', model_id: 'str', labels: 'tuple[str, ...]', temperature: 'float', split_id: 'str', calibration_split_id: 'str', n_samples: 'int', n_bins: 'int', epsilon: 'Optional[float]', before: 'CalibrationMetrics', after: 'CalibrationMetrics', outcome: 'str', override: 'Optional[Mapping[str, Any]]' = None, train_split_id: 'Optional[str]' = None, test_split_id: 'Optional[str]' = None) -> 'None'
Held-out metrics before and after calibration.
Details
``outcome`` is ``"improved"`` (neither NLL nor Brier worse, one better),
``"worsened"`` (neither better, one worse), ``"unchanged"`` or
``"mixed"``; :attr:`improved` is true only for ``"improved"``.
``override`` is the named override that accepted a label / model-id
mismatch, or ``None``. Equality follows the canonical state, so a
reloaded calibrator reproduces an equal report.
nnx.calibration.CalibrationReport.improved¶
property nnx.calibration.CalibrationReport.improved
No public description is currently available.
nnx.calibration.CalibrationReport.summary¶
nnx.calibration.CalibrationReport.summary(self) -> 'str'
One line that states the outcome — never success when the calibrated probabilities are worse.
nnx.calibration.CalibrationReport.state¶
nnx.calibration.CalibrationReport.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.calibration.CalibrationReport.from_state¶
nnx.calibration.CalibrationReport.from_state(state: 'Mapping[str, Any]') -> 'CalibrationReport'
No public description is currently available.
nnx.calibration.CalibrationReport.from_json¶
nnx.calibration.CalibrationReport.from_json(text: 'str') -> 'CalibrationReport'
No public description is currently available.
nnx.calibration.CalibrationReport.load¶
nnx.calibration.CalibrationReport.load(path: 'Union[str, os.PathLike[str]]') -> 'CalibrationReport'
No public description is currently available.
nnx.calibration.CalibrationReport.canonical_bytes¶
nnx.calibration.CalibrationReport.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.calibration.CalibrationReport.to_json¶
nnx.calibration.CalibrationReport.to_json(self) -> 'str'
No public description is currently available.
nnx.calibration.CalibrationReport.save¶
nnx.calibration.CalibrationReport.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write to_json to path atomically.
nnx.calibration.CalibrationMetrics¶
class nnx.calibration.CalibrationMetrics(nll: 'float', brier: 'float', ece: 'Optional[float]', bins: 'tuple[ReliabilityBin, ...]') -> 'None'
NLL, Brier, expected calibration error and reliability bins of one set of probabilities.
nnx.calibration.CalibrationMetrics.of¶
nnx.calibration.CalibrationMetrics.of(probabilities: 'Any', targets: 'Any', *, n_bins: 'int' = 10, epsilon: 'Optional[float]' = 1e-12) -> 'CalibrationMetrics'
All four metrics of probabilities, validated once.
nnx.calibration.CalibrationMetrics.state¶
nnx.calibration.CalibrationMetrics.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.calibration.CalibrationMetrics.from_state¶
nnx.calibration.CalibrationMetrics.from_state(state: 'Mapping[str, Any]') -> 'CalibrationMetrics'
Rebuild from state; malformed entries raise CalibrationError.
nnx.calibration.ReliabilityBin¶
class nnx.calibration.ReliabilityBin(lower: 'float', upper: 'float', count: 'int', confidence: 'Optional[float]', accuracy: 'Optional[float]') -> 'None'
One confidence bin: [lower, upper) (the last bin is closed at 1), the rows whose top-1 confidence falls in it, their mean confidence and top-1 accuracy. An empty bin has count == 0 and None means.
nnx.calibration.ReliabilityBin.state¶
nnx.calibration.ReliabilityBin.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.calibration.ReliabilityBin.from_state¶
nnx.calibration.ReliabilityBin.from_state(state: 'Mapping[str, Any]') -> 'ReliabilityBin'
Rebuild from state; malformed entries raise CalibrationError.
nnx.calibration.negative_log_likelihood¶
nnx.calibration.negative_log_likelihood(probabilities: 'Any', targets: 'Any', *, epsilon: 'Optional[float]' = 1e-12) -> 'float'
Mean -log p[target] over the rows of categorical (N, C) probabilities.
Details
``epsilon`` floors the true-class probability (default
:data:`DEFAULT_EPSILON`, the named ``nll`` metric's floor); ``epsilon=None``
is the exact value, which is ``+inf`` when any true-class probability is 0.
nnx.calibration.brier_score¶
nnx.calibration.brier_score(probabilities: 'Any', targets: 'Any') -> 'float'
Mean over rows of the class-summed squared error sum_c (p_c - 1[c == target])^2 (range [0, 2]).
nnx.calibration.reliability_bins¶
nnx.calibration.reliability_bins(probabilities: 'Any', targets: 'Any', *, n_bins: 'int' = 10) -> 'tuple[ReliabilityBin, ...]'
Equal-width top-1 reliability bins over [0, 1].
Details
Bin ``i`` covers ``[i / n_bins, (i + 1) / n_bins)``; the last bin is
closed at 1, and a confidence exactly on an inner edge belongs to the
upper bin. The prediction is the first class of maximal probability, and
empty bins are kept with ``count == 0``.
nnx.calibration.expected_calibration_error¶
nnx.calibration.expected_calibration_error(bins: 'Iterable[ReliabilityBin]') -> 'Optional[float]'
Count-weighted mean |accuracy - confidence| over the bins; None when the bins hold no rows.
nnx.calibration.model_fingerprint¶
nnx.calibration.model_fingerprint(model: 'Any') -> 'str'
sha256:<hex> of a model's weights: every state_dict() entry's name, dtype, shape and little-endian bytes, in name order (the same on every host).
Details
``model`` is an ``nn.Module`` or anything with a ``.net`` module (an
``NNModel``). Use it as a ``model_id`` so a calibrator refuses logits
from a model whose weights changed; a declared id such as
``f"{run.id}:BEST"`` works too, but is only as reliable as its naming.
nnx.calibration.CalibrationError¶
class nnx.calibration.CalibrationError
Calibration inputs, configuration or a serialized state that cannot be used as declared.
nnx.calibration.CalibrationFitError¶
class nnx.calibration.CalibrationFitError
CalibrationFit.require on a failed fit.
nnx.calibration.CalibrationMismatchError¶
class nnx.calibration.CalibrationMismatchError
Logits whose label schema or model id differ from the calibrator's, without a named override.
2.21. Abstention and risk-coverage (nnx.abstention)¶
nnx.abstention.AbstentionPolicy¶
class nnx.abstention.AbstentionPolicy(kind: 'str', threshold: 'float', labels: 'tuple[str, ...]', model_id: 'str', tuning_split_id: 'str', input_field: 'str' = 'probabilities', calibrator_id: 'Optional[str]' = None) -> 'None'
Accept a row when its score reaches threshold; abstain otherwise.
Details
Args:
kind: ``"max_probability"`` or ``"margin"``.
threshold: a finite number in ``[0, 1]``; a row is accepted when its
score is at or above it.
labels: the ordered class names the probabilities' columns carry.
model_id: the model the probabilities must come from.
tuning_split_id: the validation split the threshold was chosen on.
input_field: ``"probabilities"`` (raw) or
``"calibrated_probabilities"``.
calibrator_id: the calibrator of a calibrated field (required
then, and ``None`` for raw probabilities).
Equality, hashing and :attr:`id` follow the canonical state.
nnx.abstention.AbstentionPolicy.id¶
property nnx.abstention.AbstentionPolicy.id
sha256:<hex> of the canonical state.
nnx.abstention.AbstentionPolicy.reason¶
property nnx.abstention.AbstentionPolicy.reason
The reason an abstained row carries.
nnx.abstention.AbstentionPolicy.apply¶
nnx.abstention.AbstentionPolicy.apply(self, source: 'Any', *, labels: 'Optional[Sequence[str]]' = None, model_id: 'Optional[str]' = None, sample_ids: 'Any' = None, input_field: 'Optional[str]' = None, calibrator_id: 'Optional[str]' = None) -> 'AbstentionResult'
Accept or abstain on every row of source.
Details
``source`` is a ``CalibratedPrediction`` (the policy's
``input_field`` picks its raw or calibrated view), a categorical
``PredictionResult`` (raw; pass ``model_id=``), or ``(N, C)``
probabilities with ``labels=`` / ``model_id=`` (and
``input_field=`` / ``calibrator_id=`` for calibrated ones). The
declared schema must equal the policy's — otherwise
:class:`AbstentionSchemaError`, before any probability is read.
Malformed probabilities raise :class:`AbstentionError`. The input is
never modified.
nnx.abstention.AbstentionPolicy.state¶
nnx.abstention.AbstentionPolicy.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.abstention.AbstentionPolicy.from_state¶
nnx.abstention.AbstentionPolicy.from_state(state: 'Mapping[str, Any]') -> 'AbstentionPolicy'
Rebuild from state; anything malformed — including a non-deployable threshold such as a curve's accept-none endpoint — raises AbstentionError.
nnx.abstention.AbstentionPolicy.from_json¶
nnx.abstention.AbstentionPolicy.from_json(text: 'str') -> 'AbstentionPolicy'
No public description is currently available.
nnx.abstention.AbstentionPolicy.load¶
nnx.abstention.AbstentionPolicy.load(path: 'Any') -> 'AbstentionPolicy'
No public description is currently available.
nnx.abstention.AbstentionPolicy.canonical_bytes¶
nnx.abstention.AbstentionPolicy.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.abstention.AbstentionPolicy.to_json¶
nnx.abstention.AbstentionPolicy.to_json(self) -> 'str'
No public description is currently available.
nnx.abstention.AbstentionPolicy.save¶
nnx.abstention.AbstentionPolicy.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write to_json to path atomically.
nnx.abstention.AbstentionResult¶
class nnx.abstention.AbstentionResult(policy_id: 'str', kind: 'str', threshold: 'float', input_field: 'str', calibrator_id: 'Optional[str]', model_id: 'str', labels: 'tuple[str, ...]', sample_ids: 'np.ndarray', probabilities: 'np.ndarray', prediction: 'np.ndarray', score: 'np.ndarray', accepted: 'np.ndarray', calibration_override: 'Optional[Mapping[str, Any]]' = None) -> 'None'
A policy applied to N rows.
Details
``probabilities`` is a float64 copy of the field the policy read;
``prediction`` the first class of maximal probability (a prediction's
own ``decoded`` class, for a prediction); ``score`` the
policy score; ``accepted`` whether the row was accepted. The policy's
identity (``policy_id``, ``kind``, ``threshold``, ``input_field``,
``calibrator_id``, ``model_id``, ``labels``) travels with the result, and
so does ``calibration_override``: the named override under which a
calibrated prediction was produced, if any (``nnx.calibration``).
nnx.abstention.AbstentionResult.abstained¶
property nnx.abstention.AbstentionResult.abstained
No public description is currently available.
nnx.abstention.AbstentionResult.reasons¶
property nnx.abstention.AbstentionResult.reasons
No public description is currently available.
nnx.abstention.AbstentionResult.accepted_ids¶
property nnx.abstention.AbstentionResult.accepted_ids
No public description is currently available.
nnx.abstention.AbstentionResult.abstained_ids¶
property nnx.abstention.AbstentionResult.abstained_ids
No public description is currently available.
nnx.abstention.AbstentionResult.outcomes¶
nnx.abstention.AbstentionResult.outcomes(self) -> 'tuple[Outcome, ...]'
No public description is currently available.
nnx.abstention.AbstentionResult.accepted_rows¶
nnx.abstention.AbstentionResult.accepted_rows(self, targets: 'Any') -> 'AcceptedRows'
Sample ids, targets and predictions of the accepted rows only.
nnx.abstention.AbstentionResult.report¶
nnx.abstention.AbstentionResult.report(self, targets: 'Any') -> 'CoverageReport'
Coverage and selective risk against integer class targets.
nnx.abstention.Outcome¶
class nnx.abstention.Outcome(sample_id: 'Optional[int]', status: 'str', prediction: 'str', prediction_index: 'int', distribution: 'tuple[tuple[str, float], ...]', score: 'float', reason: 'Optional[str]') -> 'None'
One row: status "accepted" or "abstained" (with its reason); the prediction (label) and its index, the ordered distribution, the score and the sample_id are kept either way.
nnx.abstention.Outcome.accepted¶
property nnx.abstention.Outcome.accepted
No public description is currently available.
nnx.abstention.AcceptedRows¶
class nnx.abstention.AcceptedRows(sample_ids: 'np.ndarray', targets: 'np.ndarray', predictions: 'np.ndarray') -> 'None'
The accepted rows only — what a hard-label report (accuracy, a confusion matrix, VisUtils.classification_report) may receive. An abstention is never turned into a pseudo-class.
nnx.abstention.CoverageReport¶
class nnx.abstention.CoverageReport(policy_id: 'Optional[str]', total: 'int', accepted: 'int', incorrect: 'int', sample_ids: 'np.ndarray', accepted_ids: 'np.ndarray') -> 'None'
Coverage over every row and selective risk over the accepted rows.
Details
``coverage = accepted / total`` and ``risk = incorrect / accepted``.
With nothing accepted, ``risk`` is ``None`` and ``status`` is
``"no_accepted"`` — never a risk of 0; with no rows at all both are
``None`` and ``status`` is ``"empty"``. ``sample_ids`` keeps every
original id and ``accepted_ids`` the accepted ones, in row order, as
read-only int64 arrays.
nnx.abstention.CoverageReport.status¶
property nnx.abstention.CoverageReport.status
No public description is currently available.
nnx.abstention.CoverageReport.summary¶
nnx.abstention.CoverageReport.summary(self) -> 'str'
No public description is currently available.
nnx.abstention.CoverageReport.state¶
nnx.abstention.CoverageReport.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.abstention.CoverageReport.from_state¶
nnx.abstention.CoverageReport.from_state(state: 'Mapping[str, Any]') -> 'CoverageReport'
No public description is currently available.
nnx.abstention.CoverageReport.from_json¶
nnx.abstention.CoverageReport.from_json(text: 'str') -> 'CoverageReport'
No public description is currently available.
nnx.abstention.CoverageReport.load¶
nnx.abstention.CoverageReport.load(path: 'Any') -> 'CoverageReport'
No public description is currently available.
nnx.abstention.CoverageReport.coverage¶
property nnx.abstention.CoverageReport.coverage
No public description is currently available.
nnx.abstention.CoverageReport.risk¶
property nnx.abstention.CoverageReport.risk
No public description is currently available.
nnx.abstention.CoverageReport.canonical_bytes¶
nnx.abstention.CoverageReport.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.abstention.CoverageReport.to_json¶
nnx.abstention.CoverageReport.to_json(self) -> 'str'
No public description is currently available.
nnx.abstention.CoverageReport.save¶
nnx.abstention.CoverageReport.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write to_json to path atomically.
nnx.abstention.CoverageAccumulator¶
class nnx.abstention.CoverageAccumulator(*, allow_repeated_ids: 'bool' = False) -> 'None'
Coverage and selective risk aggregated over chunks.
Details
``update(result, targets)`` adds one chunk; :meth:`report` equals the
eager ``AbstentionResult.report`` over the concatenated rows — the same
accepted / incorrect / total counts and the same risk, including when
nothing is accepted. Every chunk must come from the same policy, and by
default every id must be new — within its chunk and across chunks —
since the default ``0..n-1`` repeats from chunk to chunk: pass each
chunk its own ``sample_ids=``. For a split whose ids legitimately repeat
(oversampled), pass ``allow_repeated_ids=True``; ids are then not
checked.
A ``PredictionResult`` or ``CalibratedPrediction`` chunk carries the ids
it was made with (``prediction_from_logits(sample_ids=...)``,
``TemperatureCalibrator.transform(sample_ids=...)``). Ids are kept as
int64 arrays — in row order, the accepted ones, and sorted for the repeat
check, which costs ``O(k log² n)`` for ``k`` new ids among ``n`` so far.
nnx.abstention.CoverageAccumulator.update¶
nnx.abstention.CoverageAccumulator.update(self, result: 'AbstentionResult', targets: 'Any') -> 'None'
No public description is currently available.
nnx.abstention.CoverageAccumulator.report¶
nnx.abstention.CoverageAccumulator.report(self) -> 'CoverageReport'
The aggregated report over every chunk added so far.
nnx.abstention.select_threshold¶
nnx.abstention.select_threshold(source: 'Any', targets: 'Any', *, kind: 'str', risk_ceiling: 'float', split_id: 'str', test_split_id: 'Optional[str]' = None, labels: 'Optional[Sequence[str]]' = None, model_id: 'Optional[str]' = None, input_field: 'Optional[str]' = None, calibrator_id: 'Optional[str]' = None, sample_ids: 'Any' = None, candidates: 'Optional[Iterable[float]]' = None) -> 'ThresholdSelection'
Choose a kind threshold on the validation split split_id.
Details
``source`` is declared as for :meth:`AbstentionPolicy.apply`; a
``CalibratedPrediction`` needs ``input_field=`` to say which of its two
views to tune on. Every candidate threshold is evaluated: by default
every distinct score, or, beyond :data:`MAX_DEFAULT_CANDIDATES` of them,
that many evenly spaced order statistics of the scores (pass
``candidates=`` for another set). The feasible ones accept at least one
row with ``risk <= risk_ceiling``; among them the highest coverage wins
(equal coverage is the same accepted set, so the same risk), then the
higher, more conservative threshold.
Nothing feasible gives ``status="no_feasible_policy"``. A ``split_id``
equal to ``test_split_id`` is refused: tuning on the test split would
leak it. The ceiling is empirical — met on this split, not guaranteed.
nnx.abstention.ThresholdSelection¶
class nnx.abstention.ThresholdSelection(status: 'str', policy: 'Optional[AbstentionPolicy]', points: 'tuple[CurvePoint, ...]', kind: 'str', risk_ceiling: 'float', split_id: 'str', labels: 'tuple[str, ...]', model_id: 'str', input_field: 'str' = 'probabilities', calibrator_id: 'Optional[str]' = None, test_split_id: 'Optional[str]' = None, basis: 'str' = 'empirical', calibration_override: 'Optional[Mapping[str, Any]]' = None) -> 'None'
The outcome of select_threshold.
Details
``status`` is ``"ok"`` with the chosen ``policy`` or
``"no_feasible_policy"`` with ``policy=None``. ``points`` lists every
candidate threshold with its counts, coverage and risk, then the
accept-none endpoint. ``basis`` is ``"empirical"``: the ceiling was met
on ``split_id``, which does not guarantee it anywhere else.
``labels``, ``model_id``, ``input_field`` and ``calibrator_id`` record
what was tuned on — a chosen policy carries the same — so even a
``"no_feasible_policy"`` selection says which model and view failed.
``calibration_override`` is the named override (``nnx.calibration``)
under which the calibrated probabilities tuned on were produced, if any.
nnx.abstention.ThresholdSelection.require¶
nnx.abstention.ThresholdSelection.require(self) -> 'AbstentionPolicy'
The chosen policy, or AbstentionError.
nnx.abstention.ThresholdSelection.state¶
nnx.abstention.ThresholdSelection.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.abstention.ThresholdSelection.from_state¶
nnx.abstention.ThresholdSelection.from_state(state: 'Mapping[str, Any]') -> 'ThresholdSelection'
No public description is currently available.
nnx.abstention.ThresholdSelection.from_json¶
nnx.abstention.ThresholdSelection.from_json(text: 'str') -> 'ThresholdSelection'
No public description is currently available.
nnx.abstention.ThresholdSelection.load¶
nnx.abstention.ThresholdSelection.load(path: 'Any') -> 'ThresholdSelection'
No public description is currently available.
nnx.abstention.ThresholdSelection.canonical_bytes¶
nnx.abstention.ThresholdSelection.canonical_bytes(self) -> 'bytes'
No public description is currently available.
nnx.abstention.ThresholdSelection.to_json¶
nnx.abstention.ThresholdSelection.to_json(self) -> 'str'
No public description is currently available.
nnx.abstention.ThresholdSelection.save¶
nnx.abstention.ThresholdSelection.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
Write to_json to path atomically.
nnx.abstention.risk_coverage_curve¶
nnx.abstention.risk_coverage_curve(source: 'Any', targets: 'Any', *, kind: 'str', candidates: 'Optional[Iterable[float]]' = None, input_field: 'Optional[str]' = None) -> 'tuple[CurvePoint, ...]'
Coverage and selective risk at each candidate threshold, ascending.
Details
``source`` is ``(N, C)`` probabilities (the prediction is each row's
first class of maximal probability), a categorical ``PredictionResult``
or a ``CalibratedPrediction`` with ``input_field=`` (their own
``decoded`` classes, as :func:`select_threshold` reads them). By default
the candidates are the distinct scores (bounded by
:data:`MAX_DEFAULT_CANDIDATES`). Coverage never rises as the threshold
rises; risk is measured at each point, never assumed monotone. The last
point is the accept-none endpoint.
nnx.abstention.CurvePoint¶
class nnx.abstention.CurvePoint(threshold: 'Optional[float]', accepted: 'int', incorrect: 'int', total: 'int') -> 'None'
One point of a risk-coverage curve: rows with a score at or above threshold are accepted. The final "accept_none" point (threshold=None, coverage 0, risk unavailable) exists only on the curve and is not deployable.
nnx.abstention.CurvePoint.status¶
property nnx.abstention.CurvePoint.status
No public description is currently available.
nnx.abstention.CurvePoint.deployable¶
property nnx.abstention.CurvePoint.deployable
No public description is currently available.
nnx.abstention.CurvePoint.state¶
nnx.abstention.CurvePoint.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.abstention.CurvePoint.coverage¶
property nnx.abstention.CurvePoint.coverage
No public description is currently available.
nnx.abstention.CurvePoint.risk¶
property nnx.abstention.CurvePoint.risk
No public description is currently available.
nnx.abstention.scores¶
nnx.abstention.scores(probabilities: 'Any', *, kind: 'str') -> 'tuple[np.ndarray, np.ndarray]'
(prediction, score) per row of categorical (N, C) probabilities: the first class of maximal probability and the policy score (the top probability, or the top-two margin). A prediction's own decoded class is kept by AbstentionPolicy.apply and select_threshold instead.
nnx.abstention.decide¶
nnx.abstention.decide(result: 'Any', policy: 'AbstentionPolicy', *, model_id: 'str') -> 'SelectiveDecision'
Apply policy to a typed decision (nnx.decisions).
Details
``result`` is a ``ChoiceResult`` or ``ScoreResult``; its option ids, in
order, are the labels. A provider's distribution is raw, so the policy
must read ``"probabilities"``; labels and ``model_id`` must equal the
policy's (:class:`AbstentionSchemaError` otherwise). The result is
returned unchanged inside a :class:`SelectiveDecision`, accepted or not.
nnx.abstention.SelectiveDecision¶
class nnx.abstention.SelectiveDecision(result: 'Any', outcome: 'Outcome', policy_id: 'str', input_field: 'str', calibrator_id: 'Optional[str]') -> 'None'
A typed decision with a policy's verdict: result is the provider's ChoiceResult / ScoreResult, unchanged — abstaining is a success, not an error — and outcome says whether it was accepted. input_field / calibrator_id record which probabilities the verdict read (a provider's distribution is raw).
nnx.abstention.SelectiveDecision.accepted¶
property nnx.abstention.SelectiveDecision.accepted
No public description is currently available.
nnx.abstention.SelectiveDecision.abstained¶
property nnx.abstention.SelectiveDecision.abstained
No public description is currently available.
nnx.abstention.AbstentionError¶
class nnx.abstention.AbstentionError
Malformed probabilities, targets, configuration or serialized state.
nnx.abstention.AbstentionSchemaError¶
class nnx.abstention.AbstentionSchemaError
A declared schema — labels, model id, probability field, calibrator or sample ids — that contradicts the policy's, the prediction's own, or itself (calibrated probabilities without their calibrator, say).
2.22. Experiment plans (nnx.plans)¶
nnx.plans¶
module nnx.plans
Immutable fluent experiment plans.
Details
An :class:`ExperimentPlan` holds one experiment's configuration — the
network parameters (``NNParams``), the model parameters (``NNModelParams``),
the training parameters (``NNTrainParams``), the data, the seed and the
callbacks — and compiles it to the existing ``NNModel(...)`` +
``NNModel.train(...)`` loop. Every ``with_*`` method returns a **new** plan
and leaves the source unchanged, so branching a configuration is::
base = ExperimentPlan().with_net(net).with_model(model).with_train(train).with_data(loader)
small, large = base.with_epochs(2), base.with_epochs(8) # base, small and large are all intact
Three boundaries are kept apart:
- :meth:`ExperimentPlan.validate` is pure. It aggregates every
configuration problem as a field-path :class:`Diagnostic`, and consumes
no loader, calls no factory, builds no model, reads no weights and writes
no directory.
- :meth:`ExperimentPlan.probe` is effectful but leaves no trace (bar a
CUDA model's first CUDA use, below): on one
example batch it seeds, builds a temporary model and runs a forward pass
(and the loss when the batch has a target) under ``torch.no_grad()``. It
then restores the ambient RNG streams, the cuDNN / deterministic-algorithm
settings and the seeding environment variables, even when the forward
pass fails. It never touches a callback or a run directory.
- :meth:`ExperimentPlan.fit` trains. It seeds, calls the loader and callback
factories once, builds a fresh model and calls ``NNModel.train``, the
same loop an imperative script calls. Each fit is a distinct **attempt**:
its id is folded into the run's ``salt``, so repeated fits of one plan
get distinct run ids and directories, and a plan never sets
``overwrite_existing``. It returns a :class:`FitResult` ``(model, run,
metrics, attempt_id)``.
**Ownership.** The plan's parameter objects are immutable and freely shared
between branches. Everything else is **borrowed**: shared by every branch
and every fit, never copied or cloned, and never reset by the plan.
- Borrowed: loader instances passed to :meth:`~ExperimentPlan.with_data`,
callback instances passed to :meth:`~ExperimentPlan.with_callbacks`, the
``extra_metrics`` callables, step functions, the objective, components
and the provenance manifest.
- A stateful borrowed object carries its state from one fit to the next.
``EarlyStopping`` resets itself; ``LRMonitor`` accumulates.
- For fresh objects per fit, pass zero-argument **factories**: a callable
as ``with_data(train=...)`` / ``val=``, or
:meth:`~ExperimentPlan.with_callback_factories`.
- Each fit builds a new model; ``FitResult.model`` is the trained one.
nnx.plans.ExperimentPlan¶
class nnx.plans.ExperimentPlan(net: 'Optional[NNParams]' = None, model: 'Optional[NNModelParams]' = None, train: 'Optional[NNTrainParams]' = None, train_data: 'Any' = None, val_data: 'Any' = INHERIT, data_identity: 'Any' = INHERIT, seed: 'Any' = None, callbacks: 'tuple[Any, ...]' = (), callback_factories: 'tuple[Any, ...]' = (), train_step_fn: 'Any' = None, eval_step_fn: 'Any' = None, objective: 'Any' = None, components: 'tuple[Any, ...]' = (), provenance: 'Any' = None, resume: 'Optional[_Resume]' = None, batch_adapter: 'Any' = None) -> 'None'
An immutable experiment configuration that compiles to NNModel(net, model).train(train, ...); see the module docstring for the boundaries and the ownership rules.
Details
Build it with the ``with_*`` methods, each of which returns a new plan:
- ``with_net(NNParams)``, ``with_model(NNModelParams)``,
``with_train(NNTrainParams)``;
- ``with_epochs``, ``with_optim``, ``with_scheduler``, ``with_metrics``
and ``with_extra_metrics``, which edit the training parameters;
- ``with_data(train, val, identity=...)`` (what is not given is kept,
and validation data and identity otherwise come from the training
parameters) and ``with_seed(seed)``;
- ``with_callbacks`` (borrowed instances) and
``with_callback_factories`` (called once per fit);
- ``with_step_fns``, ``with_objective``, ``with_components`` and
``with_provenance``, which are borrowed, and ``with_batch_adapter``
for a registered module's inputs;
- ``resuming(run_id, checkpoint="last", mode=None, epochs=None)``
(``mode`` and ``epochs`` default to the training parameters'
``resume_mode`` and ``resume_epochs``).
Plan-level arguments (data, seed, callbacks, factories, steps, resume)
are recorded as given and checked by :meth:`validate`, so every problem
is reported at once. The ``NNParams`` / ``NNModelParams`` /
``NNTrainParams`` objects validate themselves when constructed, so the
methods that edit the training parameters (``with_epochs``,
``with_metrics``, ...) raise :class:`PlanError` naming the field at
once.
nnx.plans.ExperimentPlan.with_net¶
nnx.plans.ExperimentPlan.with_net(self, net: 'Optional[NNParams]') -> 'ExperimentPlan'
The network parameters (None for a ModelSpec model).
nnx.plans.ExperimentPlan.with_model¶
nnx.plans.ExperimentPlan.with_model(self, model: 'NNModelParams') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.with_train¶
nnx.plans.ExperimentPlan.with_train(self, train: 'NNTrainParams') -> 'ExperimentPlan'
The training parameters. Loaders inside them are borrowed, and with_data overrides them. extra_metrics is copied into a read-only mapping.
nnx.plans.ExperimentPlan.with_epochs¶
nnx.plans.ExperimentPlan.with_epochs(self, n_epochs: 'int') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.with_optim¶
nnx.plans.ExperimentPlan.with_optim(self, optim: 'Any') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.with_scheduler¶
nnx.plans.ExperimentPlan.with_scheduler(self, scheduler: 'Any') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.with_metrics¶
nnx.plans.ExperimentPlan.with_metrics(self, metrics: 'Iterable[Any]' = KEEP, *, monitor: 'Any' = KEEP) -> 'ExperimentPlan'
Declared metrics (MetricSpec) and the monitor (MonitorSpec), as on NNTrainParams; the one not given is kept (pass monitor=None to drop a monitor).
nnx.plans.ExperimentPlan.with_extra_metrics¶
nnx.plans.ExperimentPlan.with_extra_metrics(self, extra_metrics: 'Optional[Mapping[str, Callable[..., float]]]') -> 'ExperimentPlan'
Borrowed name -> callable(y_true, y_pred) metrics, copied into a read-only mapping.
nnx.plans.ExperimentPlan.with_data¶
nnx.plans.ExperimentPlan.with_data(self, train: 'Any' = KEEP, val: 'Any' = KEEP, *, identity: 'Any' = KEEP) -> 'ExperimentPlan'
The training data and, optionally, validation data and identity.
Details
Each source is a re-iterable of batches (a ``DataLoader`` or a list,
which is borrowed) or a zero-argument factory that returns one. A
factory is called once per fit, never by :meth:`validate`.
``identity`` is the caller-supplied data identity; it becomes
``NNTrainParams.data_id``, which is part of the run id.
As with the other ``with_*`` methods, what is not given is kept (so
``with_data(val=...)`` changes only the validation source): the
current training and validation sources and identity — the plan's
own, else the training parameters' ``val_loader`` / ``data_id``,
resolved when the plan compiles, whatever the order of the calls.
``train=None`` keeps the training parameters' loader. Pass
``val=None`` to run train-only (a plan without validation data
invents none) and ``identity=`` for new data — or ``identity=None``
to record none. Without ``with_data`` the loaders and ``data_id``
inside the training parameters are used (the loaders may be
factories too).
nnx.plans.ExperimentPlan.with_seed¶
nnx.plans.ExperimentPlan.with_seed(self, seed: 'Optional[int]') -> 'ExperimentPlan'
Seed every RNG before the factories run and the model is built, and again (as NNTrainParams.seed) when training starts. The plan's seed overrides the training parameters' own seed, so branches can differ by seed; None falls back to it.
nnx.plans.ExperimentPlan.with_callbacks¶
nnx.plans.ExperimentPlan.with_callbacks(self, *callbacks: 'Any') -> 'ExperimentPlan'
Borrowed callback instances (or legacy fn(idps) callables), shared by every fit; they replace any earlier ones. A callable whose one parameter is required is a legacy callback, called with the history each epoch. One that can be called with no arguments reads as a factory and is reported: pass callback-making functions to with_callback_factories.
nnx.plans.ExperimentPlan.with_callback_factories¶
nnx.plans.ExperimentPlan.with_callback_factories(self, *factories: 'Callable[[], Any]') -> 'ExperimentPlan'
Zero-argument callables, each called once per fit to make a fresh callback. Their callbacks run after the borrowed ones, in order.
nnx.plans.ExperimentPlan.with_step_fns¶
nnx.plans.ExperimentPlan.with_step_fns(self, train_step_fn: 'Any' = KEEP, eval_step_fn: 'Any' = KEEP) -> 'ExperimentPlan'
Borrowed step functions; the one not given is kept (pass None to drop one).
nnx.plans.ExperimentPlan.with_objective¶
nnx.plans.ExperimentPlan.with_objective(self, objective: 'Any') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.with_components¶
nnx.plans.ExperimentPlan.with_components(self, *components: 'Any') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.with_batch_adapter¶
nnx.plans.ExperimentPlan.with_batch_adapter(self, adapter: 'Any') -> 'ExperimentPlan'
How a registered ModelSpec module reads a batch (nnx.models.PositionalInputs / KeywordInputs), passed to NNModel(..., batch_adapter=...) on every build.
nnx.plans.ExperimentPlan.with_provenance¶
nnx.plans.ExperimentPlan.with_provenance(self, manifest: 'Any') -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.resuming¶
nnx.plans.ExperimentPlan.resuming(self, run_id: 'str', *, checkpoint: 'str' = 'last', mode: 'Optional[str]' = None, epochs: 'Optional[str]' = None) -> 'ExperimentPlan'
Warm-resume from the run run_id's checkpoint. The source run is recorded as the new run's parent lineage and is never overwritten: the resumed fit writes its own run directory. mode ("auto" / "stateful" / "weights_only") defaults to the training parameters' own resume_mode. epochs sets what n_epochs means on the resume and defaults to the training parameters' own resume_epochs: "additional" trains n_epochs more; "planned" continues the source run's plan up to n_epochs in total, which needs its training state (a "weights_only" mode is refused).
nnx.plans.ExperimentPlan.without_resume¶
nnx.plans.ExperimentPlan.without_resume(self) -> 'ExperimentPlan'
No public description is currently available.
nnx.plans.ExperimentPlan.validate¶
nnx.plans.ExperimentPlan.validate(self) -> 'PlanValidation'
Every configuration problem as a field-path Diagnostic.
Details
Pure: it consumes no loader (iterables are never iterated), calls
no data or callback factory, builds no model, reads no weights and
writes no directory. Declared metrics and the monitor are resolved
against their registries, as ``NNModel.train`` does before
reserving a run. It constructs the built-in loss module, as
``NNModel`` does first, and resolves each borrowed callback's monitor
on a shallow copy; with ``nnx.streaming.streaming_eval_step`` it also
builds each declared metric's accumulator to check that it is
bounded.
nnx.plans.ExperimentPlan.probe¶
nnx.plans.ExperimentPlan.probe(self, example_batch: 'Any') -> 'ProbeResult'
Build a temporary model and run one forward pass on example_batch under torch.no_grad() in eval mode, plus a train-mode pass for the loss when the batch has a target (None for an all-masked task batch, which the default step records no loss for either).
Details
Effectful but restored: the plan's seed is applied first, and
afterwards the ambient Python / NumPy / torch RNG streams, the
cuDNN ``deterministic`` / ``benchmark`` flags, the
deterministic-algorithms setting and the seeding environment
variables are put back, whether or not the forward pass succeeds.
The model is discarded. No callback, factory or loader is touched,
and nothing is written. A CPU model's probe seeds no CUDA stream, so
nothing waits for CUDA's first use. A CUDA model's probe starts CUDA
first when it is not yet in use (running any seed queued for its
first use), so its streams are restored like the others. A
registered factory's module is built as ``NNModel`` builds it
(:func:`nnx.models.build_module`), which saves and restores every
RNG stream around the factory call — on a GPU host that reads, and
so initializes, CUDA even for a CPU model; its streams are then put
back as the build left them.
A probe needs only the model, the network and the seed: when those
are invalid it raises :class:`PlanError` first (the data, callbacks
and training parameters are :meth:`validate`'s and :meth:`fit`'s
concern); a failing forward pass raises its own error — including a
batch the training step cannot take either, such as a single example
for a ``BatchNorm`` net in train mode.
nnx.plans.ExperimentPlan.fit¶
nnx.plans.ExperimentPlan.fit(self, *, attempt: 'Optional[str]' = None) -> 'FitResult'
Train one attempt and return its FitResult.
Details
The steps are the ones an equally seeded imperative script takes:
validate (an invalid plan raises :class:`PlanError` before
anything runs); seed; call each data factory and callback factory
exactly once (what they return is checked next, and a
:class:`PlanError` then leaves the seed applied, as the script's
``set_seed`` would); build ``NNModel(net, model)``; and call
``model.train(train, callbacks, ...)``. The ``train`` given to
``model.train`` carries the plan's seed, data, data identity and
resume. The run's ``salt`` is ``"plan-attempt:<attempt>"``, where
``attempt`` defaults to a fresh random id, so every fit gets its own
run id and directory. ``overwrite_existing`` is never set; reusing
an ``attempt`` id for the same configuration raises
``FileExistsError`` from ``NNModel.train``.
The fit trains in this process (no ``distributed=``): inside a
``torch.distributed`` process group of more than one rank (a
``torchrun`` launch) it raises ``RuntimeError`` before seeding,
calling a factory or creating a run; a one-rank group behaves as no
group.
nnx.plans.FitResult¶
class nnx.plans.FitResult(model: 'NNModel', run: 'NNRun', metrics: 'Mapping[str, SplitMetrics]', attempt_id: 'str') -> 'None'
One fit: the trained model, its persisted run (the NNModel.train return value), the final epoch's metrics per split ("train" and "val", see SplitMetrics) and the attempt_id folded into the run's salt.
nnx.plans.SplitMetrics¶
class nnx.plans.SplitMetrics(split: 'str', available: 'bool', epoch: 'Optional[int]' = None, source: 'Optional[str]' = None, values: 'Mapping[str, float]' = <factory>, reason: 'Optional[str]' = None) -> 'None'
One split's metrics from a fit's final epoch, read from the run's own history — never from a second pass over a loader.
Details
``available`` says whether the split was measured: a plan without
validation data reports ``val`` as unavailable, with the ``reason``,
rather than inventing it. ``source`` is ``"epoch"`` for a whole-epoch
record (validation, or a training summary when the run declares metrics
or a monitor) and ``"last_batch"`` when only the final training batch's
record exists. ``values`` maps metric names to floats: the recorded
``loss`` / ``error`` / classification fields, the declared ``metrics``
and the ``extra_metrics`` as ``extra/<name>``, skipping any that were
not recorded (a non-finite value is kept). A declared metric named like
a classification field (``f1``, ``accuracy``, ...) appears as
``metrics/<name>`` beside the record's own field, as the loggers write
it; a monitor on that name compares the declared value. A whole-epoch training
summary carries the declared metrics but not the extra metrics, which
the loop computes per batch (``run.idps`` has them).
nnx.plans.ProbeResult¶
class nnx.plans.ProbeResult(output_shape: 'tuple[int, ...]', output_dtype: 'str', target_shape: 'Optional[tuple[int, ...]]', loss: 'Optional[float]', n_parameters: 'int', n_trainable: 'int') -> 'None'
A probe's forward pass on one example batch: the network's raw output shape and dtype (in eval mode, as predict sees it, before any reshaping the loss applies), the batch's target shape as given, the loss the default training step would compute on the batch (in train mode, under its mixed-precision autocast, without a gradient — so the batch is scored as a training batch: a BatchNorm net needs more than one row) when the batch has a target — None for a plan with its own train_step_fn or objective, whose loss a probe cannot know — and the temporary model's parameter counts.
nnx.plans.PlanValidation¶
class nnx.plans.PlanValidation(diagnostics: 'tuple[Diagnostic, ...]' = ()) -> 'None'
The result of ExperimentPlan.validate: every Diagnostic at once, in field order.
nnx.plans.PlanValidation.ok¶
property nnx.plans.PlanValidation.ok
No public description is currently available.
nnx.plans.PlanValidation.paths¶
property nnx.plans.PlanValidation.paths
No public description is currently available.
nnx.plans.PlanValidation.raise_for_errors¶
nnx.plans.PlanValidation.raise_for_errors(self) -> 'None'
Raise PlanError listing every diagnostic, if any.
nnx.plans.Diagnostic¶
class nnx.plans.Diagnostic(path: 'str', message: 'str') -> 'None'
One configuration problem at a field path ("data.train", "train.monitor", "callbacks[1]", ...).
nnx.plans.PlanError¶
class nnx.plans.PlanError(diagnostics: 'tuple[Diagnostic, ...]') -> 'None'
A plan that cannot be fitted or probed. diagnostics lists every problem ExperimentPlan.validate found.
Details
It is a ``ValueError`` and a ``TypeError``, the two errors the params
classes raise, so code catching either around a configuration still
catches it.
2.23. Streaming prediction and mergeable metrics (nnx.streaming)¶
nnx.streaming¶
module nnx.streaming
Streaming prediction and mergeable metrics.
Details
``NNModel.predict()`` and ``NNModel.evaluate()`` are eager: they return only
after the whole loader has run, holding every batch's outputs until then.
This module adds the bounded counterparts.
**A prediction stream.** ``NNModel.iter_predict(loader)`` returns a
:class:`PredictionStream` that yields one result per loader batch, in loader
order, and holds nothing once a batch has been handed over::
with model.iter_predict(loader) as stream:
for batch in stream: # PredictionBatch(logits, classes, sample_ids)
write(batch.sample_ids, batch.classes)
Concatenating the batches gives exactly what the eager call returns:
``predict(loader)``'s logits and classes, and ``predict_proba(loader)``'s
sample ids — the same seed-row slicing for graph loaders and the same
categorical / multilabel / continuous decoding. With a
:class:`~nnx.prediction.ProbabilitySpec`, or ``rich=True`` for a model with a
task, every batch is a :class:`~nnx.prediction.PredictionResult` instead.
- Each batch runs in eval mode under ``no_grad``, and every submodule's
training mode is restored before the batch is yielded, so the network is in
its own mode between batches — also when the forward pass raises.
- Closing the stream (leaving the ``with`` block, ``close()``, an exception)
drops its references to the loader's iterator and to the model. A closed or
consumed stream refuses to be iterated again. The loader stays the caller's:
it is never closed and can be iterated again.
- The stream holds at most the batch in flight. What the consumer keeps is
the consumer's memory: store only what you need (for example the sample
ids and classes), not the batches.
**Mergeable metrics.** :class:`StreamingMetrics` accumulates declared
:class:`~nnx.MetricSpec` metrics over any number of batches with ``update``,
combines two accumulations over disjoint samples with ``merge`` and returns a
read-only :class:`MetricSnapshot` from ``finalize``. Built-in metrics keep
sums and counts (``accuracy``, ``nll``, ``brier``, ``mae``, ``mse``) or
confusion counts (``f1``), so memory never grows with the number of samples,
and a merge adds them — the result does not depend on the order. A metric that
needs every stored score (a rank metric such as AUROC is never additive) is
rejected unless ``materialize=True`` asks for O(N) storage.
**Bounded validation.** :func:`streaming_eval_step` is a ready-made
``eval_step_fn`` for ``NNModel.train``: the default validation record from
counts and sums instead of stored predictions. ``evaluate()`` itself is
unchanged and eager.
nnx.streaming.PredictionStream¶
class nnx.streaming.PredictionStream(batches: 'Iterator[Any]') -> 'None'
A context-managed iterator of prediction batches; see the module docstring. Obtain one with NNModel.iter_predict.
nnx.streaming.PredictionStream.closed¶
property nnx.streaming.PredictionStream.closed
Whether the stream is closed or consumed (it cannot be iterated again).
nnx.streaming.PredictionStream.close¶
nnx.streaming.PredictionStream.close(self) -> 'None'
Stop the stream and drop its references to the loader's iterator and the model. Idempotent; the loader itself is left to the caller.
nnx.streaming.PredictionBatch¶
class nnx.streaming.PredictionBatch(logits: 'np.ndarray', classes: 'np.ndarray', sample_ids: 'np.ndarray') -> 'None'
One loader batch of NNModel.iter_predict.
Details
Attributes:
logits: the raw network output for the batch's rows (graph loaders:
seed rows only), as ``predict()`` returns them.
classes: the decoded predictions, as ``predict().classes``: argmax
classes, 0/1 indicators for a multilabel or ``BCEWithLogitsLoss``
model, the values themselves for a regression task.
sample_ids: ``int64`` identity of each row, as ``predict_proba()``
reports it: the position in iteration order, or the global node
index for graph seed rows.
nnx.streaming.concatenate_predictions¶
nnx.streaming.concatenate_predictions(batches: 'Iterable[Union[PredictionBatch, PredictionResult]]') -> 'Union[PredictionBatch, PredictionResult]'
Concatenate streamed batches into one result, in order — what the eager predict() / predict_proba() returns for the same loader. Materializes every batch (O(N) memory).
nnx.streaming.StreamingMetrics¶
class nnx.streaming.StreamingMetrics(metrics: 'Sequence[MetricSpec]', semantics: 'str', *, labels: 'Optional[Sequence[str]]' = None, threshold: 'float' = 0.5, materialize: 'bool' = False, ignore_index: 'Optional[int]' = None, num_outputs: 'Optional[int]' = None) -> 'None'
Mergeable accumulators for declared metrics (see the module docstring).
Details
Args:
metrics: the :class:`~nnx.MetricSpec` s to accumulate (unique names).
semantics: how predictions are read — ``"categorical"`` (class
probabilities over the last axis of ``probabilities``, argmax
labels), ``"bernoulli"`` (independent per-output probabilities,
labels at ``threshold``) or ``"continuous"`` (values). Every
metric's input must be one these semantics provide.
labels: optional ordered class / output names; a categorical
``probabilities`` width must match them. Part of the merge schema.
threshold: the bernoulli decision threshold for labels derived from
probabilities (default 0.5).
materialize: store every score of a metric that has no bounded form
(O(N) memory) instead of rejecting it.
ignore_index: categorical only — the target value that is never
scored (e.g. ``-100`` for padding).
num_outputs: the class count (categorical) or the output width along
axis 1 (bernoulli, continuous), when no ``labels`` name them;
batches of another width are refused.
Two accumulators merge only when their metrics (id, version, config and
name), semantics, labels, output count and ignore index are equal, and,
for bernoulli semantics, the decision threshold. If a
metric's update raises part-way through a batch, the accumulation no
longer describes one set of samples and refuses further use.
nnx.streaming.StreamingMetrics.for_task¶
nnx.streaming.StreamingMetrics.for_task(metrics: 'Sequence[MetricSpec]', task: 'TaskSpec', *, materialize: 'bool' = False) -> 'StreamingMetrics'
Accumulators for a model's TaskSpec: categorical (with the task's ignore_index), multilabel (bernoulli at the task's threshold) or regression (continuous), with the task's labels.
nnx.streaming.StreamingMetrics.metrics¶
property nnx.streaming.StreamingMetrics.metrics
No public description is currently available.
nnx.streaming.StreamingMetrics.semantics¶
property nnx.streaming.StreamingMetrics.semantics
No public description is currently available.
nnx.streaming.StreamingMetrics.labels¶
property nnx.streaming.StreamingMetrics.labels
No public description is currently available.
nnx.streaming.StreamingMetrics.threshold¶
property nnx.streaming.StreamingMetrics.threshold
No public description is currently available.
nnx.streaming.StreamingMetrics.ignore_index¶
property nnx.streaming.StreamingMetrics.ignore_index
No public description is currently available.
nnx.streaming.StreamingMetrics.num_outputs¶
property nnx.streaming.StreamingMetrics.num_outputs
The class count (categorical) or output width, if declared (num_outputs, else the labels').
nnx.streaming.StreamingMetrics.count¶
property nnx.streaming.StreamingMetrics.count
The valid samples scored so far.
nnx.streaming.StreamingMetrics.bounded¶
property nnx.streaming.StreamingMetrics.bounded
Whether no metric stores its scores (memory independent of N).
nnx.streaming.StreamingMetrics.update¶
nnx.streaming.StreamingMetrics.update(self, target: 'Any', *, probabilities: 'Any' = None, labels: 'Any' = None, values: 'Any' = None, valid: 'Any' = None) -> 'int'
Add one batch and return the number of valid samples scored.
Details
``target`` holds class indices ``(N,)`` (categorical), 0/1 outcomes
(bernoulli; soft targets count as 1 from 0.5) or values (continuous).
Give each input a declared metric reads: ``probabilities`` —
``(N, C)`` class probabilities (categorical) or per-output
probabilities shaped like ``target`` (bernoulli); ``labels`` — the
decoded predictions (derived from ``probabilities`` when omitted);
``values`` — continuous predictions shaped like ``target``.
``valid`` (boolean, shaped like ``target``) excludes entries; a NaN
target, or a categorical target equal to ``ignore_index``, is
excluded too. Arrays and tensors are accepted.
nnx.streaming.StreamingMetrics.update_logits¶
nnx.streaming.StreamingMetrics.update_logits(self, target: 'Any', logits: 'Any', *, valid: 'Any' = None) -> 'int'
Add one batch of raw model outputs (class axis 1): softmax (categorical), sigmoid decided at threshold (bernoulli) or the values (continuous) — the inputs evaluate() derives. Returns the number of valid samples scored.
nnx.streaming.StreamingMetrics.merge¶
nnx.streaming.StreamingMetrics.merge(self, other: 'StreamingMetrics') -> 'StreamingMetrics'
A new accumulation covering both inputs' samples; neither input changes, and later updates to either never reach the result. Raises MetricMergeError for different declarations.
nnx.streaming.StreamingMetrics.finalize¶
nnx.streaming.StreamingMetrics.finalize(self) -> 'MetricSnapshot'
The metrics over every sample seen, as a read-only snapshot. Repeatable, and never changes the accumulation.
nnx.streaming.MetricSnapshot¶
class nnx.streaming.MetricSnapshot(count: 'int', values: 'Mapping[str, float]' = <factory>, unavailable: 'tuple[str, ...]' = ()) -> 'None'
A finalized StreamingMetrics: read-only values that later updates never change.
Details
Attributes:
count: the valid samples scored (categorical rows; multilabel and
continuous entries). ``0`` for an empty or fully masked stream.
values: each available metric's value, by name.
unavailable: the declared metrics with no value (every one when
``count`` is 0).
nnx.streaming.MetricSnapshot.available¶
property nnx.streaming.MetricSnapshot.available
No public description is currently available.
nnx.streaming.streaming_eval_step¶
nnx.streaming.streaming_eval_step(ctx: 'EvalStepContext') -> 'NNEvaluationDataPoint'
An eval_step_fn for NNModel.train that builds the default validation record from counts and sums.
Details
The record is the one NNx's own ``evaluate()`` returns — the same task
counts and status, loss (each batch's numerator over the summed loss
denominators), classification or task metrics and declared metrics —
without storing a target or prediction per sample, so memory does not
grow with the validation set. A model subclass that overrides
``evaluate()`` keeps its override only on the default validation step. It cannot compute ``extra_metrics`` (callables on the
full arrays) or a declared metric that needs every stored score; both are
rejected when training starts. ``evaluate()`` and the default validation
step are unchanged.
nnx.streaming.StreamClosedError¶
class nnx.streaming.StreamClosedError
A closed or already consumed PredictionStream was used again.
nnx.streaming.MetricMergeError¶
class nnx.streaming.MetricMergeError
Two StreamingMetrics with different declarations were merged: other metrics (id, version or config), other probability semantics, other task labels, another ignore index or (bernoulli) another decision threshold.
2.24. Run bundles (nnx.bundles)¶
nnx.bundles.export_bundle¶
nnx.bundles.export_bundle(run_id: 'str', destination: 'Union[str, os.PathLike[str]]', *, checkpoint: 'str' = 'last', root: 'Optional[str]' = None, calibrators: 'Iterable[TemperatureCalibrator]' = ()) -> 'BundleInfo'
Publish run_id's checkpoint as a run bundle at destination.
Details
Args:
run_id: a run under ``<root>/runs`` (the current directory by default)
— your own run: its pickle checkpoint is read locally.
destination: an empty or missing directory, or an existing bundle,
whose generation this export replaces.
checkpoint: a ``Checkpoints`` tag (``"last"``, ``"best"``, …) or a
``ModelCheckpoint`` file stem ``"<tag>_e<epoch>"``.
calibrators: ``TemperatureCalibrator`` records to ship with the
model. Their labels must match the task's, and a fingerprint
``model_id`` (``nnx.calibration.model_fingerprint``) must be the
bundled weights'.
Returns the published bundle's :class:`BundleInfo` (``verified=True``).
A checkpoint written by ``train()`` gives a ``"resume"`` bundle; a
weights-only one an ``"inference"`` bundle. A runtime-only module, a
recorded topology transform NNx cannot replay from data (anything but an
``nnx.transforms`` recipe operation of a known version or a torchao QAT
conversion — the error names its index and id), module extra state that
is not a tensor, and training state that is not tensors and JSON
primitives are refused before anything is written.
nnx.bundles.inspect_bundle¶
nnx.bundles.inspect_bundle(path: 'Union[str, os.PathLike[str]]') -> 'BundleInfo'
Summarize the bundle at path from its manifest, state.json and calibrator records (each checked against its manifest size and SHA-256). Tensor payloads are not read or checked — validate_bundle checks them. Never unpickles, calls a factory or downloads.
nnx.bundles.validate_bundle¶
nnx.bundles.validate_bundle(path: 'Union[str, os.PathLike[str]]') -> 'BundleInfo'
Check the whole bundle at path and summarize it: every listed payload present, a regular file inside the published generation directory with the manifest's size and SHA-256, nothing unlisted, one generation id throughout, strict JSON without duplicate keys, and safetensors headers holding exactly the tensors state.json references. No tensor is read before all of that passes (a calibrator with a fingerprint model_id is then checked against the weights); nothing is unpickled, and no factory is called. Raises BundleIntegrityError (or BundleError for a path that is not a bundle).
nnx.bundles.reconstruct_bundle¶
nnx.bundles.reconstruct_bundle(path: 'Union[str, os.PathLike[str]]', *, factories: 'Optional[Mapping[tuple[str, int], ModelFactory]]' = None, components: 'Optional[Iterable[Any]]' = None, device: 'Any' = None, batch_adapter: 'Optional[BatchAdapter]' = None, exclude_submodules: 'Sequence[str]' = ()) -> 'ReconstructedBundle'
Validate the bundle at path and rebuild its model.
Details
Args:
factories: the model factories to rebuild a registered ``ModelSpec``
from, ``{(id, version): factory}`` — used instead of the process
registry. ``None`` uses the process registry
(``nnx.models.register_model_factory``). A built-in net needs none.
components: for a ``"resume"`` bundle, the stateful components the
resumed run will register (e.g. the same callbacks): their saved
state is checked against them. ``None`` skips the check (for
inference use; :meth:`ReconstructedBundle.resume` checks again).
device: a ``Devices`` member to build the model on instead of the
saved one.
batch_adapter: how a registered module sees a batch
(``nnx.models.BatchAdapter``), as passed to ``NNModel`` — it is
runtime-only, never stored.
exclude_submodules: top-level submodules the training run attached
to the net that are not part of the rebuilt architecture — a
JEPA predictor registered as ``model.net._jepa_predictor``, say
— whose weights are left out, as in
``NNModel.from_checkpoint``; every other weight still loads
strictly. A model rebuilt without them serves inference and
cannot :meth:`~ReconstructedBundle.resume`.
Everything missing — the model factory, a required or incompatible
component — is reported in one :class:`BundleReconstructionError` before
any model is allocated. Calibrators with a fingerprint ``model_id`` are
checked against the rebuilt weights.
nnx.bundles.BundleInfo¶
class nnx.bundles.BundleInfo(path: 'str', version: 'int', generation: 'str', capability: 'str', source_run_id: 'str', source_checkpoint: 'str', epoch: 'int', model: 'str', model_params: 'Mapping[str, Any]', net_params: 'Optional[Mapping[str, Any]]' = None, components: 'Mapping[str, Any]' = <factory>, calibrators: 'tuple[Mapping[str, Any], ...]' = (), payloads: 'Mapping[str, Any]' = <factory>, verified: 'bool' = False) -> 'None'
A bundle's summary, from inspect_bundle (verified=False: only the manifest and state.json were checked) or validate_bundle / reconstruct_bundle (verified=True).
Details
Attributes:
path: the bundle directory.
version: the bundle format version.
generation: the published generation id.
capability: ``"resume"`` (weights and training state) or
``"inference"`` (weights only).
source_run_id / source_checkpoint: the run and checkpoint it came from.
epoch: the checkpoint's epoch.
model: the network descriptor (a built-in ``Nets`` name or a
registered ``ModelSpec``).
model_params: ``NNModelParams.state()`` of the checkpoint.
net_params: the built-in net's params ``state()`` (its dimensions:
the architecture reconstruction builds); ``None`` for a
registered module, whose architecture its factory builds.
components: saved component metadata, ``{name: {"version", "required"}}``.
calibrators: each calibrator's ``{"payload", "id", "labels", "model_id"}``.
payloads: ``{name: {"sha256", "size"}}`` from the manifest.
verified: whether every payload was checked.
nnx.bundles.ReconstructedBundle¶
class nnx.bundles.ReconstructedBundle(model: 'NNModel', info: 'BundleInfo', calibrators: 'tuple[TemperatureCalibrator, ...]', _checkpoint: 'NNCheckpoint', _training_state: 'Optional[dict[str, Any]]', _excluded: 'tuple[str, ...]' = ()) -> 'None'
A model rebuilt from a validated bundle, its calibrators and — for a "resume" bundle — the training state resume continues from.
nnx.bundles.ReconstructedBundle.capability¶
property nnx.bundles.ReconstructedBundle.capability
No public description is currently available.
nnx.bundles.ReconstructedBundle.resume_checkpoint¶
property nnx.bundles.ReconstructedBundle.resume_checkpoint
The checkpoint label a resumed run records as its parent checkpoint ("bundle-<generation[:16]>_e<epoch>": the first 16 hex characters of the generation).
nnx.bundles.ReconstructedBundle.resume¶
nnx.bundles.ReconstructedBundle.resume(self, params: 'NNTrainParams', **train_kwargs: 'Any') -> 'NNRun'
Continue training model from the bundle's training state: model.train(params, **train_kwargs) as a stateful resume of the bundle's run and checkpoint — the optimizer, scheduler, GradScaler, RNG and component state are restored and epoch numbering continues, exactly as resuming the original run on disk would. params must not name a resume source of its own.
Details
An ``"inference"`` bundle raises :class:`BundleCapabilityError`
before anything changes, and so does a ``Trainer`` run's bundle (named
optimizers): this continues ``NNModel.train`` runs. ``resume_mode`` is
always ``"stateful"`` here.
nnx.bundles.BundleError¶
class nnx.bundles.BundleError
A run bundle that cannot be written, read or used.
nnx.bundles.BundleIntegrityError¶
class nnx.bundles.BundleIntegrityError
A bundle that fails validation: a missing, extra, altered, symlinked or out-of-root payload, a generation mismatch, duplicate JSON keys, an unknown format or version, or malformed state.
nnx.bundles.BundleCapabilityError¶
class nnx.bundles.BundleCapabilityError
An operation the bundle's capability does not support — resuming an "inference" bundle, which has no training state.
nnx.bundles.BundleReconstructionError¶
class nnx.bundles.BundleReconstructionError(problems: 'Iterable[str]') -> 'None'
What a reconstruction needs and was not given (a model factory, a component), every problem listed in problems; raised before any model is allocated.
2.25. Replayable transformation recipes (nnx.transforms)¶
nnx.transforms.TransformRecipe¶
class nnx.transforms.TransformRecipe(operations: 'tuple[TransformOp, ...]', materialization: "Literal['fresh', 'in_place']" = 'in_place') -> 'None'
An ordered, immutable list of TransformOp.
Details
``materialization`` says what :meth:`materialize` does: ``"in_place"``
(the default) transforms the given model itself — its trained weights
are what LoRA wraps and SVD factorizes; ``"fresh"`` builds a fresh,
randomly initialized registered base from the given model's descriptor
and transforms that, leaving the given model untouched. The operations
are copied: changing the list they came from changes nothing here.
nnx.transforms.TransformRecipe.checkpoint_transforms¶
nnx.transforms.TransformRecipe.checkpoint_transforms(self) -> 'tuple[NNCheckpointTransform, ...]'
The operations as checkpoint transforms, in order.
nnx.transforms.TransformRecipe.validate¶
nnx.transforms.TransformRecipe.validate(self, model: 'NNModel', *, optimizers: 'Iterable[torch.optim.Optimizer]' = ()) -> 'None'
Check the whole recipe as materialize would — against model in place, or against the fresh base model describes — mutating nothing, and raise one RecipeError naming every problem. optimizers are checked for an in-place materialization (a fresh base's parameters belong to no existing optimizer).
nnx.transforms.TransformRecipe.materialize¶
nnx.transforms.TransformRecipe.materialize(self, model: 'NNModel', *, optimizers: 'Iterable[torch.optim.Optimizer]' = ()) -> 'NNModel'
Validate the whole recipe (as validate), then apply it — to model itself ("in_place", returned) or to a fresh registered base built from model's descriptor ("fresh", returned). The operations are recorded on the returned model, so its checkpoints and Hub saves rebuild the same topology. Transactional: a failure leaves the model — and the global random streams a fresh base or a LoRA initialization draws from — as they were.
nnx.transforms.TransformOp¶
class nnx.transforms.TransformOp(id: 'str', targets: 'tuple[str, ...]', config: 'Mapping[str, Any]' = <factory>, version: 'int' = 1) -> 'None'
One recorded operation: id ("lora" or "low_rank"), version, the explicit targets (dotted paths under model.net) and an immutable config. Build one with lora or low_rank.
nnx.transforms.TransformOp.state¶
nnx.transforms.TransformOp.state(self) -> 'dict[str, Any]'
The JSON-ready form a checkpoint records (options of its NNCheckpointTransform).
nnx.transforms.TransformOp.checkpoint_transform¶
nnx.transforms.TransformOp.checkpoint_transform(self) -> 'NNCheckpointTransform'
No public description is currently available.
nnx.transforms.TransformOp.from_checkpoint_transform¶
nnx.transforms.TransformOp.from_checkpoint_transform(transform: 'NNCheckpointTransform') -> 'TransformOp'
The operation a checkpoint recorded; malformed options raise a RecipeError.
nnx.transforms.lora¶
nnx.transforms.lora(*targets: 'str', r: 'int' = 8, alpha: 'float' = 16.0, dropout: 'float' = 0.0) -> 'TransformOp'
Wrap each nn.Linear at targets in a LoRALinear (r, alpha, dropout as for apply_lora_to). Targets are explicit module paths, not globs.
nnx.transforms.low_rank¶
nnx.transforms.low_rank(*targets: 'str', rank: 'int', method: 'str' = 'svd') -> 'TransformOp'
Replace each nn.Linear at targets by its rank-rank factorization (nnx.surgery.low_rank_factorize).
nnx.transforms.check_optimizer¶
nnx.transforms.check_optimizer(model: 'NNModel', optimizer: 'torch.optim.Optimizer') -> 'None'
Refuse an optimizer built before the model's recipe: one holding parameters the recipe replaced (a low-rank operation's layers, or the model a fresh materialization started from), or holding a LoRA target's frozen base weights but none of the adapter built around them — whenever it was built, since the adapter would never train. An optimizer built afterwards over the trainable parameters — all, a subset, with parameters outside model.net, or BitFit-style over unfrozen base tensors — passes.
nnx.transforms.RecipeError¶
class nnx.transforms.RecipeError(problems: 'Sequence[_Problem]') -> 'None'
A recipe that does not fit a model, or a malformed operation. problems lists (operation index, operation id, target, reason); the message names each.
2.26. Execution precision (nnx.precision)¶
nnx.precision¶
module nnx.precision
Explicit FP32 / FP16 / BF16 execution policy.
Details
``NNModelParams.precision`` takes a :class:`PrecisionPolicy`: the precision
training, evaluation and prediction run in. A policy is *resolved* against
the device it runs on — once per run, before the run is reserved — into a
:class:`ResolvedPrecision`, which says what is requested, what runs, and
why a fallback happened:
======== =========================================== ==================
mode runs as where
======== =========================================== ==================
``fp32`` full precision (no autocast, no scaler) every device
``fp16`` float16 autocast + a ``GradScaler`` CUDA
``bf16`` bfloat16 autocast, no scaler CPU; CUDA devices
with bf16 support
======== =========================================== ==================
A request the device cannot run fails before any run is reserved, unless
the policy names ``fallback="fp32"``: the run then trains in full precision
and records why. The model's parameters stay float32 in every mode
(autocast only — never ``model.half()``), the backward pass runs outside
autocast, and the update keeps its unscale → normalize → finite-check →
clip → step order.
The legacy ``NNModelParams(mixed_precision=True)`` keeps its meaning — FP16
autocast with a scaler on CUDA, silently full precision elsewhere — and its
run id; combining it with a policy other than ``fp16`` is refused.
The policy covers ``NNModel.train`` (the default step and objectives),
``Trainer.train(objective=...)``, ``evaluate()`` and prediction. It does
not cover :mod:`nnx.lr_finder`, :mod:`nnx.diffusion.sampling` or
generation, and TF32 is reported (never set) as separate metadata.
nnx.precision.PrecisionPolicy¶
class nnx.precision.PrecisionPolicy(mode: 'PrecisionMode' = 'fp32', fallback: 'PrecisionFallback' = 'error') -> 'None'
The requested execution precision.
Details
Args:
mode: ``"fp32"`` (the default), ``"fp16"`` or ``"bf16"``.
fallback: what happens when the device cannot run ``mode``:
``"error"`` (the default) fails before any run is reserved;
``"fp32"`` runs in full precision and records why.
nnx.precision.PrecisionPolicy.state¶
nnx.precision.PrecisionPolicy.state(self) -> 'dict[str, str]'
{"mode": ...}, plus "fallback" when it is not the default "error".
nnx.precision.PrecisionPolicy.from_state¶
nnx.precision.PrecisionPolicy.from_state(state: 'Mapping[str, Any]') -> 'PrecisionPolicy'
No public description is currently available.
nnx.precision.PrecisionPolicy.resolve¶
nnx.precision.PrecisionPolicy.resolve(self, device: 'Union[str, torch.device]') -> 'ResolvedPrecision'
What this policy runs as on device (see the module table).
nnx.precision.ResolvedPrecision¶
class nnx.precision.ResolvedPrecision(requested: 'str', effective: 'str', device_type: 'str', source: 'str' = 'default', fallback_reason: 'Optional[str]' = None, tf32: 'Mapping[str, Optional[bool]]' = <factory>, covers: 'Optional[tuple[str, ...]]' = None) -> 'None'
A policy resolved against a device: what was requested, what runs, and why a fallback happened.
Details
``autocast()`` is the forward context, ``build_scaler()`` the
``GradScaler`` (FP16 only) and ``record()`` the run-inspection
mapping (``NNRun.precision``).
nnx.precision.ResolvedPrecision.scoped¶
nnx.precision.ResolvedPrecision.scoped(self, covers: 'tuple[str, ...]') -> 'ResolvedPrecision'
This resolution, recording that it applies to covers only.
nnx.precision.ResolvedPrecision.applies_to¶
property nnx.precision.ResolvedPrecision.applies_to
No public description is currently available.
nnx.precision.ResolvedPrecision.autocast_dtype¶
property nnx.precision.ResolvedPrecision.autocast_dtype
No public description is currently available.
nnx.precision.ResolvedPrecision.reduced¶
property nnx.precision.ResolvedPrecision.reduced
Whether anything runs below full precision.
nnx.precision.ResolvedPrecision.uses_scaler¶
property nnx.precision.ResolvedPrecision.uses_scaler
No public description is currently available.
nnx.precision.ResolvedPrecision.autocast¶
nnx.precision.ResolvedPrecision.autocast(self) -> 'contextlib.AbstractContextManager[Any]'
The forward-pass context: autocast to the effective dtype, or nothing in full precision.
nnx.precision.ResolvedPrecision.build_scaler¶
nnx.precision.ResolvedPrecision.build_scaler(self) -> 'Optional[Any]'
A fresh GradScaler for FP16, None otherwise.
nnx.precision.ResolvedPrecision.output¶
nnx.precision.ResolvedPrecision.output(self, tensor: 'torch.Tensor') -> 'torch.Tensor'
tensor as a caller sees it: a reduced-precision floating output of an autocast forward is returned as float32, so outputs keep their full-precision schema (and BF16 converts to a dtype NumPy can hold).
nnx.precision.ResolvedPrecision.record¶
nnx.precision.ResolvedPrecision.record(self) -> 'dict[str, Any]'
The run-inspection mapping: requested and effective precision, the fallback reason, TF32 (separately), and what the policy covers — never nnx.lr_finder or nnx.diffusion.sampling.
nnx.precision.ResolvedPrecision.from_record¶
nnx.precision.ResolvedPrecision.from_record(record: 'Mapping[str, Any]') -> 'ResolvedPrecision'
No public description is currently available.
nnx.precision.PrecisionUnsupportedError¶
class nnx.precision.PrecisionUnsupportedError
A precision the device cannot run, with no fallback chosen — or a step function that cannot apply a reduced precision.
nnx.precision.ReferenceTolerance¶
class nnx.precision.ReferenceTolerance(loss: 'float', weights: 'float') -> 'None'
How far a reduced-precision run may drift from its FP32 reference on NNx's seeded fixtures: the largest absolute difference of a per-epoch training loss, and of any trained parameter.
nnx.precision.resolve_precision¶
nnx.precision.resolve_precision(params: 'Any', device: 'Union[str, torch.device, Any]') -> 'ResolvedPrecision'
The precision a model's parameters (NNModelParams) run in on device: its policy, the legacy mixed_precision flag (FP16 on CUDA, full precision elsewhere — its historical meaning), or full precision.
nnx.precision.precision_support¶
nnx.precision.precision_support(device: 'Union[str, torch.device, None]' = None) -> 'dict[str, dict[str, str]]'
The precision matrix as this host sees it: for each device type ("cpu", "cuda", "mps", or just device's), each mode's status —
Details
- ``"verified"``: NNx's seeded reference fixtures hold it within
:data:`REFERENCE_TOLERANCES` (CPU FP32 and BF16);
- ``"supported"``: the hardware is present and the mode resolves, but
no NNx fixture vouches for it here;
- ``"unverified"``: the hardware is absent, so nothing is claimed;
- ``"unsupported"``: the device type cannot run the mode.
nnx.precision.full_precision_only¶
nnx.precision.full_precision_only(step: 'Any') -> 'Any'
Mark step as running in full precision only and return it.
Details
``NNModel.train`` refuses a reduced precision policy for a marked step
before anything is reserved. A wrapper around a built-in paradigm step
keeps that refusal when it is marked too (or built with
``functools.wraps``, which copies the marker).
2.27. Multi-seed summaries and paired comparisons (nnx.comparison)¶
nnx.comparison.Metric¶
class nnx.comparison.Metric(name: 'str', direction: 'str', unit: 'Optional[str]' = None) -> 'None'
What a value measures: its name, whether larger or smaller is better (direction: "maximize" / "minimize") and its unit (None: unitless or not declared).
nnx.comparison.Metric.state¶
nnx.comparison.Metric.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.Metric.from_state¶
nnx.comparison.Metric.from_state(state: 'Mapping[str, Any]') -> 'Metric'
No public description is currently available.
nnx.comparison.Observation¶
class nnx.comparison.Observation(run_id: 'str', metric: 'Metric', value: 'Optional[float]', split: 'str', selection: 'str', status: 'str' = 'completed', attempt_id: 'Optional[str]' = None, config: 'Optional[str]' = None, data: 'Optional[str]' = None, split_id: 'Optional[str]' = None, replicate: 'Optional[str]' = None, evaluation: 'Optional[str]' = None) -> 'None'
One attempt's committed evaluation of one metric.
Details
Attributes:
run_id: the run that produced it.
attempt_id: the attempt within that run, when known.
metric: what ``value`` measures.
value: the value (``None`` when the attempt has none; NaN and
infinities are kept and counted as non-finite).
status: the attempt's completion status: ``"completed"``,
``"failed"``, ``"cancelled"``, ``"running"`` or ``"unknown"``.
split: the evaluation split the value was measured on (for example
``"validation"``).
selection: the checkpoint-selection rule behind the value (for
example ``"last"``, or ``"best:loss"`` for a monitor's election).
config: the configuration identity.
data: the data identity.
split_id: the split identity (a split manifest digest, for example).
replicate: the replicate key (for example ``"seed=3"``) — what
:func:`compare` pairs on.
evaluation: the committed evaluation behind the value (for example
the epoch record and the checkpoint it was committed with).
``config``, ``data``, ``split_id``, ``replicate``, ``attempt_id`` and
``evaluation`` are ``None`` when unknown.
nnx.comparison.Observation.id¶
property nnx.comparison.Observation.id
The attempt id, or run:<run id> when the attempt is unknown.
nnx.comparison.Observation.finite¶
property nnx.comparison.Observation.finite
Whether this is a completed attempt with a finite value.
nnx.comparison.Observation.state¶
nnx.comparison.Observation.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.Observation.from_state¶
nnx.comparison.Observation.from_state(state: 'Mapping[str, Any]') -> 'Observation'
No public description is currently available.
nnx.comparison.summarize¶
nnx.comparison.summarize(observations: 'Iterable[Observation]') -> 'Summary'
Group observations into poolable groups (same metric, split, selection rule, configuration, data and split identity) and summarize each; see GroupSummary. Duplicate attempts are refused: an attempt gives one observation per metric, so the same attempt observed on two splits or under two selection rules belongs in two summaries. A known identity never pools with an unknown one.
nnx.comparison.Summary¶
class nnx.comparison.Summary(groups: 'tuple[GroupSummary, ...]') -> 'None'
A stratified summary: one GroupSummary per poolable group, and the fields that differ between groups (differing).
nnx.comparison.Summary.differing¶
property nnx.comparison.Summary.differing
No public description is currently available.
nnx.comparison.Summary.observations¶
property nnx.comparison.Summary.observations
No public description is currently available.
nnx.comparison.Summary.pooled¶
nnx.comparison.Summary.pooled(self) -> 'GroupSummary'
The single group, when every observation may be pooled; otherwise IncompatibleObservations naming the differing fields.
nnx.comparison.Summary.state¶
nnx.comparison.Summary.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.GroupSummary¶
class nnx.comparison.GroupSummary(metric: 'Metric', split: 'str', selection: 'str', config: 'Optional[str]', data: 'Optional[str]', split_id: 'Optional[str]', observations: 'tuple[Observation, ...]') -> 'None'
One poolable group: the shared fields (metric, split, selection, config, data, split_id), every observation (sorted), and the statistics over its completed, finite values.
Details
``n`` counts those values; ``n_attempts`` every observation,
``n_failed`` those not completed (failed, cancelled, running or of
unknown status) and ``n_nonfinite`` completed ones without a finite
value. ``std`` is the sample standard deviation (``n - 1``): ``None``
when ``n < 2``; ``mean`` is ``None`` when ``n == 0``. ``unknown`` names
the identity fields that are not known.
nnx.comparison.GroupSummary.unknown¶
property nnx.comparison.GroupSummary.unknown
No public description is currently available.
nnx.comparison.GroupSummary.label¶
nnx.comparison.GroupSummary.label(self) -> 'str'
No public description is currently available.
nnx.comparison.GroupSummary.state¶
nnx.comparison.GroupSummary.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.compare¶
nnx.comparison.compare(a: 'Iterable[Observation]', b: 'Iterable[Observation]', *, pairing: 'str', bootstrap: 'Optional[Bootstrap]' = None) -> 'PairedComparison'
Pair configuration b against a by replicate key.
Details
Each side must be one poolable group with a unique replicate key per
observation, and the sides must agree on everything but the
configuration (metric, split, selection rule, data and split identity).
``pairing`` declares what a shared replicate key means (for example
``"same seed for initialisation and data order, same split"``): NNx
cannot prove it, so it is required and recorded. Deltas are signed
``B - A``.
nnx.comparison.PairedComparison¶
class nnx.comparison.PairedComparison(a: 'GroupSummary', b: 'GroupSummary', pairing: 'str', bootstrap: 'Optional[Bootstrap]' = None) -> 'None'
Configuration B against A over shared replicate keys.
Details
``pairs`` holds every matched replicate (``delta`` is ``B - A``, never
flipped: for a metric to minimize, a negative delta means B is lower);
``unmatched_a`` / ``unmatched_b`` list the replicate keys found on one
side only, with their attempt ids. ``n``, ``mean`` and ``std`` (``n - 1``)
summarize the finite deltas; ``n_nonfinite`` counts the pairs of two
finite values whose delta overflows a float (excluded from ``n`` and
reported); ``interval`` is the bootstrap's, when one was requested and
``n >= 2``.
nnx.comparison.PairedComparison.metric¶
property nnx.comparison.PairedComparison.metric
No public description is currently available.
nnx.comparison.PairedComparison.state¶
nnx.comparison.PairedComparison.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.Pair¶
class nnx.comparison.Pair(replicate: 'str', a: 'str', b: 'str', delta: 'Optional[float]') -> 'None'
One replicate key on both sides: the observations' ids and the signed delta b - a (None unless both are completed and finite).
nnx.comparison.Pair.state¶
nnx.comparison.Pair.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.Bootstrap¶
class nnx.comparison.Bootstrap(seed: 'int', resamples: 'int' = 2000, level: 'float' = 0.95) -> 'None'
A percentile bootstrap of the mean paired delta: resamples resamples of the replicate pairs (the unit), drawn with replacement by numpy.random.default_rng(seed), giving a level interval. It describes seed variability over these replicates — not a population-level or significance claim.
nnx.comparison.Bootstrap.state¶
nnx.comparison.Bootstrap.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.Bootstrap.interval¶
nnx.comparison.Bootstrap.interval(self, deltas: 'Sequence[float]') -> 'Optional[tuple[float, float]]'
No public description is currently available.
nnx.comparison.ComparisonReport¶
class nnx.comparison.ComparisonReport(summary: 'Summary', comparisons: 'tuple[PairedComparison, ...]' = ()) -> 'None'
A summary of observations and any paired comparisons, as one artifact: table (machine-readable rows), text (a concise view), state / save / load (strict JSON, nnx.comparison/1).
nnx.comparison.ComparisonReport.build¶
nnx.comparison.ComparisonReport.build(observations: 'Iterable[Observation]', comparisons: 'Sequence[tuple[Iterable[Observation], Iterable[Observation], str]]' = (), *, bootstrap: 'Optional[Bootstrap]' = None) -> 'ComparisonReport'
Summarize observations and run each (a, b, pairing) comparison (every compared observation must be among them).
nnx.comparison.ComparisonReport.table¶
nnx.comparison.ComparisonReport.table(self) -> 'list[dict[str, Any]]'
One row per observation, group, pair and unmatched replicate, plus one comparison summary row per comparison (n, mean, std, bootstrap, interval), each with a row kind; values are plain JSON types.
nnx.comparison.ComparisonReport.text¶
nnx.comparison.ComparisonReport.text(self) -> 'str'
A concise, fixed-order text view.
nnx.comparison.ComparisonReport.state¶
nnx.comparison.ComparisonReport.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.comparison.ComparisonReport.to_json¶
nnx.comparison.ComparisonReport.to_json(self) -> 'str'
No public description is currently available.
nnx.comparison.ComparisonReport.save¶
nnx.comparison.ComparisonReport.save(self, path: 'Union[str, os.PathLike[str]]') -> 'None'
No public description is currently available.
nnx.comparison.ComparisonReport.from_state¶
nnx.comparison.ComparisonReport.from_state(state: 'Mapping[str, Any]') -> 'ComparisonReport'
Rebuild a report from its observations and comparison specs; the stored results must equal the re-derived ones.
nnx.comparison.ComparisonReport.load¶
nnx.comparison.ComparisonReport.load(path: 'Union[str, os.PathLike[str]]') -> 'ComparisonReport'
No public description is currently available.
nnx.comparison.observations_from_runs¶
nnx.comparison.observations_from_runs(run_ids: 'Iterable[str]', *, metric: 'Metric', split: 'str' = 'validation', selection: 'str' = 'last', root: 'Optional[str]' = None, replicate: 'Optional[str]' = 'seed', config: 'Optional[Mapping[str, str]]' = None) -> 'list[Observation]'
Observations of metric read from saved runs, without loading a model or a checkpoint (run.yaml, idps.csv — or, for a run trained with a HistoryJournal, the journal's committed records, each chunk checked — and the provenance files are read once each — plus, for a run with a parent, its metadata.yaml and every ancestor's run.yaml and provenance files, whose data and split identities join the configuration identity and whose attempt records decide which parent attempt a child started from; an ancestor's unreadable file is refused — and nothing is written).
Details
Args:
run_ids: the runs (under ``root``'s ``runs/``).
metric: the metric to read — a named metric (``train.metrics``), a
built-in field (``loss``, ``error``, ``accuracy``, ``f1``,
``recall``, ``precision``) or an ``extra_metrics`` name.
split: ``"validation"`` (the epoch's validation record) or
``"train"`` (its whole-epoch training summary — recorded only by
runs that declare metrics or a monitor; never a last batch).
selection: ``"last"`` (the last committed epoch) or ``"best"`` (the
last committed epoch the run's monitor elected; the rule names
the declared monitor, and the value is unknown for a run without
one or whose monitor elected nothing).
replicate: ``"seed"`` keys each observation ``seed=<seed>`` (the
training or ``Trainer`` seed), plus ``init_seed=<seed>`` for a
registered ``ModelSpec`` whose own seed differs; unknown for a
run without a training seed. ``None`` leaves it unknown.
config: run id → declared configuration label, for every run; by
default, a digest of the run's configuration without its salt,
seeds and device, in which a parent run's id is replaced by the
parent's own identity, the checkpoint tag, the resume mode and,
when it is part of the procedure, the epoch the run started at —
so a continuation, fine-tune or later generation pools only
with runs descended the same way from the same configuration
(an unknown lineage pools with nothing).
Runs should be trained with ``provenance=``: the status and
attempt id come from the run's attempt record (``"unknown"`` without
one) and the data and split identities from its manifest. A run's
history is committed by its LAST checkpoint, so without an attempt
record naming that checkpoint's epoch the committed epoch — and so the
value — is unknown. The recorded epoch is used as is (LAST itself is
never opened, so nothing is unpickled), and a history journal is read
only up to it — not at all while it is unknown, so an uncommitted tail
is never read. A legacy run (no commit marker) keeps its value,
but its status is unknown, so it is never counted in a group's ``n``.
A NaN metric is read back as missing: ``idps.csv`` writes NaN as an
empty cell.
nnx.comparison.ComparisonError¶
class nnx.comparison.ComparisonError
A malformed observation, a duplicate attempt, an invalid pairing or a report file that does not hold what it claims.
nnx.comparison.IncompatibleObservations¶
class nnx.comparison.IncompatibleObservations(message: 'str', *, fields: 'Sequence[str]') -> 'None'
Observations that may not be pooled or paired: fields names what differs.
2.28. Causal language-model task (nnx.lm_tasks)¶
nnx.lm_tasks¶
module nnx.lm_tasks
A causal language-model task: declared alignment, masks and valid-token normalization.
Details
A next-token step written by hand flattens ``(B, T, V)`` logits before any
mask exists, divides by the batch count when it accumulates and returns
placeholder classification fields. :class:`CausalLMTask` declares instead:
- **alignment** — ``"shift_inputs"``: a batch is token ids ``(B, T)`` (or
``(ids, loss_mask)``) and the task shifts once, predicting token ``t + 1``
from tokens ``<= t``; ``"pre_shifted"``: a batch is ``(inputs, targets)``
(or ``(inputs, targets, loss_mask)``) already aligned, never shifted again;
- the **vocabulary** size and axis (the logits' last axis), the **ignore**
id (``-100``) and an optional **padding** id, an optional per-position
**loss mask** (a loss mask is never an attention mask), label
**smoothing** for the objective, and the task **version**.
The ignore id, the padding id and the loss mask select one set of **valid
positions** — the objective's denominator, the NLL and the token accuracy
all use exactly those. Shapes, dtypes, lengths and every non-ignored id are
checked before any forward pass.
- :meth:`CausalLMTask.objective` — an objective: one ``"token_ce"``
term per microbatch, the (smoothed) cross-entropy **sum** over valid
positions with the valid count as its denominator, so accumulation
windows normalize by the window's valid tokens. An all-masked microbatch
contributes nothing, and an all-masked window takes no optimizer or
scaler step (its gradients are cleared; schedulers still advance once per
epoch on the default epoch clock, or once per committed update on the
``optimizer_update`` clock, and a plateau scheduler skips an epoch whose validation record is
unavailable). The task configuration is
checkpointed component state (``"lm.causal_task"``): a resume with another
configuration is refused before the first resumed update.
- :meth:`CausalLMTask.eval_step` — an ``eval_step_fn``: the epoch's
**unsmoothed** NLL is the total valid-token negative log likelihood over
the valid-token count (independent of batching and right padding; the
model never receives the attention mask, so a mapping batch whose
``attention_mask`` shows left padding is refused; ``pad_id`` alone does not
detect left padding, so right-pad the batch), perplexity is
``exp(NLL)`` (``+inf`` when it overflows, never a clipped finite value) and
token accuracy is the top-1 rate. They are the record's ``loss`` and its
``metrics`` ``nll`` / ``perplexity`` / ``token_accuracy``; an epoch whose
every position is masked is reported unavailable (no loss, no metrics).
No classification field is fabricated.
nnx.lm_tasks.CausalLMTask¶
class nnx.lm_tasks.CausalLMTask(vocab_size: 'int', alignment: 'str' = 'shift_inputs', ignore_id: 'int' = -100, pad_id: 'Optional[int]' = None, smoothing: 'float' = 0.0, vocab_axis: 'int' = -1, tokenizer: 'Optional[str]' = None, version: 'int' = 1) -> 'None'
A next-token prediction task.
Details
Args:
vocab_size: the vocabulary size; the logits' class axis must have
exactly this width.
alignment: ``"shift_inputs"`` (the task shifts token ids once) or
``"pre_shifted"`` (batches carry aligned inputs and targets).
ignore_id: the target id that marks a position as not scored.
pad_id: an optional padding id: target positions holding it are not
scored either (input positions may hold it).
smoothing: label smoothing for the **objective** only, in ``[0, 1)``;
reported NLL and perplexity are always unsmoothed.
vocab_axis: the logits' vocabulary axis (``-1``, the last).
tokenizer: an optional declared tokenizer identity (for example the
tokenizer file's digest), checkpointed with the task.
version: the task schema version.
nnx.lm_tasks.CausalLMTask.state¶
nnx.lm_tasks.CausalLMTask.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.lm_tasks.CausalLMTask.from_state¶
nnx.lm_tasks.CausalLMTask.from_state(state: 'Mapping[str, Any]') -> 'CausalLMTask'
No public description is currently available.
nnx.lm_tasks.CausalLMTask.split¶
nnx.lm_tasks.CausalLMTask.split(self, batch: 'Any') -> 'tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]'
(inputs, targets, loss_mask) for one batch, aligned once (see the module docstring), with shapes, dtypes and ids checked.
Details
``"shift_inputs"`` takes token ids ``(B, T)``, ``(ids,)``, ``(ids,
loss_mask)`` with a boolean mask, or a mapping ``{"input_ids",
"labels"?, "loss_mask"?}`` — HuggingFace-style ``labels`` sit at the
ids' positions (unshifted, ``ignore_id`` where not scored) and are
shifted with them. ``"pre_shifted"`` takes ``(inputs, targets[,
loss_mask])`` or a mapping ``{"inputs", "targets", "loss_mask"?}``.
A loss mask (boolean, or 0/1 except in the ``(ids, loss_mask)``
tuple) covers the ids (``"shift_inputs"``) or the targets
(``"pre_shifted"``). A mapping's ``attention_mask`` aligns the same
way (the ids, or the targets when pre-shifted). It is never a loss
mask, but a position it pads that would still be scored is refused,
and so is padding before a real token (left or interior padding): the
model never receives the mask. Without one the task cannot detect
left padding (``pad_id`` only masks pad targets), so right-pad the
batch.
nnx.lm_tasks.CausalLMTask.valid¶
nnx.lm_tasks.CausalLMTask.valid(self, targets: 'torch.Tensor', loss_mask: 'Optional[torch.Tensor]') -> 'torch.Tensor'
The scored positions: not ignore_id, not pad_id and, when given, inside the loss mask.
nnx.lm_tasks.CausalLMTask.logits¶
nnx.lm_tasks.CausalLMTask.logits(self, model: 'Any', inputs: 'torch.Tensor', targets: 'torch.Tensor') -> 'torch.Tensor'
The model's logits (B, T, V) for inputs, checked.
nnx.lm_tasks.CausalLMTask.token_sums¶
nnx.lm_tasks.CausalLMTask.token_sums(self, logits: 'torch.Tensor', targets: 'torch.Tensor', valid: 'torch.Tensor', *, with_ce: 'bool' = True) -> 'tuple[torch.Tensor, float, int, int]'
(smoothed CE sum (differentiable), unsmoothed NLL sum, correct, valid count) over the valid positions; flattened once, here. with_ce=False (evaluation) skips the cross-entropy pass and returns an empty sum in its place.
nnx.lm_tasks.CausalLMTask.objective¶
nnx.lm_tasks.CausalLMTask.objective(self, *, nonfinite: 'str' = 'fail') -> 'CausalLMObjective'
No public description is currently available.
nnx.lm_tasks.CausalLMTask.eval_step¶
nnx.lm_tasks.CausalLMTask.eval_step(self) -> 'CausalLMEval'
No public description is currently available.
nnx.lm_tasks.CausalLMObjective¶
class nnx.lm_tasks.CausalLMObjective(task: 'CausalLMTask', *, nonfinite: 'str' = 'fail') -> 'None'
The task's training objective (see CausalLMTask.objective): a "token_ce" term per microbatch — the cross-entropy sum over valid positions (smoothed when the task smooths) over the valid count. Its record's loss is that microbatch's term value and metrics its unsmoothed nll and token_accuracy. Checkpointed as component "lm.causal_task": the task configuration and the valid tokens the run's objective has scored (counted from 0 for each fresh run, restored on a stateful resume; a window the engine then skips — a non-finite loss, an AMP overflow — is still counted). Pass a fresh objective to a weights_only resume: it restores no component state, so a reused instance keeps its own count.
nnx.lm_tasks.CausalLMObjective.component_spec¶
nnx.lm_tasks.CausalLMObjective.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.lm_tasks.CausalLMObjective.component_state¶
nnx.lm_tasks.CausalLMObjective.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.lm_tasks.CausalLMObjective.check_component_state¶
nnx.lm_tasks.CausalLMObjective.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.lm_tasks.CausalLMObjective.load_component_state¶
nnx.lm_tasks.CausalLMObjective.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.lm_tasks.CausalLMEval¶
class nnx.lm_tasks.CausalLMEval(task: 'CausalLMTask') -> 'None'
The task's eval_step_fn (see CausalLMTask.eval_step): total valid-token NLL over the valid-token count for the whole validation loader, its perplexity and token accuracy. The model runs in eval mode and every submodule's training flag is restored.
nnx.lm_tasks.perplexity¶
nnx.lm_tasks.perplexity(nll: 'Optional[float]') -> 'Optional[float]'
exp(nll); +inf when it overflows (never a clipped value).
nnx.lm_tasks.LMTaskError¶
class nnx.lm_tasks.LMTaskError
A task configuration or a batch the task rejects.
2.29. Query-grouped ranking (nnx.ranking)¶
nnx.ranking¶
module nnx.ranking
Query-grouped ranking and retrieval evaluation.
Details
A scorer ranks the **candidates of each query**; it is judged per query,
never per batch. :class:`RankingTask` declares the relevance grades, the
cutoffs ``k``, the relevance threshold, how the pairwise objective weights
queries and pairs, the candidate-set kind and the buffer limit, and checks
every batch — query ids, candidate ids, relevance grades and an optional
mask — before any update.
A batch is ``(features, query_ids, candidate_ids, relevance)`` or
``(..., mask)`` (or a mapping with those keys): the scorer maps
``features`` to one score per candidate row; ids are integer tensors or
lists of strings; relevance grades are integers in ``0..max_relevance``.
- :meth:`RankingTask.objective` — the **pairwise logistic** objective: for each query, every pair of its candidates with unequal
relevance contributes ``log(1 + exp(-(s_hi - s_lo)))``; equal-relevance
pairs (ties) contribute nothing and a query without such a pair is
skipped. ``weighting="query"`` averages each query's pairs and then the
queries (a query weighs the same however many pairs it has);
``weighting="pair"`` weighs every pair equally. Pairs are formed within a
microbatch. Checkpointed as component ``"ranking.task"``.
- :meth:`RankingTask.eval_step` — an ``eval_step_fn``: it buffers the whole
validation stream (at most ``max_buffered`` candidate rows), groups it by
query — so a query split across batches is joined, in any order — and
rejects a candidate seen twice for one query. Each query's candidates are
ranked by score, **ties broken by candidate id** (ascending); then
``MRR@k`` (the reciprocal rank of the first relevant candidate in the top
``k``, else 0), ``Recall@k`` (relevant candidates in the top ``k`` over
all relevant) and ``NDCG@k`` (gain ``2^rel - 1``, discount
``log2(rank + 1)``) are averaged over queries. A candidate is relevant
when its grade is at least ``relevance_threshold``; a query with no
relevant candidate is **excluded** (and counted), never scored 0. ``k``
above a query's candidate count is rejected.
The record (``kind="ranking"``) carries ``mrr_at_<k>``, ``recall_at_<k>``,
``ndcg_at_<k>``, the counts of scored and excluded queries and candidates,
and ``exhaustive_candidates`` (``1.0`` when every query's candidate set is
its full corpus, ``0.0`` when it is sampled — then the metrics measure the
candidate set, not full-corpus retrieval). :meth:`RankingTask.metric_specs`
declares the metrics so a :class:`~nnx.MonitorSpec` can name them (for
example ``MonitorSpec("ndcg_at_10")``) for BEST, early stopping and plateau
scheduling.
nnx.ranking.RankingTask¶
class nnx.ranking.RankingTask(k: 'tuple[int, ...]' = (10,), max_relevance: 'int' = 1, relevance_threshold: 'int' = 1, weighting: 'str' = 'query', candidate_sets: 'str' = 'sampled', max_buffered: 'int' = 1000000, group_size: 'Optional[int]' = None, version: 'int' = 1) -> 'None'
A query-grouped ranking task.
Details
Args:
k: the cutoffs the metrics are reported at (each ``>= 1``).
max_relevance: grades are integers in ``0..max_relevance``
(``1``: binary relevance).
relevance_threshold: a candidate is relevant (for MRR and Recall,
and for excluding queries) when its grade is at least this.
weighting: the objective's ``"query"`` or ``"pair"`` weighting.
candidate_sets: ``"sampled"`` (each query sees a sample of its
corpus) or ``"exhaustive"`` (its whole corpus); recorded with
the metrics.
max_buffered: the most candidate rows the evaluation buffers to
join each query's candidates.
group_size: when set, every evaluated query must have exactly this
many candidates — an incomplete (or over-full) query is
rejected rather than scored.
version: the task schema version.
nnx.ranking.RankingTask.state¶
nnx.ranking.RankingTask.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.ranking.RankingTask.from_state¶
nnx.ranking.RankingTask.from_state(state: 'Mapping[str, Any]') -> 'RankingTask'
No public description is currently available.
nnx.ranking.RankingTask.metric_names¶
nnx.ranking.RankingTask.metric_names(self) -> 'list[str]'
No public description is currently available.
nnx.ranking.RankingTask.metric_specs¶
nnx.ranking.RankingTask.metric_specs(self) -> 'list[Any]'
MetricSpecs declaring this task's metrics — registered ids ranking.mrr / ranking.recall / ranking.ndcg with config={"k": k}, reported as <metric>_at_<k> — so a MonitorSpec can select on them. eval_step computes them; a declaration under another name finds no value.
nnx.ranking.RankingTask.split¶
nnx.ranking.RankingTask.split(self, batch: 'Any') -> 'tuple[Any, list[Id], list[Id], torch.Tensor, Optional[torch.Tensor]]'
(features, query_ids, candidate_ids, relevance, mask), checked.
nnx.ranking.RankingTask.scores¶
nnx.ranking.RankingTask.scores(self, model: 'Any', features: 'Any', n: 'int') -> 'torch.Tensor'
The scorer's one score per row, checked.
nnx.ranking.RankingTask.objective¶
nnx.ranking.RankingTask.objective(self, *, nonfinite: 'str' = 'fail') -> 'RankingObjective'
No public description is currently available.
nnx.ranking.RankingTask.eval_step¶
nnx.ranking.RankingTask.eval_step(self) -> 'RankingEval'
No public description is currently available.
nnx.ranking.RankingTask.evaluate_rows¶
nnx.ranking.RankingTask.evaluate_rows(self, rows: 'Iterable[tuple[Id, Id, float, int]]') -> 'tuple[dict[str, float], int, int, int]'
Metrics over complete queries from (query, candidate, score, relevance) rows in any order: (means, scored queries, excluded queries, candidates of the scored queries). Duplicate candidates of a query raise.
nnx.ranking.RankingObjective¶
class nnx.ranking.RankingObjective(task: 'RankingTask', *, nonfinite: 'str' = 'fail') -> 'None'
The pairwise logistic objective of a RankingTask (see the module docstring); checkpointed as component "ranking.task" with the task configuration and the run's training units — pairs (weighting="pair") or queries ("query") scored, from 0 for each fresh run and restored on resume (a window the engine then skips is still counted).
nnx.ranking.RankingObjective.component_spec¶
nnx.ranking.RankingObjective.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.ranking.RankingObjective.component_state¶
nnx.ranking.RankingObjective.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.ranking.RankingObjective.check_component_state¶
nnx.ranking.RankingObjective.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.ranking.RankingObjective.load_component_state¶
nnx.ranking.RankingObjective.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.ranking.RankingEval¶
class nnx.ranking.RankingEval(task: 'RankingTask') -> 'None'
The RankingTask's eval_step_fn: buffers the validation stream (at most max_buffered rows), groups it by query and reports the per-query metrics (see the module docstring), with the pairwise loss over the same queries as the record's loss.
nnx.ranking.pairwise_logistic_loss¶
nnx.ranking.pairwise_logistic_loss(scores: 'torch.Tensor', relevance: 'torch.Tensor', query_ids: 'Sequence[Id]', *, weighting: 'str' = 'query') -> 'tuple[torch.Tensor, int]'
(numerator, denominator) of the pairwise logistic loss over the unequal-relevance pairs within each query (see the module docstring).
Details
Pairs are formed grade by grade (each grade's candidates against every
lower-graded one) in blocks of at most :data:`PAIR_BLOCK` pairs, so no
query's full pair matrix is ever held — in training too: a block of at
least :data:`CHECKPOINT_PAIRS` pairs is recomputed in backward
(``torch.utils.checkpoint``) rather than kept, while a smaller one costs
less to keep than to recompute. ``relevance`` may live on any device; it
is moved to the scores'.
nnx.ranking.rank¶
nnx.ranking.rank(scores: 'Sequence[float]', candidate_ids: 'Sequence[Id]') -> 'list[int]'
Positions of the candidates in ranked order: by score, highest first, ties broken by candidate id (ascending) — a stable, documented order. A non-finite score has no place in that order and is refused.
nnx.ranking.mrr_at_k¶
nnx.ranking.mrr_at_k(relevance: 'Sequence[int]', scores: 'Sequence[float]', candidate_ids: 'Sequence[Id]', k: 'int', *, threshold: 'int' = 1) -> 'float'
Reciprocal rank of the first relevant candidate within the top k (0 when none is).
nnx.ranking.recall_at_k¶
nnx.ranking.recall_at_k(relevance: 'Sequence[int]', scores: 'Sequence[float]', candidate_ids: 'Sequence[Id]', k: 'int', *, threshold: 'int' = 1) -> 'float'
Relevant candidates in the top k over all relevant candidates (the query must have at least one).
nnx.ranking.ndcg_at_k¶
nnx.ranking.ndcg_at_k(relevance: 'Sequence[int]', scores: 'Sequence[float]', candidate_ids: 'Sequence[Id]', k: 'int') -> 'float'
DCG@k / IDCG@k with gain 2^rel - 1 and discount log2(rank + 1) (the query must have a nonzero grade).
nnx.ranking.RankingError¶
class nnx.ranking.RankingError
A configuration, batch or query the ranking task rejects.
2.30. Graph-level classification (nnx.graph_tasks)¶
nnx.graph_tasks.GraphCollection¶
class nnx.graph_tasks.GraphCollection(graphs: 'Sequence[Any]', ids: 'Sequence[int]', *, targets: 'Optional[Sequence[Optional[int]]]' = None, unlabeled: 'Sequence[int]' = (), num_classes: 'Optional[int]' = None) -> 'None'
Whole graphs with stable ids and graph-level targets.
Details
Args:
graphs: PyG ``Data`` objects with ``x`` (``(nodes, features)``,
floating) and ``edge_index`` (``(2, edges)``, node indices of the
same graph); each graph's target is its ``y`` (one integer
class) unless ``targets`` gives them.
ids: one stable non-negative integer id per graph (unique).
targets: optional class per graph, overriding ``y``.
unlabeled: ids of graphs deliberately without a target (their
target becomes :data:`IGNORE`); any other graph without one is
refused.
num_classes: when given, every target must be below it.
nnx.graph_tasks.GraphCollection.labeled¶
property nnx.graph_tasks.GraphCollection.labeled
How many graphs carry a target.
nnx.graph_tasks.GraphCollection.subset¶
nnx.graph_tasks.GraphCollection.subset(self, ids: 'Sequence[int]') -> 'GraphCollection'
The graphs with these ids, in the order given.
nnx.graph_tasks.GraphCollection.loader¶
nnx.graph_tasks.GraphCollection.loader(self, batch_size: 'int', *, shuffle: 'bool' = False, seed: 'Optional[int]' = None) -> 'Any'
A PyG loader over whole graphs (in collection order unless shuffle; a seed makes the shuffle reproducible and is refused without it).
nnx.graph_tasks.GraphClassifier¶
class nnx.graph_tasks.GraphClassifier(encoder: 'nn.Module', pool: 'Any' = 'mean', head: 'Optional[nn.Module]' = None, *, input_dim: 'Optional[int]' = None) -> 'None'
encoder (node rows from (x, edge_index)) → GraphPool → optional head: one output row per graph.
Details
``unpack_batch`` validates a graph-collection batch (see
:func:`check_graph_batch`) and returns its graph-level targets;
``sample_ids`` returns its graph ids, which ``predict_proba`` reports as
the rows' sample ids.
nnx.graph_tasks.GraphClassifier.forward¶
nnx.graph_tasks.GraphClassifier.forward(self, x: 'torch.Tensor', edge_index: 'torch.Tensor', batch: 'torch.Tensor', ptr: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.graph_tasks.GraphClassifier.unpack_batch¶
nnx.graph_tasks.GraphClassifier.unpack_batch(self, batch: 'Any') -> 'tuple[tuple[torch.Tensor, ...], Optional[torch.Tensor]]'
No public description is currently available.
nnx.graph_tasks.GraphClassifier.sample_ids¶
nnx.graph_tasks.GraphClassifier.sample_ids(self, batch: 'Any') -> 'torch.Tensor'
The batch's graph ids, one per output row, in batch order.
nnx.graph_tasks.GraphPool¶
class nnx.graph_tasks.GraphPool(mode: 'str' = 'mean') -> 'None'
Pools node rows into one row per graph: "mean" or "sum" over each graph's own nodes (graph-local, invariant to node order).
nnx.graph_tasks.GraphPool.forward¶
nnx.graph_tasks.GraphPool.forward(self, x: 'torch.Tensor', batch: 'torch.Tensor', num_graphs: 'int') -> 'torch.Tensor'
No public description is currently available.
nnx.graph_tasks.GraphPool.extra_repr¶
nnx.graph_tasks.GraphPool.extra_repr(self) -> 'str'
No public description is currently available.
nnx.graph_tasks.graph_classifier_spec¶
nnx.graph_tasks.graph_classifier_spec(*, input_dim: 'int', num_classes: 'int', hidden_dims: 'Sequence[int]' = (32,), encoder: 'str' = 'graph_conv', pool: 'str' = 'mean', activation: 'str' = 'relu', dropout: 'float' = 0.0, seed: 'int' = 0) -> 'Any'
The registered recipe of a GraphClassifier: encoder convolutions ("graph_conv" GCN, "graph_sage", "graph_att" GAT) over hidden_dims, pool and a linear head to num_classes. Pass it as NNModelParams(net=...); a reload rebuilds the same module from the run.
nnx.graph_tasks.check_graph_batch¶
nnx.graph_tasks.check_graph_batch(batch: 'Any') -> 'tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]'
(x, edge_index, batch_vector, ptr) of a graph-collection batch, checked: graph ids (unique), at least one node per graph, ptr and the batch vector consistent, and no edge across graphs. An edge-label (link) batch is refused: pooling it would score candidate edges as graphs.
nnx.graph_tasks.is_graph_collection_batch¶
nnx.graph_tasks.is_graph_collection_batch(batch: 'Any') -> 'bool'
Whether batch is a batch of whole graphs (it carries graph ids).
nnx.graph_tasks.GraphTaskError¶
class nnx.graph_tasks.GraphTaskError
A graph collection, batch or classifier setting NNx rejects.
2.31. Link and edge prediction (nnx.link_tasks)¶
nnx.link_tasks¶
module nnx.link_tasks
Leakage-aware link and edge prediction on a static homogeneous graph.
Details
Node masks cannot hold out edges: a held-out edge left in the message graph
is a leaked answer. This module splits **edges**, keeps the graph the model
passes messages over apart from the edges it is asked about, and checks
every batch against that split before any forward pass.
- :class:`LinkSplit` — a versioned manifest: the node count, canonical
positive edges per split (``train`` / ``val`` / ``test``; an undirected
edge is ``(min, max)``, so a duplicate, or in an undirected split a
reverse, can never sit in two splits), fixed ``val`` / ``test`` negatives, directedness, the self-loop
policy, the seed, and — for edge-label tasks — the categories, each
edge's category and whether labelled edges' **existence** is visible
context. :func:`split_links` derives one from an ``edge_index`` with a
local seeded generator (never the global RNG); :meth:`LinkSplit.replay`
re-derives it and refuses a mismatch.
- **Topology.** :meth:`LinkSplit.message_edge_index` is the only graph
messages pass over: the training edges (and their reverses when
undirected) — for every split, so evaluation never sees held-out
positives. In edge-label mode with ``label_existence="visible"`` every
labelled edge's existence is context (their categories never are).
- **Negatives** come from the static complement — never a positive of any
split, its reverse (in an undirected split), a duplicate or a barred
self-loop. A request beyond
the complement's capacity fails before anything is sampled; ``val`` /
``test`` negatives are fixed in the manifest; training negatives are
re-drawn every pass from ``(seed, pass)``.
- :class:`LinkTask` — ``"binary"`` (link existence: one logit per
candidate, targets 0/1) or ``"edge_label"`` (categories in their declared
order; no negatives, so a non-edge is never a category). ``loader()``
builds batches of candidates over the message graph; ``objective()``,
``eval_step()`` and ``predict()`` check each batch against the split —
a message edge outside the training topology (a hidden positive) or a
candidate outside its split fails before any update. Evaluation
materialises every candidate (up to ``max_candidates``) for exact
AUROC / AP, scored from the float64 logits (a probability rounds to 0 or
1 for a confident model). For a one-class set, ``link_metrics`` reports
them unavailable with the reason, and the evaluator's record omits them.
The manifest is checkpointed component state ``"link.task"``: a resume
with another split fails before the first resumed update.
- :func:`link_predictor_spec` — a registered recipe
(``ModelSpec("nnx.link_predictor", 1, config)``): a GCN / GraphSAGE /
GAT encoder and a ``"dot"`` or ``"mlp"`` edge decoder, rebuilt on reload.
Homogeneous static graphs only: no temporal or heterogeneous graphs, and
unobserved pairs are negatives only as the split declares them.
nnx.link_tasks.split_links¶
nnx.link_tasks.split_links(edge_index: 'Any', num_nodes: 'int', *, val: 'Any' = 0.1, test: 'Any' = 0.1, seed: 'int', directed: 'bool' = False, self_loops: 'str' = 'bar', negatives: 'int' = 1, edge_labels: 'Optional[Sequence[int]]' = None, categories: 'Optional[Sequence[str]]' = None, label_existence: 'str' = 'hidden') -> 'LinkSplit'
Split the graph's edges with a local generator seeded by seed.
Details
Edges are canonicalised and de-duplicated first (an undirected reverse
is the same edge), so no duplicate, and in an undirected split no
reverse, can cross splits (a directed split keeps ``u→v`` and ``v→u``
apart).
``val`` / ``test`` are fractions of the edges or exact counts.
``negatives`` fixed non-edges per ``val`` / ``test`` positive come from
the complement of every positive (and barred self-loops); a request
beyond the complement's capacity fails before sampling. With
``categories`` (and one category per input edge in ``edge_labels``) the
split is an edge-label split: no negatives, and a duplicate edge must
keep its category.
nnx.link_tasks.LinkSplit¶
class nnx.link_tasks.LinkSplit(num_nodes: 'int', train: 'tuple[Edge, ...]', val: 'tuple[Edge, ...]' = (), test: 'tuple[Edge, ...]' = (), val_negatives: 'tuple[Edge, ...]' = (), test_negatives: 'tuple[Edge, ...]' = (), directed: 'bool' = False, self_loops: 'str' = 'bar', seed: 'Optional[int]' = None, categories: 'Optional[tuple[str, ...]]' = None, edge_labels: 'tuple[tuple[Edge, int], ...]' = (), label_existence: 'str' = 'hidden', version: 'int' = 1) -> 'None'
A versioned edge split (see the module docstring). Every edge is canonical; validation refuses an edge outside the graph, a barred self-loop, a positive in two splits (or twice in one), and a negative that is a positive or repeats.
nnx.link_tasks.LinkSplit.mode¶
property nnx.link_tasks.LinkSplit.mode
No public description is currently available.
nnx.link_tasks.LinkSplit.positives¶
nnx.link_tasks.LinkSplit.positives(self, name: 'str') -> 'tuple[Edge, ...]'
No public description is currently available.
nnx.link_tasks.LinkSplit.negatives¶
nnx.link_tasks.LinkSplit.negatives(self, name: 'str') -> 'tuple[Edge, ...]'
No public description is currently available.
nnx.link_tasks.LinkSplit.candidate_ids¶
nnx.link_tasks.LinkSplit.candidate_ids(self, name: 'str') -> 'dict[Edge, int]'
Stable integer ids of a split's fixed candidates (positives, then negatives): one global numbering over the manifest.
nnx.link_tasks.LinkSplit.labels¶
nnx.link_tasks.LinkSplit.labels(self) -> 'dict[Edge, int]'
No public description is currently available.
nnx.link_tasks.LinkSplit.message_edge_index¶
nnx.link_tasks.LinkSplit.message_edge_index(self) -> 'torch.Tensor'
The only edges messages pass over, for every split: the training edges (with reverses when undirected) — plus every labelled edge's existence in an edge-label split with label_existence="visible".
nnx.link_tasks.LinkSplit.state¶
nnx.link_tasks.LinkSplit.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.link_tasks.LinkSplit.from_state¶
nnx.link_tasks.LinkSplit.from_state(state: 'Mapping[str, Any]') -> 'LinkSplit'
No public description is currently available.
nnx.link_tasks.LinkSplit.digest¶
nnx.link_tasks.LinkSplit.digest(self) -> 'str'
No public description is currently available.
nnx.link_tasks.LinkSplit.replay¶
nnx.link_tasks.LinkSplit.replay(self, edge_index: 'Any', *, edge_labels: 'Optional[Sequence[int]]' = None) -> 'LinkSplit'
Re-derive this split from the graph with its own seed and policy; a mismatch (different graph, settings or code) raises.
nnx.link_tasks.LinkTask¶
class nnx.link_tasks.LinkTask(split: 'LinkSplit', mode: 'Optional[str]' = None, train_negatives: 'Optional[int]' = None, max_candidates: 'int' = 1000000, version: 'int' = 1) -> 'None'
Link existence (mode="binary") or edge categories ("edge_label") over a LinkSplit (see the module docstring).
Details
Args:
split: the manifest.
mode: ``"binary"`` or ``"edge_label"``; by default the split's (an
edge-label split has categories).
train_negatives: non-edges sampled per training positive, per pass
(binary only; by default 1, and 0 in edge-label mode).
max_candidates: the most candidates evaluation materialises.
Training negatives are drawn per epoch from ``(seed, epoch)``, so a
stateful resume continues them as if uninterrupted — when the training
loop's ``set_epoch`` reaches the task's loader (directly, or through a
wrapper that forwards it). The objective checkpoints the seed and
refuses a resumed loader with another one.
nnx.link_tasks.LinkTask.state¶
nnx.link_tasks.LinkTask.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.link_tasks.LinkTask.loader¶
nnx.link_tasks.LinkTask.loader(self, name: 'str', x: 'torch.Tensor', batch_size: 'int', *, seed: 'Optional[int]' = None) -> '_Loader'
Batches of name's candidates over the message graph. Training negatives and order are re-drawn every pass from (seed, pass) (seed, an integer >= 0, defaults to the split's); a seed on a val / test loader, whose candidates are fixed, is refused.
nnx.link_tasks.LinkTask.check_batch¶
nnx.link_tasks.LinkTask.check_batch(self, batch: 'Any') -> 'None'
Refuse a batch that would leak or does not fit the split: message edges outside the training topology (a held-out positive is named); candidates outside the batch's split, with the wrong target or id; a barred self-loop; and a training negative that is a held-out (val / test) negative.
nnx.link_tasks.LinkTask.objective¶
nnx.link_tasks.LinkTask.objective(self, *, nonfinite: 'str' = 'fail') -> 'LinkObjective'
No public description is currently available.
nnx.link_tasks.LinkTask.eval_step¶
nnx.link_tasks.LinkTask.eval_step(self) -> 'LinkEval'
No public description is currently available.
nnx.link_tasks.LinkTask.metric_specs¶
nnx.link_tasks.LinkTask.metric_specs(self) -> 'list[Any]'
MetricSpecs for MonitorSpec("auroc") / ("ap") (binary) — a one-class set leaves them unavailable, so it never wins BEST.
nnx.link_tasks.LinkTask.predict¶
nnx.link_tasks.LinkTask.predict(self, model: 'Any', batches: 'Any') -> 'LinkPrediction'
Every candidate's prediction, in batch order, aligned with its id, pair and target (eval mode, no gradients, modes restored).
nnx.link_tasks.LinkPredictor¶
class nnx.link_tasks.LinkPredictor(encoder: 'nn.Module', decoder: 'str' = 'dot', *, width: 'int' = 0, outputs: 'int' = 1) -> 'None'
encoder (node rows from (x, edge_index)) and an edge decoder: "dot" (one logit per candidate: h_u · h_v) or "mlp" (a linear head on [h_u * h_v, |h_u - h_v|]: one logit, or one per category).
nnx.link_tasks.LinkPredictor.forward¶
nnx.link_tasks.LinkPredictor.forward(self, x: 'torch.Tensor', edge_index: 'torch.Tensor', edge_label_index: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.link_tasks.LinkPredictor.unpack_batch¶
nnx.link_tasks.LinkPredictor.unpack_batch(self, batch: 'Any') -> 'tuple[tuple[torch.Tensor, ...], Optional[torch.Tensor]]'
Refused: the default train / evaluate / predict paths cannot check a batch against its split, so a leak would pass unseen.
nnx.link_tasks.LinkPredictor.sample_ids¶
nnx.link_tasks.LinkPredictor.sample_ids(self, batch: 'Any') -> 'torch.Tensor'
The candidates' ids, one per output row.
nnx.link_tasks.link_predictor_spec¶
nnx.link_tasks.link_predictor_spec(*, input_dim: 'int', hidden_dims: 'Sequence[int]' = (32,), encoder: 'str' = 'graph_conv', decoder: 'str' = 'dot', num_categories: 'Optional[int]' = None, activation: 'str' = 'relu', dropout: 'float' = 0.0, seed: 'int' = 0) -> 'Any'
The registered recipe of a LinkPredictor (pass it as NNModelParams(net=...); a reload rebuilds the same module).
nnx.link_tasks.LinkObjective¶
class nnx.link_tasks.LinkObjective(task: 'LinkTask', *, nonfinite: 'str' = 'fail') -> 'None'
The task's training objective: binary cross-entropy (or categorical cross-entropy) summed over a batch's candidates, over the candidate count, after LinkTask.check_batch. Checkpointed as component "link.task" with the manifest (topology policy, candidate ids, fixed negatives).
nnx.link_tasks.LinkObjective.component_spec¶
nnx.link_tasks.LinkObjective.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.link_tasks.LinkObjective.component_state¶
nnx.link_tasks.LinkObjective.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.link_tasks.LinkObjective.check_component_state¶
nnx.link_tasks.LinkObjective.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.link_tasks.LinkObjective.load_component_state¶
nnx.link_tasks.LinkObjective.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
The manifest is configuration (checked, not restored). The training negatives' seed is restored and must match the resumed loader's. The passes need no state: the training loop announces each epoch to the task's loader (set_epoch), so a resumed run draws what the uninterrupted one would have — through any wrapper that forwards set_epoch.
nnx.link_tasks.LinkEval¶
class nnx.link_tasks.LinkEval(task: 'LinkTask') -> 'None'
The task's eval_step_fn: every validation candidate materialised (at most max_candidates), then exact metrics over the whole set — never per-batch averages — from the float64 logits, never rounded probabilities (link_metrics with from_logits=True). The record (kind="link") counts the candidates; its loss is the BCE (or NLL).
nnx.link_tasks.MetricValue¶
class nnx.link_tasks.MetricValue(value: 'Optional[float]', reason: 'Optional[str]' = None) -> 'None'
A metric's value, or None with the reason it is unavailable.
nnx.link_tasks.is_link_batch¶
nnx.link_tasks.is_link_batch(batch: 'Any') -> 'bool'
Whether batch asks about candidate edges (it has edge_label_index).
nnx.link_tasks.link_metrics¶
nnx.link_tasks.link_metrics(probabilities: 'Any', targets: 'Any', *, from_logits: 'bool' = False) -> 'dict[str, MetricValue]'
Exact metrics over a whole candidate set. Binary (probabilities of shape (K,)): bce, auroc, ap (unavailable, with the reason, for a one-class set) and accuracy at 0.5. Categorical ((K, C)): nll and accuracy.
Details
``from_logits=True`` reads raw logits instead (as
:meth:`LinkTask.predict` gives them) and never forms probabilities: in
float64, the BCE comes from the logits, the NLL from their log-softmax,
and AUROC / AP rank the logits themselves. A probability rounds to
exactly 0 or 1 for a confident model (beyond ``|logit|`` ≈ 16.6 in
float32, 36.7 in float64), tying candidates the logits tell apart and
clipping their loss; logits stay exact at any confidence. The task's
evaluation scores the logits.
nnx.link_tasks.LinkPrediction¶
class nnx.link_tasks.LinkPrediction(ids: 'np.ndarray', pairs: 'np.ndarray', logits: 'np.ndarray', probabilities: 'np.ndarray', targets: 'np.ndarray', split: 'str') -> 'None'
Candidates and their predictions, aligned row for row: ids, pairs (K, 2), logits / probabilities (float64; (K,) binary, (K, C) categorical) and targets — all of one split. Score it with link_metrics(logits, targets, from_logits=True): a confident model's probabilities round to 0 or 1.
nnx.link_tasks.LinkTaskError¶
class nnx.link_tasks.LinkTaskError
A split, batch or setting that would leak or cannot be served.
3. Params¶
nnx.nn.params.nn_params.NNParams¶
class nnx.nn.params.nn_params.NNParams(*, dropout_prob: 'float', n_heads: 'Optional[int]' = None, activation: 'Optional[Activations]' = leaky_relu, activations: 'Optional[list[Activations]]' = None, dropout_probs: 'Optional[list[float]]' = None, input_dim: 'int', output_dim: 'int', hidden_dims: 'Optional[list[int]]' = None) -> 'None'
Architecture of a built-in feed-forward or graph network: input and output sizes, hidden layer widths, activation and dropout (net-wide or per layer), and attention heads (n_heads, GRAPH_ATT only). Subclasses add the convolutional, MoE, ViT and transformer fields.
nnx.nn.params.nn_params.NNParams.dims¶
property nnx.nn.params.nn_params.NNParams.dims
No public description is currently available.
nnx.nn.params.nn_params.NNParams.activation_for¶
nnx.nn.params.nn_params.NNParams.activation_for(self, layer_idx: 'int') -> 'Activations'
The activation for hidden layer layer_idx — the per-layer entry when activations is set, else the net-wide scalar.
nnx.nn.params.nn_params.NNParams.dropout_for¶
nnx.nn.params.nn_params.NNParams.dropout_for(self, layer_idx: 'int') -> 'float'
The dropout prob for hidden layer layer_idx — per-layer entry when dropout_probs is set, else the net-wide scalar.
nnx.nn.params.nn_params.NNParams.state¶
nnx.nn.params.nn_params.NNParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_params.NNParams.from_state¶
nnx.nn.params.nn_params.NNParams.from_state(state: 'dict') -> 'NNParams'
No public description is currently available.
nnx.nn.params.nn_params.NNParams.resolve_from_state¶
nnx.nn.params.nn_params.NNParams.resolve_from_state(state: 'dict') -> 'NNParams'
Dispatch to the params subclass that wrote state.
Details
``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.
nnx.nn.params.nn_model_params.NNModelParams¶
class nnx.nn.params.nn_model_params.NNModelParams(*, net: 'Optional[Union[Nets, ModelSpec, RuntimeModule]]' = None, device: 'Devices' = cpu, loss: 'Losses' = cross_entropy, mixed_precision: 'bool' = False, task: 'Optional[TaskSpec]' = None, precision: 'Optional[PrecisionPolicy]' = None) -> 'None'
How an NNModel is built and run: the network (a Nets member, a registered ModelSpec or a runtime module), the device, the loss and the task.
nnx.nn.params.nn_model_params.NNModelParams.is_valid¶
nnx.nn.params.nn_model_params.NNModelParams.is_valid(self) -> 'bool'
No public description is currently available.
nnx.nn.params.nn_model_params.NNModelParams.builtin¶
property nnx.nn.params.nn_model_params.NNModelParams.builtin
Whether the net is a built-in Nets member.
nnx.nn.params.nn_model_params.NNModelParams.state¶
nnx.nn.params.nn_model_params.NNModelParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_model_params.NNModelParams.from_state¶
nnx.nn.params.nn_model_params.NNModelParams.from_state(state: 'dict') -> 'NNModelParams'
No public description is currently available.
nnx.nn.params.nn_train_params.NNTrainParams¶
class nnx.nn.params.nn_train_params.NNTrainParams(*, n_epochs: 'int', scheduler: 'NNSchedulerParams' = NNSchedulerParams(min_lr=1e-07, factor=0.95, patience=8, cooldown=2, threshold=0.001, kind=None, step_size=None, T_max=None, max_lr=None, total_steps=None, warmup_steps=None, clock='epoch'), optim: 'Union[NNOptimParams, NNOptimFactoryParams]' = NNOptimParams(name=adam, max_lr=0.01, weight_decay=5e-05, momentum=(0.9, 0.999), grad_clip_norm=None, accumulate_grad_batches=1, param_groups=None, eps=1e-08), seed: 'Optional[int]' = None, data_id: 'Optional[str]' = None, save_phase_checkpoints: 'bool' = True, train_loader: 'Optional[Iterable[Any]]' = None, val_loader: 'Optional[Iterable[Any]]' = None, metrics: 'tuple[MetricSpec, ...]' = (), monitor: 'Optional[MonitorSpec]' = None, extra_metrics: 'Optional[Mapping[str, Callable]]' = None, resume_from_run_id: 'Optional[str]' = None, resume_from_checkpoint: 'Optional[str]' = 'last', parent_run_id: 'Optional[str]' = None, overwrite_existing: 'bool' = False, resume_mode: 'str' = 'auto', resume_epochs: 'str' = 'additional') -> 'None'
Training configuration.
Details
`seed` reseeds the global RNG streams (Python random, NumPy, torch
CPU+CUDA, cuDNN) when NNModel.train() runs. None disables seeding
(default). A loader or sampler with its own `torch.Generator` is not
reseeded: rebuild it (or reseed its generator) before each fit for
the same shuffle order.
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`.
`extra_metrics` maps a name to ``callable(y_true, y_pred) -> float``
(truth first, decoded class predictions second). The default
classification training step calls each per batch, and `evaluate()`
(the default validation pass) calls each once on the aggregate
predictions; a custom `train_step_fn` / `eval_step_fn` decides whether
and how to call them. Runtime-only: not part of `state()`.
`metrics` declares named, registered metrics (:class:`~nnx.MetricSpec`)
computed over the full sample of each epoch — on the validation set by
`evaluate()` and on the training epoch by the default step — and
reported in the records' `metrics` under each spec's name. `monitor`
(:class:`~nnx.MonitorSpec`) names the split and metric that BEST
selection and a `ReduceLROnPlateau` scheduler track, with one shared
improvement rule. Both are serialized only when set, so a
configuration without them keeps its `state()` and run id.
nnx.nn.params.nn_train_params.NNTrainParams.with_train_loader¶
nnx.nn.params.nn_train_params.NNTrainParams.with_train_loader(self, value: 'Iterable[Any]') -> 'NNTrainParams'
No public description is currently available.
nnx.nn.params.nn_train_params.NNTrainParams.with_val_loader¶
nnx.nn.params.nn_train_params.NNTrainParams.with_val_loader(self, value: 'Iterable[Any]') -> 'NNTrainParams'
No public description is currently available.
nnx.nn.params.nn_train_params.NNTrainParams.state¶
nnx.nn.params.nn_train_params.NNTrainParams.state(self)
No public description is currently available.
nnx.nn.params.nn_train_params.NNTrainParams.from_state¶
nnx.nn.params.nn_train_params.NNTrainParams.from_state(state: 'dict') -> 'NNTrainParams'
No public description is currently available.
nnx.nn.params.nn_optim_params.NNOptimParams¶
class nnx.nn.params.nn_optim_params.NNOptimParams(*, 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]]' = None, eps: 'float' = 1e-08) -> 'None'
Optimizer config.
Details
`momentum` is overloaded by optimizer kind:
- For SGD / SGD_NESTEROV: a single float, the SGD momentum coefficient.
- For ADAM / ADAM_AMSGRAD / ADAMW: a (beta1, beta2) tuple, passed as
the `betas=` argument. The name is retained for backwards
compatibility — `is_valid()` enforces the per-optim shape.
`weight_decay` follows the chosen optimizer: ADAM / ADAM_AMSGRAD add it
to the gradient (coupled L2 penalty, rescaled by the adaptive
denominator), while ADAMW decays the weights directly
(``p -= lr * weight_decay * p``) before the Adam update — the
decoupled decay of Loshchilov & Hutter.
`eps` is the Adam-family denominator term (``1e-8``, torch's default).
It is only meaningful for ADAM / ADAM_AMSGRAD / ADAMW; setting it on
an SGD variant raises. It is omitted from `state()` at its default so
existing run ids are unchanged.
`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. Each update
averages the loss over every valid sample in the window (sum of
numerators over sum of denominators), so uneven or partial windows give
the full-batch gradient.
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.
nnx.nn.params.nn_optim_params.NNOptimParams.state¶
nnx.nn.params.nn_optim_params.NNOptimParams.state(self)
No public description is currently available.
nnx.nn.params.nn_optim_params.NNOptimParams.from_state¶
nnx.nn.params.nn_optim_params.NNOptimParams.from_state(state: 'dict') -> 'NNOptimParams'
No public description is currently available.
nnx.nn.params.nn_optim_params.NNOptimParams.is_valid¶
nnx.nn.params.nn_optim_params.NNOptimParams.is_valid(self) -> 'bool'
No public description is currently available.
nnx.nn.params.nn_optim_params.NNOptimParams.builder¶
nnx.nn.params.nn_optim_params.NNOptimParams.builder() -> 'NNOptimParamsBuilder'
Return a variant-aware builder. See NNOptimParamsBuilder.
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder¶
class nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder() -> 'None'
Variant-aware builder for NNOptimParams.
Details
Reach via `NNOptimParams.builder()`. Pick exactly one variant
method (`adam`, `adam_amsgrad`, `adamw`, `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.
`copy()` branches a (possibly partial) builder and `from_params()`
rebuilds one from an existing `NNOptimParams`, so shared setup is
written once and varied per branch.
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.copy¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.copy(self) -> 'NNOptimParamsBuilder'
Return an independent branch of this builder, complete or partial.
Details
The branch starts with the same fields; afterwards setters on
either builder never affect the other, and neither touches values
already built. Configuration containers (the `param_groups` list)
are copied; the immutable `NNParamGroupSpec` rows are shared.
Nothing is built or validated here — `build()` validates each
branch on its own.
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.from_params¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.from_params(params: 'NNOptimParams') -> 'NNOptimParamsBuilder'
Return a builder pre-loaded with every field of params.
Details
`from_params(params).build()` equals `params`, with the same
`state()` (key order and omitted defaults included). The builder
is ordinary afterwards: a new variant call replaces the variant
fields (`name` / `max_lr` / `momentum` / `weight_decay` / `eps`)
and keeps the modifiers, exactly as in a hand-written chain.
Raises:
TypeError: if `params` is not exactly an `NNOptimParams` — a
registered-factory `NNOptimFactoryParams` has no builder
(pass it straight to `NNTrainerParamsBuilder.optimizer`),
and subclasses are rejected rather than downgraded.
ValueError: if `params` carries a field this builder cannot
reproduce.
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.adam¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.adam(self, *, max_lr: 'float', betas: 'tuple[float, float]' = (0.9, 0.999), weight_decay: 'float' = 0.0, eps: 'Optional[float]' = None) -> '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). weight_decay is Adam's coupled L2 term (added to the gradient); use adamw() for decoupled decay. eps defaults to torch's 1e-8 and is then omitted from state().
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.adam_amsgrad¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.adam_amsgrad(self, *, max_lr: 'float', betas: 'tuple[float, float]' = (0.9, 0.999), weight_decay: 'float' = 0.0, eps: 'Optional[float]' = None) -> 'NNOptimParamsBuilder'
torch.optim.Adam with amsgrad=True. Same betas / eps mapping as adam().
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.adamw¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.adamw(self, *, max_lr: 'float', betas: 'tuple[float, float]' = (0.9, 0.999), weight_decay: 'float' = 0.01, eps: 'Optional[float]' = None) -> 'NNOptimParamsBuilder'
torch.optim.AdamW — Adam with decoupled weight decay: each step first scales the weights by 1 - lr * weight_decay and then applies the Adam update, instead of adding weight_decay * p to the gradient. weight_decay defaults to torch's 1e-2; betas / eps map as in adam().
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.sgd¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.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.
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.sgd_nesterov¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.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().
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.grad_clip¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.grad_clip(self, norm: 'Optional[float]') -> 'NNOptimParamsBuilder'
Global-L2 gradient-norm clipping. None = no clipping (the dataclass default; this method is the opt-in path).
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.accumulate_grad¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.accumulate_grad(self, batches: 'int') -> 'NNOptimParamsBuilder'
Accumulate gradients over batches mini-batches before stepping. Default (no call) leaves the dataclass at 1.
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.param_groups¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.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).
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.build¶
nnx.nn.params.nn_optim_params_builder.NNOptimParamsBuilder.build(self) -> 'NNOptimParams'
Construct the dataclass from the fields the user touched.
Details
Checks for the required fields first: when no variant method was
called, it raises a ValueError that names the variant methods,
rather than the dataclass's missing-argument TypeError. Every
params Builder's ``build()`` behaves this way.
Only the fields the caller set are passed to the dataclass, so
every untouched optional field keeps its default. ``state()``
omits fields at their default, so the built value has the same
``state()`` and run id as one constructed directly.
Raises:
ValueError: if no variant method (`.adam`, `.adam_amsgrad`,
`.adamw`, `.sgd`, `.sgd_nesterov`) was called before
`.build()`. The message names the five methods so the user
can fix the chain without consulting the dataclass schema.
nnx.optimizers.NNOptimFactoryParams¶
class nnx.optimizers.NNOptimFactoryParams(*, factory: 'OptimizerFactorySpec', max_lr: 'float', weight_decay: 'float' = 0.0, grad_clip_norm: 'Optional[float]' = None, accumulate_grad_batches: 'int' = 1, param_groups: 'Optional[list[NNParamGroupSpec]]' = None) -> 'None'
Optimizer config backed by a registered factory.
Details
The registered-variant counterpart of
:class:`~nnx.nn.params.nn_optim_params.NNOptimParams`: it shares
``max_lr`` / ``weight_decay`` (defaults for every resolved group),
``grad_clip_norm``, ``accumulate_grad_batches`` and ``param_groups``
with the same validation, but names its optimizer through ``factory``
instead of an :class:`~nnx.nn.enum.optims.Optims` value and carries no
``momentum`` / ``name`` fields. Its ``state()`` holds a ``factory``
entry, which is how run decoding tells the two variants apart.
nnx.optimizers.NNOptimFactoryParams.is_valid¶
nnx.optimizers.NNOptimFactoryParams.is_valid(self) -> 'bool'
Always True: every field was validated at construction (the built-in variant's momentum-shape check has no counterpart here). Whether the factory is registered is checked when training starts.
nnx.optimizers.NNOptimFactoryParams.state¶
nnx.optimizers.NNOptimFactoryParams.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.optimizers.NNOptimFactoryParams.from_state¶
nnx.optimizers.NNOptimFactoryParams.from_state(state: 'Mapping[str, Any]') -> 'NNOptimFactoryParams'
No public description is currently available.
nnx.nn.params.nn_scheduler_params.NNSchedulerParams¶
class nnx.nn.params.nn_scheduler_params.NNSchedulerParams(*, min_lr: 'float', factor: 'float', patience: 'int', cooldown: 'int', threshold: 'float', kind: 'Optional[Schedulers]' = None, step_size: 'Optional[int]' = None, T_max: 'Optional[int]' = None, max_lr: 'Optional[float]' = None, total_steps: 'Optional[int]' = None, warmup_steps: 'Optional[int]' = None, clock: 'SchedulerClockUnit' = 'epoch') -> 'None'
A learning-rate scheduler: its kind (ReduceLROnPlateau when kind is None), the plateau fields, the kind-specific horizons and the clock that one step counts (an epoch or an optimizer update).
nnx.nn.params.nn_scheduler_params.NNSchedulerParams.state¶
nnx.nn.params.nn_scheduler_params.NNSchedulerParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_scheduler_params.NNSchedulerParams.from_state¶
nnx.nn.params.nn_scheduler_params.NNSchedulerParams.from_state(state: 'dict') -> 'NNSchedulerParams'
No public description is currently available.
nnx.nn.params.nn_scheduler_params.NNSchedulerParams.builder¶
nnx.nn.params.nn_scheduler_params.NNSchedulerParams.builder() -> 'NNSchedulerParamsBuilder'
Return a variant-aware builder. See NNSchedulerParamsBuilder.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder¶
class nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder() -> 'None'
Variant-aware builder for NNSchedulerParams.
Details
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.
`copy()` branches a (possibly empty) builder and `from_params()`
rebuilds one from an existing `NNSchedulerParams`.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.copy¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.copy(self) -> 'NNSchedulerParamsBuilder'
Return an independent branch of this builder.
Details
The branch starts with the same fields; a later variant call on
either builder replaces only that builder's fields. Nothing is
built or validated here — `build()` validates each branch.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.from_params¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.from_params(params: 'NNSchedulerParams') -> 'NNSchedulerParamsBuilder'
Return a builder pre-loaded with every field of params.
Details
`from_params(params).build()` equals `params`, with the same
`state()` (key order and omitted defaults included). A variant
call afterwards replaces the whole configuration, as usual,
except the clock, which it keeps.
Raises:
TypeError: if `params` is not exactly an `NNSchedulerParams`.
ValueError: if `params` carries a field this builder cannot
reproduce.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.clock¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.clock(self, clock: 'str') -> 'NNSchedulerParamsBuilder'
What one scheduler step counts: "epoch" (the default) or "optimizer_update" — once per committed update of the scheduler's optimizer, with horizons counted in updates. Call before or after the variant method; a variant call keeps the clock (a loaded one too).
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.reduce_on_plateau¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.reduce_on_plateau(self, *, min_lr: 'float', factor: 'float', patience: 'int', cooldown: 'int', threshold: 'float') -> 'NNSchedulerParamsBuilder'
ReduceLROnPlateau — the default scheduler.
Details
Sets the five plateau fields (``factor`` in ``(0, 1)``: it reduces
the LR). `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.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.step¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.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.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.cosine_annealing¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.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.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.one_cycle¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.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.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.linear_warmup_decay¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.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 the optimizer's LR over warmup_steps, linear decay to 0 over the remaining total_steps - warmup_steps. Used by most transformer training recipes.
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.build¶
nnx.nn.params.nn_scheduler_params_builder.NNSchedulerParamsBuilder.build(self) -> 'NNSchedulerParams'
Construct the dataclass from the fields the user touched.
Details
Checks for the required fields first: when no variant method was
called, it raises a ValueError that names the variant methods,
rather than the dataclass's missing-argument TypeError. Every
params Builder's ``build()`` behaves this way.
Only the fields the caller set are passed to the dataclass, so
every untouched field keeps its default. ``state()`` omits fields
at their default, so the built value has the same ``state()`` and
run id as one constructed directly.
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.
nnx.nn.params.nn_transformer_params.NNTransformerParams¶
class nnx.nn.params.nn_transformer_params.NNTransformerParams(*, dropout_prob: 'float', n_heads: 'Optional[int]' = None, activation: 'Optional[Activations]' = leaky_relu, activations: 'Optional[list[Activations]]' = None, dropout_probs: 'Optional[list[float]]' = None, input_dim: 'int', output_dim: 'int', hidden_dims: 'Optional[list[int]]' = None, vocab_size: 'int', n_layers: 'int', d_model: 'int', max_seq_len: 'int', ffn_mult: 'int' = 4, rope_base: 'float' = 10000.0, tie_embeddings: 'bool' = True, attn_dropout: 'float' = 0.0, resid_dropout: 'float' = 0.0) -> 'None'
Architecture of the decoder-only TransformerNN: vocabulary size, depth, width, attention heads and the maximum sequence length.
nnx.nn.params.nn_transformer_params.NNTransformerParams.state¶
nnx.nn.params.nn_transformer_params.NNTransformerParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_transformer_params.NNTransformerParams.from_state¶
nnx.nn.params.nn_transformer_params.NNTransformerParams.from_state(state: 'dict') -> 'NNTransformerParams'
No public description is currently available.
nnx.nn.params.nn_transformer_params.NNTransformerParams.builder¶
nnx.nn.params.nn_transformer_params.NNTransformerParams.builder() -> 'NNTransformerParamsBuilder'
Return a fluent LM-path builder. See NNTransformerParamsBuilder.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder¶
class nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder() -> 'None'
Builder for NNTransformerParams.
Details
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.
`copy()` branches a (possibly partial) builder and `from_params()`
rebuilds one from an existing `NNTransformerParams` — including the
inherited `NNParams` fields (`hidden_dims`, `activation(s)`,
`dropout_prob(s)`, custom `input_dim` / `output_dim`) no setter
exposes.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.copy¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.copy(self) -> 'NNTransformerParamsBuilder'
Return an independent branch of this builder, complete or partial.
Details
The branch starts with the same fields; afterwards setters on
either builder never affect the other. List fields carried from
`from_params` are copied. Nothing is built or validated here —
`build()` validates each branch on its own.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.from_params¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.from_params(params: 'NNTransformerParams') -> 'NNTransformerParamsBuilder'
Return a builder pre-loaded with every field of params.
Details
`from_params(params).build()` equals `params`, with the same
`state()` (key order and omitted defaults included) — the
inherited `hidden_dims`, `activation`, `activations`,
`dropout_prob` and `dropout_probs` survive because `build()`
applies its LM-path defaults *under* the carried fields. Setters
keep their usual rules afterwards: `.vocab(size)` still resets
`input_dim` / `output_dim` to `size`.
Raises:
TypeError: if `params` is not exactly an `NNTransformerParams`
(base `NNParams` or a subclass).
ValueError: if `params` carries a field this builder cannot
reproduce.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.vocab¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.vocab(self, size: 'int') -> 'NNTransformerParamsBuilder'
Set the vocabulary size. Mirrors into both input_dim and output_dim on the parent NNParams (the LM convention).
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.layers¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.layers(self, *, n: 'int', heads: 'int', d_model: 'int') -> 'NNTransformerParamsBuilder'
Set depth (n_layers), attention head count (n_heads), and hidden dimension (d_model). All three must be positive integers — Python or NumPy integers, normalized to plain int; floats (even 128.0), booleans and strings raise ValueError naming the argument — and the check runs before the divisibility arithmetic so the builder rejects exactly what the direct constructor rejects. Enforces d_model % heads == 0 and an even head width (d_model / heads, required by RoPE) 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.
Details
Raises:
ValueError: for a non-integral or non-positive `n`, `heads`
or `d_model`, for `d_model % heads != 0`, or for an odd
`d_model // heads`.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.ffn¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.ffn(self, *, mult: 'int') -> 'NNTransformerParamsBuilder'
FFN expansion ratio. Default is 4 (the SwiGLU-friendly ratio); only call this method to override.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.context¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.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.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.dropout¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.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).
Details
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.
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.tied_embeddings¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.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).
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.build¶
nnx.nn.params.nn_transformer_params_builder.NNTransformerParamsBuilder.build(self) -> 'NNTransformerParams'
Construct the dataclass.
Details
Checks for the required fields first: when a required setter was
not called, it raises a ValueError that names the missing setter
methods, rather than the dataclass's missing-argument TypeError.
Every params Builder's ``build()`` behaves this way.
Supplies `dropout_prob=0.0` — the one dead parent-NNParams field
the TransformerNN net never reads but the parent dataclass
requires — *under* the set fields, so a value carried by
`from_params` wins; `hidden_dims` / `activation` fall to the
NNParams defaults (`None` / `LEAKY_RELU`) unless carried. Using
the parent's own defaults keeps the builder's `state()` /
`run.id` identical to the direct-kwarg ctor's.
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.
nnx.nn.params.nn_conv_params.NNConvParams¶
class nnx.nn.params.nn_conv_params.NNConvParams(*, dropout_prob: 'float', n_heads: 'Optional[int]' = None, activation: 'Optional[Activations]' = leaky_relu, activations: 'Optional[list[Activations]]' = None, dropout_probs: 'Optional[list[float]]' = None, input_dim: 'int', output_dim: 'int', hidden_dims: 'Optional[list[int]]' = None, conv_channels: 'list[int]', in_channels: 'int' = 1, kernel_size: 'int' = 5, stride: 'int' = 1, padding: 'int' = 0, pool_size: 'int' = 2) -> 'None'
Parameters for a ConvNN with a required conv-block activation.
nnx.nn.params.nn_conv_params.NNConvParams.image_side¶
nnx.nn.params.nn_conv_params.NNConvParams.image_side(self) -> 'int'
Spatial side of the (square) input image.
Details
Always a plain ``int``: every count that feeds the conv arithmetic
(``input_dim``, ``in_channels``, ``kernel_size``, ``stride``,
``padding``, ``pool_size``, ``conv_channels``) is validated at
construction as a non-boolean integer — NumPy integers accepted and
normalized — so a fractional or boolean knob raises ``ValueError``
naming the field before any shape helper or layer runs.
nnx.nn.params.nn_conv_params.NNConvParams.spatial_sizes¶
nnx.nn.params.nn_conv_params.NNConvParams.spatial_sizes(self) -> 'list[int]'
Feature-map side after each Conv→Pool block (floor arithmetic, matching Conv2d/MaxPool2d).
nnx.nn.params.nn_conv_params.NNConvParams.flatten_dim¶
nnx.nn.params.nn_conv_params.NNConvParams.flatten_dim(self) -> 'int'
Input width of the first FC layer: last block's channels × side².
nnx.nn.params.nn_conv_params.NNConvParams.state¶
nnx.nn.params.nn_conv_params.NNConvParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_conv_params.NNConvParams.from_state¶
nnx.nn.params.nn_conv_params.NNConvParams.from_state(state: 'dict') -> 'NNConvParams'
No public description is currently available.
nnx.nn.params.nn_vit_params.NNViTParams¶
class nnx.nn.params.nn_vit_params.NNViTParams(*, dropout_prob: 'float', n_heads: 'Optional[int]' = None, activation: 'Optional[Activations]' = leaky_relu, activations: 'Optional[list[Activations]]' = None, dropout_probs: 'Optional[list[float]]' = None, input_dim: 'int', output_dim: 'int', hidden_dims: 'Optional[list[int]]' = None, image_size: 'int', patch_size: 'int', d_model: 'int', n_layers: 'int', in_channels: 'int' = 3, ffn_mult: 'int' = 4, attn_dropout: 'float' = 0.0, resid_dropout: 'float' = 0.0) -> 'None'
Parameters for ViTNN, built by Nets.VIT.
nnx.nn.params.nn_vit_params.NNViTParams.n_patches¶
property nnx.nn.params.nn_vit_params.NNViTParams.n_patches
Patch tokens per image (the encoder adds one CLS token).
nnx.nn.params.nn_vit_params.NNViTParams.state¶
nnx.nn.params.nn_vit_params.NNViTParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_vit_params.NNViTParams.from_state¶
nnx.nn.params.nn_vit_params.NNViTParams.from_state(state: 'dict') -> 'NNViTParams'
No public description is currently available.
nnx.nn.params.nn_moe_params.NNMoEParams¶
class nnx.nn.params.nn_moe_params.NNMoEParams(*, dropout_prob: 'float', n_heads: 'Optional[int]' = None, activation: 'Optional[Activations]' = leaky_relu, activations: 'Optional[list[Activations]]' = None, dropout_probs: 'Optional[list[float]]' = None, input_dim: 'int', output_dim: 'int', hidden_dims: 'Optional[list[int]]' = None, num_experts: 'int', top_k: 'int' = 2) -> 'None'
Serializable parameters for an expert-bearing feed-forward MoE.
Details
Unlike base :class:`NNParams`, ``hidden_dims`` must contain at least one
layer because only hidden layers are replaced by ``MoELinear`` modules.
nnx.nn.params.nn_moe_params.NNMoEParams.state¶
nnx.nn.params.nn_moe_params.NNMoEParams.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_moe_params.NNMoEParams.from_state¶
nnx.nn.params.nn_moe_params.NNMoEParams.from_state(state: 'dict') -> 'NNMoEParams'
No public description is currently available.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams¶
class nnx.nn.params.nn_tokenizer_params.NNTokenizerParams(*, path: 'str', tokenizer: 'object') -> 'None'
Frozen dataclass holding a tokenizer + its on-disk pointer.
Details
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.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.of¶
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.of(tokenizer: 'object', path: 'str') -> 'NNTokenizerParams'
Construct from a live Tokenizer instance and persist it to path.
Details
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.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.from_state¶
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.from_state(state: 'dict') -> 'NNTokenizerParams'
Load from a state dict produced by state. The single required key is path; the tokenizer is reconstructed from the file the path points to.
Details
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.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.state¶
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.state(self) -> 'dict'
Return the serializable view — only the path goes into run.yaml.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.vocab_size¶
property nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.vocab_size
No public description is currently available.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.encode¶
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.encode(self, text: 'str') -> 'list[int]'
No public description is currently available.
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.decode¶
nnx.nn.params.nn_tokenizer_params.NNTokenizerParams.decode(self, ids: 'list[int]', skip_special_tokens: 'bool' = True) -> 'str'
No public description is currently available.
nnx.nn.params.nn_tokenizer_params.train_bpe¶
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) -> 'TokenizerType'
Train a BPE tokenizer on either a list of files or a list of texts.
Details
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. ``["<unk>", "<pad>", "<bos>", "<eos>"]`` (the
default). Included in the vocab and not split during
tokenization. ``"<unk>"`` is the model's unknown token: it is
appended when missing, so an unseen character encodes instead
of raising.
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.
nnx.nn.params.nn_run.NNRun¶
class nnx.nn.params.nn_run.NNRun(*, net: 'Optional[NNParams]', train: 'NNTrainParams', model: 'NNModelParams', trainer: 'Optional[NNTrainerParams]' = None, salt: 'Optional[str]' = None, transforms: 'tuple[NNCheckpointTransform, ...]' = (), idps: 'Optional[list[NNIterationDataPoint]]' = None, resume_status: 'Optional[ResumeStatus]' = None, provenance: 'Optional[ProvenanceRecord]' = None, history: 'Optional[str]' = None, precision: 'Optional[ResolvedPrecision]' = None, compile: 'Optional[CompileRecord]' = None) -> 'None'
A training run: its network, training and model parameters, its id and its persisted history and checkpoints under runs/<id>/. Load one with NNRun.load.
nnx.nn.params.nn_run.NNRun.id¶
property nnx.nn.params.nn_run.NNRun.id
No public description is currently available.
nnx.nn.params.nn_run.NNRun.state¶
nnx.nn.params.nn_run.NNRun.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.with_resume_status¶
nnx.nn.params.nn_run.NNRun.with_resume_status(self, value: 'Optional[ResumeStatus]') -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.with_precision¶
nnx.nn.params.nn_run.NNRun.with_precision(self, value: 'Optional[ResolvedPrecision]') -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.with_compile¶
nnx.nn.params.nn_run.NNRun.with_compile(self, value: 'Optional[CompileRecord]') -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.with_idps¶
nnx.nn.params.nn_run.NNRun.with_idps(self, value: 'list[NNIterationDataPoint]') -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.with_provenance¶
nnx.nn.params.nn_run.NNRun.with_provenance(self, value: 'Optional[ProvenanceRecord]') -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.with_history¶
nnx.nn.params.nn_run.NNRun.with_history(self, value: 'Optional[str]') -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.ensure_writable¶
nnx.nn.params.nn_run.NNRun.ensure_writable(self, root: 'Optional[str]' = None, *, overwrite: 'bool' = False) -> 'None'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.writable_lease¶
nnx.nn.params.nn_run.NNRun.writable_lease(self, root: 'Optional[str]' = None, *, overwrite: 'bool' = False) -> 'Iterator[None]'
Reserve this run ID and hold exclusive ownership until training ends.
nnx.nn.params.nn_run.NNRun.checkpoints¶
nnx.nn.params.nn_run.NNRun.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.
Details
BEST is not a cadence checkpoint, so it is not in this list (its
epoch is often none of the phase epochs). Load it directly via
``NNCheckpoint.load(run=run.id, type=Checkpoints.BEST)``.
nnx.nn.params.nn_run.NNRun.save¶
nnx.nn.params.nn_run.NNRun.save(self, root: 'Optional[str]' = None, *, update_best: 'bool' = True) -> 'NNRun'
No public description is currently available.
nnx.nn.params.nn_run.NNRun.load¶
nnx.nn.params.nn_run.NNRun.load(id: 'str', root: 'Optional[str]' = None) -> 'NNRun'
Load a saved run. Its history is idps.csv read whole — or, for a run with a history journal, the journal's committed tail: the last retention records, from the chunks that hold them. Either way records past the LAST checkpoint's epoch are left out.
nnx.nn.params.nn_run.NNRun.all¶
nnx.nn.params.nn_run.NNRun.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.
nnx.nn.params.nn_checkpoint.NNCheckpointTransform¶
class nnx.nn.params.nn_checkpoint.NNCheckpointTransform(*, name: 'str', version: 'int' = 1, options: 'dict[str, Any]' = <factory>) -> 'None'
A versioned recipe for rebuilding a checkpoint's module topology.
nnx.nn.params.nn_checkpoint.NNCheckpointTransform.state¶
nnx.nn.params.nn_checkpoint.NNCheckpointTransform.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.nn.params.nn_checkpoint.NNCheckpointTransform.from_state¶
nnx.nn.params.nn_checkpoint.NNCheckpointTransform.from_state(state: 'dict[str, Any]') -> 'NNCheckpointTransform'
No public description is currently available.
nnx.nn.params.nn_checkpoint.NNCheckpoint¶
class nnx.nn.params.nn_checkpoint.NNCheckpoint(*, net_params: 'Optional[NNParams]', net_state: 'dict[str, Any]', model_params: 'NNModelParams', idp: 'NNIterationDataPoint', transforms: 'tuple[NNCheckpointTransform, ...]' = (), training_state_id: 'Optional[str]' = None, training_state_present: 'Optional[bool]' = None) -> 'None'
Model state plus the recipes needed to rebuild its module topology.
Details
``transforms`` is empty for ordinary and legacy checkpoints. Training
callbacks that replace modules at ``on_train_end`` can persist ordered,
versioned recipes here; :meth:`NNModel.from_checkpoint` replays recognized
recipes before loading ``net_state``.
nnx.nn.params.nn_checkpoint.NNCheckpoint.reconstructible¶
property nnx.nn.params.nn_checkpoint.NNCheckpoint.reconstructible
Whether a model can be rebuilt from this checkpoint alone: False for a runtime-only module, whose weights need the module again (NNModel.from_checkpoint(ckpt, module=...)).
nnx.nn.params.nn_checkpoint.NNCheckpoint.to_file¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.to_file(self, path: 'str', format: "Literal['pickle', 'safetensors']" = 'pickle') -> 'None'
Atomically write this NNCheckpoint to path.
Details
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 + transform
recipes
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. Portable, so a runtime-only
module (``reconstructible=False``) is rejected
with ``MissingModelFactoryError`` before anything is
written.
Both formats write to a hidden temporary file beside ``path``
(``.<name>.<random>``), flush it to disk and rename it into place,
so an interrupted save can never leave a half-written checkpoint at
the destination — matching the atomicity guarantee NNRun.save
offers for YAML/CSV.
nnx.nn.params.nn_checkpoint.NNCheckpoint.save¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.save(self, run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None, optimizer_state: 'Optional[dict[str, Any]]' = None, scheduler_state: 'Optional[dict[str, Any]]' = None, scaler_state: 'Optional[dict[str, Any]]' = None, rng_state: 'Optional[dict[str, Any]]' = None, completed_epoch: 'Optional[int]' = None, resume_net_state: 'Optional[dict[str, Any]]' = None, optimizer_type: 'Optional[str]' = None, scheduler_type: 'Optional[str]' = None, optimizer_topology: 'Optional[list[list[dict[str, Any]]]]' = None, optimizer_factory: 'Optional[dict[str, Any]]' = None, components: 'Optional[dict[str, Any]]' = None, optimizers_state: 'Optional[dict[str, Any]]' = None, schedulers_state: 'Optional[dict[str, Any]]' = None, optimizer_types: 'Optional[dict[str, str]]' = None, scheduler_types: 'Optional[dict[str, str]]' = None, optimizer_topologies: 'Optional[dict[str, list[list[dict[str, Any]]]]]' = None, optimizer_factories: 'Optional[dict[str, Optional[dict[str, Any]]]]' = None, precision: 'Optional[dict[str, Any]]' = None, compile: 'Optional[dict[str, Any]]' = None, distributed: 'Optional[dict[str, Any]]' = None, counters: 'Optional[dict[str, Any]]' = None) -> 'None'
Save the checkpoint to disk atomically.
Details
Every save is a resume point: a manifest
(``<tag>.pt.manifest.json``, written last) stamps it with the tag's
generation ordinal, the checkpoint id and the SHA-256 of each file it
consists of, so a resume refuses a torn, mixed or corrupted point
(:meth:`verify`). Files are flushed to stable storage
(``F_FULLFSYNC`` on macOS, ``fsync`` elsewhere) before they are
renamed into place, and a Ctrl-C arriving meanwhile is held until the point is
complete. ``counters`` (the run's logical position: ``global_step``,
``committed_updates``, ``planned_n_epochs``) rides in the training
state.
``components`` is the ``ComponentRegistry.collect()``
mapping of every registered component's versioned state; it lives
in the same generation sidecar as the optimizer state, so a model
and component state from different generations are never paired.
``optimizers_state`` / ``schedulers_state`` (with their
``*_types``, ``optimizer_topologies`` and ``optimizer_factories``)
carry a multi-optimizer ``Trainer``'s name-keyed states; either
``optimizer_state`` or ``optimizers_state`` makes the checkpoint
stateful.
When `optimizer_state` is supplied, a generation-addressed sibling
file holds the training state, plus a fixed-name compatibility copy.
This sidecar is used by NNModel.train when NNTrainParams'
resume_from_run_id / resume_from_checkpoint name this checkpoint, to
warm-resume with the prior optimizer momentum / Adam state.
The immutable generation sidecar is committed first and the checkpoint
second. The checkpoint names the sidecar it owns, so interruption
between replacements leaves the previous generation resumable.
nnx.nn.params.nn_checkpoint.NNCheckpoint.resume_point¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.resume_point(run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None) -> 'Optional[dict[str, Any]]'
The manifest stamping this checkpoint's resume point — generation, checkpoint_id, completed_epoch, global_step, committed_updates, planned_n_epochs and the files with their SHA-256 — or None for a checkpoint without a manifest. Read as is: after an interrupted publish the live manifest may name the previous checkpoint while a staged one describes the published files — verify checks both and returns the one that matches.
nnx.nn.params.nn_checkpoint.NNCheckpoint.verify¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.verify(run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None) -> 'Optional[dict[str, Any]]'
Check a resume point before anything is restored from it and return its manifest (None for a checkpoint without one, which loads as before). Raises ResumePointError naming the reason: a listed file missing or a manifest naming another checkpoint (torn), files of different generations (mixed), or a file whose SHA-256 differs (digest mismatch). Every listed file is hashed; only the matching training-state sidecar is deserialized, with torch.load(weights_only=True), to compare its generation; no checkpoint data is written (only the checkpoint's lock file is taken, created when missing): a point whose writer stopped between publishing the checkpoint and its live manifest is recognized by its staged manifest — every digest matching — and that manifest is returned (the writer's next save makes it live).
nnx.nn.params.nn_checkpoint.NNCheckpoint.load_training_state¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.load_training_state(run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None, map_location: 'Any' = 'cpu') -> 'Optional[dict[str, Any]]'
Load and validate the resumable optimizer/scheduler/scaler bundle.
Details
Legacy optimizer-only sidecars are normalized into the new mapping so
checkpoints written by older NNx versions remain resumable.
Returns ``None`` — without creating ``runs/<run>/`` — when the run
directory does not exist; the run ID is still validated first.
nnx.nn.params.nn_checkpoint.NNCheckpoint.load_with_training_state¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.load_with_training_state(run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None, map_location: 'Any' = 'cpu') -> 'tuple[Optional[NNCheckpoint], Optional[dict[str, Any]]]'
Atomically load a checkpoint and its matching training-state bundle.
Details
Returns ``(None, None)`` — without creating ``runs/<run>/`` — when
the run directory does not exist; the run ID is still validated.
nnx.nn.params.nn_checkpoint.NNCheckpoint.load_optimizer_state¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.load_optimizer_state(run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None) -> 'Optional[dict[str, Any]]'
Load the optimizer state sidecar for a checkpoint. Returns None when no sidecar exists (e.g., checkpoints written without resume support).
Details
The sidecar itself is read with ``weights_only=True``. Locating it
reads the checkpoint through :meth:`load_training_state`, which
loads a pickle-format checkpoint with ``NNCheckpoint.from_file``:
trust the run directory as you would for ``from_file``.
nnx.nn.params.nn_checkpoint.NNCheckpoint.from_file¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.from_file(path: 'str', map_location: 'Any' = 'cpu') -> 'Optional[NNCheckpoint]'
Load an NNCheckpoint from disk, auto-detecting pickle vs safetensors.
Details
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.
nnx.nn.params.nn_checkpoint.NNCheckpoint.load¶
nnx.nn.params.nn_checkpoint.NNCheckpoint.load(run: 'str', type: 'Checkpoints', root: 'Optional[str]' = None, map_location: 'Any' = 'cpu') -> 'Optional[NNCheckpoint]'
No public description is currently available.
nnx.nn.params.nn_checkpoint.ResumePointError¶
class nnx.nn.params.nn_checkpoint.ResumePointError
A resume point that cannot be trusted: torn (its manifest names another checkpoint), mixed (files from different generations) or corrupted (a file's digest differs). Raised before anything is restored; the message says which and why.
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint¶
class nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint(*, lr: 'float', iter_idx: 'int', epoch_idx: 'int', batch_idx: 'int', train_edp: 'NNEvaluationDataPoint', val_edp: 'Optional[NNEvaluationDataPoint]' = None, train_summary: 'Optional[NNEvaluationDataPoint]' = None, selection: 'Optional[MonitorRecord]' = None, update_count: 'Optional[int]' = None) -> 'None'
One row in the per-iteration training log.
Details
`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.
**Epoch summary.** A run that declares metrics or a monitor
also records, on the same last idp of each epoch, `train_summary` —
the whole-epoch training record (loss and error averaged with every
batch's own denominator, declared metrics over the full sample) — and
`selection`, the epoch's :class:`~nnx.MonitorRecord` (monitor identity,
value, status and whether the epoch improved). Both are omitted from
`state()` otherwise, so legacy history is unchanged.
**Update counter.** Objective runs also record
`update_count` — the committed optimizer updates so far — distinct from
the microbatch counter `iter_idx` and the epoch `epoch_idx`; omitted
otherwise.
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.with_val_edp¶
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.with_val_edp(self, value: 'Optional[NNEvaluationDataPoint]') -> 'NNIterationDataPoint'
No public description is currently available.
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.with_epoch_summary¶
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.with_epoch_summary(self, train_summary: 'Optional[NNEvaluationDataPoint]', selection: 'Optional[MonitorRecord]') -> 'NNIterationDataPoint'
No public description is currently available.
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.monitored_train_edp¶
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.monitored_train_edp(self) -> 'NNEvaluationDataPoint'
The epoch's training record for monitors: the whole-epoch summary when recorded, else this (last) batch's record.
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.state¶
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.from_state¶
nnx.nn.params.nn_iteration_data_point.NNIterationDataPoint.from_state(state: 'dict') -> 'NNIterationDataPoint'
No public description is currently available.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint¶
class nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint(*, f1: 'Optional[float]' = None, recall: 'Optional[float]' = None, accuracy: 'Optional[float]' = None, precision: 'Optional[float]' = None, loss: 'Optional[float]' = None, error: 'Optional[float]' = None, extra: 'Mapping[str, float]' = <factory>, kind: 'Optional[str]' = None, count: 'Optional[int]' = None, status: 'Optional[str]' = None, metrics: 'Mapping[str, float]' = <factory>) -> 'None'
Per-batch / per-epoch evaluation metrics.
Details
The four classification 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(extra_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).
**Task records.** A model with a ``TaskSpec`` writes records
that also carry the task `kind` (``"categorical"`` / ``"multilabel"`` /
``"regression"``), the `count` of valid (unmasked) targets they
summarize, a `status` (``"ok"``, or ``"empty"`` when every target was
masked) and the task's own `metrics` (``mse`` / ``mae`` for regression,
``subset_accuracy`` / ``element_accuracy`` for multilabel). A
regression record leaves the classification fields and `error` as
``None`` rather than fabricating them; an empty record has no loss or
metrics at all. All four task fields are omitted from `state()` on
legacy records, whose serialization is unchanged.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.with_loss¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.with_loss(self, value: 'float')
No public description is currently available.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.with_error¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.with_error(self, value: 'float')
No public description is currently available.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.with_extra¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.with_extra(self, name: 'str', value: 'float') -> 'NNEvaluationDataPoint'
No public description is currently available.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.of¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.of(Y: 'np.ndarray', Y_hat: 'np.ndarray', average: 'str' = 'macro', extra_metrics: 'Optional[Mapping[str, Callable]]' = None)
Compute per-batch evaluation metrics.
Details
`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_true, y_pred) -> float} map
of user-supplied custom metrics, called as ``fn(Y, Y_hat)`` — truth
first, decoded predictions second — once on the arrays passed here
(one batch in the default training step, the aggregate in
`NNModel.evaluate`) and stored in the returned object's `extra` dict.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.mean_of¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.mean_of(edps: 'list[NNEvaluationDataPoint]') -> 'NNEvaluationDataPoint'
Unweighted-mean reduce a list of EDPs across every metric.
Details
.. 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).
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.state¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.state(self) -> 'dict'
No public description is currently available.
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.from_state¶
nnx.nn.params.nn_evaluation_data_point.NNEvaluationDataPoint.from_state(state: 'dict') -> 'NNEvaluationDataPoint'
No public description is currently available.
3.1. Bounded history journal (nnx.history)¶
nnx.history.HistoryJournal¶
class nnx.history.HistoryJournal(retention: 'int' = 1000, chunk_size: 'Optional[int]' = None) -> 'None'
Opt-in bounded history for one training run.
Details
Args:
retention: how many of the most recent records stay in memory (the
``ctx.idps`` window, ``NNRun.idps`` of the returned and of a
loaded run).
chunk_size: records per journal chunk file, at most ``retention``
(so the records waiting for their chunk are always inside the
window). ``None`` (the default) means ``min(250, retention)``;
:attr:`chunk` is the resolved size.
nnx.history.HistoryJournal.chunk¶
property nnx.history.HistoryJournal.chunk
Records per chunk file: chunk_size, or min(250, retention).
nnx.history.iter_history¶
nnx.history.iter_history(run_id: 'str', root: 'Optional[str]' = None, *, lineage: 'bool' = False) -> 'Iterator[NNIterationDataPoint]'
Stream a run's committed history in order: a journal chunk by chunk (each checked against its SHA-256), a legacy idps.csv whole (floats parsed round-trip, so its export reproduces the file). Records past the LAST checkpoint's epoch are never yielded.
Details
With ``lineage=True`` a resumed run is preceded by the committed records
of the run it resumed from, up to and including the epoch it resumed
from (recorded in its resume status; for runs recorded before that, the
epoch of that checkpoint now) and never past the run's own first record,
recursively: each epoch once, never the parent's later epochs. When
that epoch cannot be known (the checkpoint, or the parent's LAST, is
unreadable), the parent's records before the run's first record are
used, with a ``RuntimeWarning``; when neither is known, or the parent
run is gone, the lineage starts at the run, with a ``RuntimeWarning``.
Lineage follows resumes: a run that started fresh (a born-again
generation naming its teacher's run as ``parent_run_id``, say) starts
its own epochs, and its lineage is its own history.
nnx.history.export_history_csv¶
nnx.history.export_history_csv(run_id: 'str', path: 'Union[str, os.PathLike[str]]', root: 'Optional[str]' = None, *, lineage: 'bool' = False) -> 'int'
Write a run's committed history (see iter_history) to path in the legacy idps.csv layout — the same columns, order and index NNRun.save writes — atomically. Returns the number of records. The export holds the records in memory while it builds the table; stream iter_history for histories that do not fit.
nnx.history.migrate_history¶
nnx.history.migrate_history(run_id: 'str', root: 'Optional[str]' = None, *, spec: 'Optional[HistoryJournal]' = None, discard_uncommitted: 'bool' = False) -> 'None'
Move a legacy run's idps.csv into a history journal, explicitly.
Details
The committed records are written as a journal (with per-epoch summary
rows), published and read back, and only then is ``idps.csv`` removed;
``run.yaml`` — and so the run id — is untouched. A leftover journal from
an interrupted migration (never published) is replaced. Refused: a run
that already keeps a journal, a run that is being trained (its lease is
held), and — unless ``discard_uncommitted=True`` — a CSV holding records
past the LAST checkpoint's epoch, which readers hide but a migration
would delete.
nnx.history.HistoryCorruptionError¶
class nnx.history.HistoryCorruptionError
A history journal whose committed files do not match its manifest: a chunk whose bytes changed, a broken index chain, a missing file.
4. Networks¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN¶
class nnx.nn.net.feed_fwd_nn.FeedFwdNN(params: 'NNParams')
Feed-forward classifier built from NNParams: flatten, a stack of hidden nn.Linear layers (per-layer activation and dropout), then a linear head. forward(x) returns logits.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.params¶
property nnx.nn.net.feed_fwd_nn.FeedFwdNN.params
No public description is currently available.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.forward¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN.forward(self, X: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.seed_count¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN.seed_count(self, batch) -> 'Optional[int]'
The seed rows of a PyG NeighborLoader / full-graph split batch, as GraphNNBase.seed_count: this net also reads PyG node batches, whose other rows belong to other splits and must not be scored.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.unpack_batch¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN.unpack_batch(self, batch)
No public description is currently available.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.to_file¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN.to_file(self, path: 'str') -> 'None'
No public description is currently available.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.from_file¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN.from_file(path: 'str', params: 'NNParams', map_location='cpu') -> 'FeedFwdNN'
No public description is currently available.
nnx.nn.net.feed_fwd_nn.FeedFwdNN.from_state¶
nnx.nn.net.feed_fwd_nn.FeedFwdNN.from_state(state_dict: 'dict', params: 'NNParams') -> 'FeedFwdNN'
No public description is currently available.
nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN¶
class nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN(params: 'NNMoEParams')
Feed-forward classifier built from NNMoEParams whose hidden layers are top-k routed MoELinear experts; the head stays a plain nn.Linear. forward(x) returns logits.
nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN.forward¶
nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN.forward(self, X: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN.params¶
property nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN.params
No public description is currently available.
nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN.unpack_batch¶
nnx.nn.net.feed_fwd_moe_nn.FeedFwdMoENN.unpack_batch(self, batch)
No public description is currently available.
nnx.nn.net.conv_nn.ConvNN¶
class nnx.nn.net.conv_nn.ConvNN(params: 'NNConvParams')
LeNet-style convolutional classifier built from NNConvParams: conv and pool blocks, then fully connected layers. forward(x) takes (B, C, H, W) images or flattened rows and returns logits.
nnx.nn.net.conv_nn.ConvNN.params¶
property nnx.nn.net.conv_nn.ConvNN.params
The net's configuration; setting it (as model surgery does) re-resolves the per-layer state the forward reads.
nnx.nn.net.conv_nn.ConvNN.forward¶
nnx.nn.net.conv_nn.ConvNN.forward(self, X: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.nn.net.conv_nn.ConvNN.unpack_batch¶
nnx.nn.net.conv_nn.ConvNN.unpack_batch(self, batch)
No public description is currently available.
nnx.nn.net.graph_nn_base.GraphNNBase¶
class nnx.nn.net.graph_nn_base.GraphNNBase(params: 'NNParams')
Abstract base for GNN architectures.
Details
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.
nnx.nn.net.graph_nn_base.GraphNNBase.forward¶
nnx.nn.net.graph_nn_base.GraphNNBase.forward(self, X: 'torch.Tensor', E: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.nn.net.graph_nn_base.GraphNNBase.unpack_batch¶
nnx.nn.net.graph_nn_base.GraphNNBase.unpack_batch(self, batch) -> 'tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]'
No public description is currently available.
nnx.nn.net.graph_nn_base.GraphNNBase.seed_count¶
nnx.nn.net.graph_nn_base.GraphNNBase.seed_count(self, batch) -> 'Optional[int]'
Number of seed rows at the head of a NeighborLoader subgraph.
Details
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.
nnx.nn.net.graph_conv_nn.GraphConvNN¶
class nnx.nn.net.graph_conv_nn.GraphConvNN(params: 'NNParams')
Node-level graph net of PyG GCNConv layers (the graph extra: pip install "thekaveh-nnx[graph]").
nnx.nn.net.graph_sage_nn.GraphSageNN¶
class nnx.nn.net.graph_sage_nn.GraphSageNN(params: 'NNParams')
Node-level graph net of PyG SAGEConv layers (the graph extra: pip install "thekaveh-nnx[graph]").
nnx.nn.net.graph_att_nn.GraphAttNN¶
class nnx.nn.net.graph_att_nn.GraphAttNN(params: 'NNParams')
Node-level graph net of PyG GATConv layers (the graph extra: pip install "thekaveh-nnx[graph]").
nnx.nn.net.transformer_nn.TransformerNN¶
class nnx.nn.net.transformer_nn.TransformerNN(params: 'NNTransformerParams')
Decoder-only language model built from NNTransformerParams. forward(tokens) maps (B, T) token ids to (B, T, vocab) logits; forward_with_cache decodes with a per-layer KV cache.
nnx.nn.net.transformer_nn.TransformerNN.forward¶
nnx.nn.net.transformer_nn.TransformerNN.forward(self, tokens: 'torch.Tensor') -> 'torch.Tensor'
Args: tokens: (batch, seq) long tensor of token ids.
Details
Returns:
(batch, seq, vocab_size) logits — pre-softmax.
nnx.nn.net.transformer_nn.TransformerNN.forward_with_cache¶
nnx.nn.net.transformer_nn.TransformerNN.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.
Details
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.
nnx.nn.net.transformer_nn.TransformerNN.unpack_batch¶
nnx.nn.net.transformer_nn.TransformerNN.unpack_batch(self, batch)
Make TransformerNN compatible with the standard supervised NNModel training loop.
Details
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).
nnx.nn.net.transformer_nn.TransformerNN.to_file¶
nnx.nn.net.transformer_nn.TransformerNN.to_file(self, path: 'str') -> 'None'
No public description is currently available.
nnx.nn.net.transformer_nn.TransformerNN.from_file¶
nnx.nn.net.transformer_nn.TransformerNN.from_file(path: 'str', params: 'NNTransformerParams', map_location='cpu') -> 'TransformerNN'
No public description is currently available.
nnx.nn.net.transformer_nn.TransformerNN.from_state¶
nnx.nn.net.transformer_nn.TransformerNN.from_state(state_dict: 'dict', params: 'NNTransformerParams') -> 'TransformerNN'
No public description is currently available.
nnx.nn.net.vit_nn.ViTNN¶
class nnx.nn.net.vit_nn.ViTNN(*, 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)
Small Vision Transformer encoder.
Details
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.
nnx.nn.net.vit_nn.ViTNN.patch_positions¶
nnx.nn.net.vit_nn.ViTNN.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).
Details
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``).
nnx.nn.net.vit_nn.ViTNN.forward¶
nnx.nn.net.vit_nn.ViTNN.forward(self, x: 'torch.Tensor', mask: 'Optional[torch.Tensor]' = None) -> 'torch.Tensor'
Run the encoder.
Details
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.
nnx.nn.net.vit_nn.ViTNN.unpack_batch¶
nnx.nn.net.vit_nn.ViTNN.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.
nnx.nn.net.vit_nn.ViTBlock¶
class nnx.nn.net.vit_nn.ViTBlock(d_model: 'int', n_heads: 'int', ffn_mult: 'int' = 4, attn_dropout: 'float' = 0.0, resid_dropout: 'float' = 0.0)
Pre-norm ViT block: x = x + attn(RMSNorm(x)); x = x + ffn(RMSNorm(x)).
Details
Same shape as :class:`nnx.nn.net.transformer_layers.TransformerBlock`
but with bidirectional attention instead of causal. SwiGLU is
reused unchanged.
nnx.nn.net.vit_nn.ViTBlock.forward¶
nnx.nn.net.vit_nn.ViTBlock.forward(self, x: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.nn.moe.MoELinear¶
class nnx.nn.moe.MoELinear(in_features: 'int', out_features: 'int', *, num_experts: 'int', top_k: 'int' = 2)
Sparse top-k Mixture-of-Experts drop-in for nn.Linear.
Details
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``.
nnx.nn.moe.MoELinear.forward¶
nnx.nn.moe.MoELinear.forward(self, x: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.nn.moe.MoELinear.extra_repr¶
nnx.nn.moe.MoELinear.extra_repr(self) -> 'str'
No public description is currently available.
5. Datasets¶
nnx.nn.dataset.nn_dataset_base.NNDatasetBase¶
class nnx.nn.dataset.nn_dataset_base.NNDatasetBase() -> 'None'
Base of the NNx datasets: a name, the input and output sizes and the train / validation / test loaders, all set by the subclass at construction.
nnx.nn.dataset.nn_dataset_base.NNDatasetBase.state¶
nnx.nn.dataset.nn_dataset_base.NNDatasetBase.state(self) -> 'dict'
No public description is currently available.
nnx.nn.dataset.nn_dataset.NNDataset¶
class nnx.nn.dataset.nn_dataset.NNDataset(*, ds_class: 'type[VisionDataset]', root_dir: 'str' = './data', download: 'bool' = True, transform: 'Optional[Callable]' = None, batch_sizes: 'tuple[Optional[int], Optional[int], Optional[int]]' = (None, None, None), val_proportion: 'float' = 0.1, seed: 'Optional[int]' = None, train_transform: 'Optional[Callable]' = None, eval_transform: 'Optional[Callable]' = None) -> 'None'
Vision dataset wrapper (the vision extra: pip install "thekaveh-nnx[vision]"). 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).
Details
``batch_sizes`` is a ``(train, val, test)`` tuple. ``None`` — the
default for every slot — means *one batch holding the complete split*:
the default train loader therefore yields a single full-split batch per
epoch, i.e. **one optimizer step per epoch**. That suits small
full-batch fits but not stochastic training (diffusion, MoE routing,
self-supervision); pass an explicit train size such as
``batch_sizes=(128, None, None)`` for mini-batches. A positive integer
is used verbatim (a size larger than the split simply yields one smaller
batch); NumPy integers are normalized to ``int``. Zero, ``False``,
negatives, floats, strings and anything but a 3-tuple are rejected with
the ``batch_sizes[i] (split)`` slot named before ``ds_class`` is
instantiated — zero never disables a split; ``val_proportion=0.0`` does
(``val_loader`` is then ``None``; the resolved val size is a placeholder
``1`` when requested as ``None``, and an explicit size is kept).
``transform`` is handed to both torchvision factories and applies to
every split. Split transforms: ``train_transform`` and
``eval_transform`` add a per-split input transform on top of it — e.g. a
random augmentation for training and a fixed resize for validation and
test. Each split then becomes a ``nnx.preprocessing.SplitView`` over
the untouched base dataset: targets pass through, the base
``transform`` / labels / order never change, and views are safe to read
alternately and from DataLoader workers. ``input_dim`` is probed through
``eval_transform`` when one is set. Both default to ``None`` (the
existing single-transform behavior, plain ``Subset`` splits); ``state()``
gains ``train_transform`` / ``eval_transform`` descriptions only when
set — runtime callables are recorded by name and must be passed again
to rebuild the dataset.
nnx.nn.dataset.nn_graph_dataset.NNGraphDataset¶
class nnx.nn.dataset.nn_graph_dataset.NNGraphDataset(*, ds_class: 'type[Dataset]', n_neighbors: 'Optional[list[int]]' = None, root_dir: 'str' = './data', transform: 'Optional[Callable]' = None, n_workers: 'int' = 4, batch_sizes: 'tuple[Optional[int], Optional[int], Optional[int]]' = (None, None, None), seed: 'Optional[int]' = None, sampler: "Literal['neighbor', 'full']" = 'neighbor') -> 'None'
Single-graph node-classification wrapper over a PyG dataset class (the graph extra: pip install "thekaveh-nnx[graph]").
Details
``sampler="neighbor"`` (default) builds one ``NeighborLoader`` per split
from the graph's ``train_mask`` / ``val_mask`` / ``test_mask``;
``sampler="full"`` yields the whole graph as one batch per split (no
pyg-lib / torch-sparse needed).
``batch_sizes`` is a ``(train, val, test)`` tuple of *seed-node* counts
for the neighbor sampler. ``None`` (the default for every slot) means
*every node of the split mask in one batch* — one optimizer step per
epoch for the default train loader; pass an explicit train size such as
``batch_sizes=(256, None, None)`` for mini-batches. A positive integer
is kept verbatim; NumPy integers are normalized to ``int``. Zero,
``False``, negatives, floats, strings and anything but a 3-tuple are
rejected with the ``batch_sizes[i] (split)`` slot named before
``ds_class`` is instantiated. ``sampler="full"`` rejects every explicit
size (the complete split is always one batch).
An empty ``val_mask`` / ``test_mask`` is an *absent* split in both
sampler modes: that loader is ``None`` — the optional-loader
contract the tabular / preference wrappers already follow, which
``NNModel.train`` / ``Trainer.train`` honour by skipping validation —
its resolved size is ``0`` (derived from absence, never an accepted
explicit zero) and ``state()`` reports ``"0"``. An empty ``train_mask``
raises ``ValueError`` at construction. The caller's masks are only read.
nnx.nn.dataset.nn_tabular_dataset.NNTabularDataset¶
class nnx.nn.dataset.nn_tabular_dataset.NNTabularDataset(*, df: 'pd.DataFrame', feature_cols: 'list[str]', target_col: 'str', 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' = torch.float32, target_dtype: 'Optional[torch.dtype]' = None, seed: 'Optional[int]' = None, split: 'Optional[SplitManifest]' = None, id_col: 'Optional[str]' = None, source_identity: 'Optional[Union[IdentityRef, str]]' = None, standardize: 'Union[bool, Standardizer]' = False) -> 'None'
Wrap a pandas DataFrame as train/val/test DataLoaders.
Details
`feature_cols` columns are stacked into the input tensor; `target_col`
is the target column. By default, targets are coerced to int64 (long)
and validated as contiguous integer classes 0..K-1 (classification);
the loaders yield 1-D class-index targets `(batch,)` (the
`CrossEntropyLoss` convention). Set `target_dtype` to a floating-point
dtype (e.g. `torch.float32`) to skip the integer cast and contiguity
check and fix `output_dim=1` for regression; the loaders then yield
targets of shape `(batch, 1)` so they line up with a model whose
final linear layer has one output. Integer dtypes are rejected —
leave `target_dtype` unset (`None`) for classification.
``batch_sizes`` is a ``(train, val, test)`` tuple. ``None`` (the
default for every slot) means *one batch holding the complete split* —
the default train loader yields one full-split batch per epoch, i.e.
one optimizer step per epoch; pass an explicit train size such as
``batch_sizes=(64, None, None)`` for mini-batches. A positive integer
is kept verbatim (larger than the split → one smaller batch); NumPy
integers are normalized to ``int``. Zero, ``False``, negatives, floats,
strings and anything but a 3-tuple are rejected with the
``batch_sizes[i] (split)`` slot named before any tensor conversion or
split. Zero never disables a split: `val_proportion=0.0` /
`test_proportion=0.0` do, and that split's loader is then ``None``. Its
resolved ``batch_sizes`` slot is a placeholder ``1`` when it was
requested as ``None``; an explicit size is kept.
Admission: only the selected feature and target columns are inspected.
NaN in any of them, ``±inf`` in a feature (checked in source precision,
before the ``feature_dtype`` cast that could absorb it), finite features
outside an integer ``feature_dtype``'s range (the cast would wrap them)
and a non-finite target raise ``ValueError`` naming the columns and
dtypes; so does a finite value that overflows a narrower floating /
complex ``feature_dtype`` or floating ``target_dtype`` during conversion.
All checks, the classification label check included, run before the
split, so a rejection leaves the DataFrame and the global RNG untouched
(the one exception: with ``standardize``, a standardized feature that
overflows ``feature_dtype`` can only be detected after the split, once
the training statistics exist).
Explicit membership: ``split=`` takes a ``SplitManifest``
(``nnx.data_splits``) instead of the seeded ``random_split``. Rows are
matched to the plan by the sample ids in the ``id_col`` column (required;
never inferred from the index), so reordered rows keep their split; a
positional manifest matches row positions and needs its recorded
``source_identity``. Duplicate, missing or unexpected ids, a changed
``source_identity`` and an empty train split raise ``SplitError`` before
any tensor or loader exists. Only the planned rows are admitted and
counted (``output_dim`` included): rows the plan excluded (a
chronological ``gap``) may be absent and are never inspected. An empty
validation or test membership gives a ``None`` loader. ``seed``,
``val_proportion`` and ``test_proportion`` do not apply and must be left
at their defaults; the train loader still shuffles, and validation /
test follow the manifest's order. ``state()`` gains ``split`` (the
manifest digest) only then.
Train-only standardization: ``standardize=True`` fits a
``nnx.preprocessing.Standardizer`` on the training split's rows of the
feature columns only — after the split, so validation and test values
never reach the statistics — and applies those frozen statistics to
every split. ``standardize=<fitted Standardizer>`` reuses one without
refitting; its columns must equal ``feature_cols`` (same order) and its
dtype ``feature_dtype``, checked before the split. Either way the
standardizer in use is ``ds.standardizer`` (save it to serve the model).
Standardizing needs a floating ``feature_dtype`` and unique
``feature_cols``; the raw features are never converted, targets are never
transformed, and the split membership (and RNG draw) is the same as
without standardization. ``state()`` gains ``standardizer`` (its digest)
only then.
nnx.nn.dataset.nn_preference_dataset.NNPreferenceDataset¶
class nnx.nn.dataset.nn_preference_dataset.NNPreferenceDataset(*, prompts: 'list[str]', chosen: 'list[str]', rejected: 'list[str]', tokenizer: 'object', 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]' = None) -> 'None'
Wrap parallel lists of (prompt, chosen, rejected) strings as DPO loaders.
Details
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``.
``pad_token_id`` defaults to ``0``. Under :func:`train_bpe`'s default
special tokens, id 0 is ``"<unk>"`` and ``"<pad>"`` is id 1: pass the
tokenizer's pad id (``tokenizer.token_to_id("<pad>")``), or the DPO
step's padding mask would also hide every real ``<unk>`` token.
``batch_sizes`` is a ``(train, val, test)`` tuple. ``None`` (the
default for every slot) means *one batch holding the complete split* —
one DPO step per epoch for the default train loader; pass an explicit
train size such as ``batch_sizes=(8, None, None)`` for mini-batches. A
positive integer is kept verbatim (larger than the split → one smaller
batch); NumPy integers are normalized to ``int``. Zero, ``False``,
negatives, floats, strings and anything but a 3-tuple are rejected with
the ``batch_sizes[i] (split)`` slot named before ``tokenizer.encode``
is called. Zero never disables a split: `val_proportion=0.0` /
`test_proportion=0.0` do (that loader is then ``None``; its resolved
``batch_sizes`` slot is a placeholder ``1`` when requested as ``None``,
and an explicit size is kept).
6. Enums¶
nnx.nn.enum.activations.Activations¶
class nnx.nn.enum.activations.Activations(Enum)
Enum values: ELU, SELU, TANH, RELU, SOFTMAX, SIGMOID, SOFTPLUS, LEAKY_RELU.
nnx.nn.enum.activations.Activations.ELU¶
nnx.nn.enum.activations.Activations.ELU = 'elu'
Enum value elu.
nnx.nn.enum.activations.Activations.SELU¶
nnx.nn.enum.activations.Activations.SELU = 'selu'
Enum value selu.
nnx.nn.enum.activations.Activations.TANH¶
nnx.nn.enum.activations.Activations.TANH = 'tanh'
Enum value tanh.
nnx.nn.enum.activations.Activations.RELU¶
nnx.nn.enum.activations.Activations.RELU = 'relu'
Enum value relu.
nnx.nn.enum.activations.Activations.SOFTMAX¶
nnx.nn.enum.activations.Activations.SOFTMAX = 'softmax'
Enum value softmax.
nnx.nn.enum.activations.Activations.SIGMOID¶
nnx.nn.enum.activations.Activations.SIGMOID = 'sigmoid'
Enum value sigmoid.
nnx.nn.enum.activations.Activations.SOFTPLUS¶
nnx.nn.enum.activations.Activations.SOFTPLUS = 'softplus'
Enum value softplus.
nnx.nn.enum.activations.Activations.LEAKY_RELU¶
nnx.nn.enum.activations.Activations.LEAKY_RELU = 'leaky_relu'
Enum value leaky_relu.
nnx.nn.enum.checkpoints.Checkpoints¶
class nnx.nn.enum.checkpoints.Checkpoints(Enum)
Enum values: Q1, Q2, Q3, BEST, LAST, FIRST.
nnx.nn.enum.checkpoints.Checkpoints.Q1¶
nnx.nn.enum.checkpoints.Checkpoints.Q1 = 'q1'
Enum value q1.
nnx.nn.enum.checkpoints.Checkpoints.Q2¶
nnx.nn.enum.checkpoints.Checkpoints.Q2 = 'q2'
Enum value q2.
nnx.nn.enum.checkpoints.Checkpoints.Q3¶
nnx.nn.enum.checkpoints.Checkpoints.Q3 = 'q3'
Enum value q3.
nnx.nn.enum.checkpoints.Checkpoints.BEST¶
nnx.nn.enum.checkpoints.Checkpoints.BEST = 'best'
Enum value best.
nnx.nn.enum.checkpoints.Checkpoints.LAST¶
nnx.nn.enum.checkpoints.Checkpoints.LAST = 'last'
Enum value last.
nnx.nn.enum.checkpoints.Checkpoints.FIRST¶
nnx.nn.enum.checkpoints.Checkpoints.FIRST = 'first'
Enum value first.
nnx.nn.enum.devices.Devices¶
class nnx.nn.enum.devices.Devices(Enum)
Enum values: CPU, MPS, CUDA.
nnx.nn.enum.devices.Devices.CPU¶
nnx.nn.enum.devices.Devices.CPU = 'cpu'
Enum value cpu.
nnx.nn.enum.devices.Devices.MPS¶
nnx.nn.enum.devices.Devices.MPS = 'mps'
Enum value mps.
nnx.nn.enum.devices.Devices.CUDA¶
nnx.nn.enum.devices.Devices.CUDA = 'cuda'
Enum value cuda.
nnx.nn.enum.devices.Devices.torch_device¶
nnx.nn.enum.devices.Devices.torch_device(self) -> 'torch.device'
Explicit alias for self() — more readable in code that mixes the enum and torch.device usage.
nnx.nn.enum.devices.Devices.get¶
nnx.nn.enum.devices.Devices.get() -> 'Devices'
No public description is currently available.
nnx.nn.enum.devices.Devices.get_torch_device¶
nnx.nn.enum.devices.Devices.get_torch_device() -> 'torch.device'
Convenience: auto-detect and return the corresponding torch.device directly. Equivalent to Devices.get().torch_device().
nnx.nn.enum.losses.Losses¶
class nnx.nn.enum.losses.Losses(Enum)
Enum values: CROSS_ENTROPY, MEAN_SQUARED_ERROR, BINARY_CROSS_ENTROPY, NEGATIVE_LOG_LIKELIHOOD.
nnx.nn.enum.losses.Losses.CROSS_ENTROPY¶
nnx.nn.enum.losses.Losses.CROSS_ENTROPY = 'cross_entropy'
Enum value cross_entropy.
nnx.nn.enum.losses.Losses.MEAN_SQUARED_ERROR¶
nnx.nn.enum.losses.Losses.MEAN_SQUARED_ERROR = 'mean_squared_error'
Enum value mean_squared_error.
nnx.nn.enum.losses.Losses.BINARY_CROSS_ENTROPY¶
nnx.nn.enum.losses.Losses.BINARY_CROSS_ENTROPY = 'binary_cross_entropy'
Enum value binary_cross_entropy.
nnx.nn.enum.losses.Losses.NEGATIVE_LOG_LIKELIHOOD¶
nnx.nn.enum.losses.Losses.NEGATIVE_LOG_LIKELIHOOD = 'negative_log_likelihood'
Enum value negative_log_likelihood.
nnx.nn.enum.nets.Nets¶
class nnx.nn.enum.nets.Nets(Enum)
Enum values: CONV, FEED_FWD, FEED_FWD_MOE, GRAPH_ATT, GRAPH_CONV, GRAPH_SAGE, TRANSFORMER, VIT.
nnx.nn.enum.nets.Nets.CONV¶
nnx.nn.enum.nets.Nets.CONV = 'conv'
Enum value conv.
nnx.nn.enum.nets.Nets.FEED_FWD¶
nnx.nn.enum.nets.Nets.FEED_FWD = 'feed_fwd'
Enum value feed_fwd.
nnx.nn.enum.nets.Nets.FEED_FWD_MOE¶
nnx.nn.enum.nets.Nets.FEED_FWD_MOE = 'feed_fwd_moe'
Enum value feed_fwd_moe.
nnx.nn.enum.nets.Nets.GRAPH_ATT¶
nnx.nn.enum.nets.Nets.GRAPH_ATT = 'graph_att'
Enum value graph_att.
nnx.nn.enum.nets.Nets.GRAPH_CONV¶
nnx.nn.enum.nets.Nets.GRAPH_CONV = 'graph_conv'
Enum value graph_conv.
nnx.nn.enum.nets.Nets.GRAPH_SAGE¶
nnx.nn.enum.nets.Nets.GRAPH_SAGE = 'graph_sage'
Enum value graph_sage.
nnx.nn.enum.nets.Nets.TRANSFORMER¶
nnx.nn.enum.nets.Nets.TRANSFORMER = 'transformer'
Enum value transformer.
nnx.nn.enum.nets.Nets.VIT¶
nnx.nn.enum.nets.Nets.VIT = 'vit'
Enum value vit.
nnx.nn.enum.optims.Optims¶
class nnx.nn.enum.optims.Optims(Enum)
Enum values: SGD, ADAM, ADAM_AMSGRAD, SGD_NESTEROV, ADAMW.
nnx.nn.enum.optims.Optims.SGD¶
nnx.nn.enum.optims.Optims.SGD = 'sgd'
Enum value sgd.
nnx.nn.enum.optims.Optims.ADAM¶
nnx.nn.enum.optims.Optims.ADAM = 'adam'
Enum value adam.
nnx.nn.enum.optims.Optims.ADAM_AMSGRAD¶
nnx.nn.enum.optims.Optims.ADAM_AMSGRAD = 'adam_amsgrad'
Enum value adam_amsgrad.
nnx.nn.enum.optims.Optims.SGD_NESTEROV¶
nnx.nn.enum.optims.Optims.SGD_NESTEROV = 'sgd_nesterov'
Enum value sgd_nesterov.
nnx.nn.enum.optims.Optims.ADAMW¶
nnx.nn.enum.optims.Optims.ADAMW = 'adamw'
Enum value adamw.
nnx.nn.enum.schedulers.Schedulers¶
class nnx.nn.enum.schedulers.Schedulers(Enum)
Enum values: REDUCE_LR_ON_PLATEAU, STEP, COSINE_ANNEALING, ONE_CYCLE, LINEAR_WARMUP_DECAY.
nnx.nn.enum.schedulers.Schedulers.REDUCE_LR_ON_PLATEAU¶
nnx.nn.enum.schedulers.Schedulers.REDUCE_LR_ON_PLATEAU = 'reduce_lr_on_plateau'
Enum value reduce_lr_on_plateau.
nnx.nn.enum.schedulers.Schedulers.STEP¶
nnx.nn.enum.schedulers.Schedulers.STEP = 'step'
Enum value step.
nnx.nn.enum.schedulers.Schedulers.COSINE_ANNEALING¶
nnx.nn.enum.schedulers.Schedulers.COSINE_ANNEALING = 'cosine_annealing'
Enum value cosine_annealing.
nnx.nn.enum.schedulers.Schedulers.ONE_CYCLE¶
nnx.nn.enum.schedulers.Schedulers.ONE_CYCLE = 'one_cycle'
Enum value one_cycle.
nnx.nn.enum.schedulers.Schedulers.LINEAR_WARMUP_DECAY¶
nnx.nn.enum.schedulers.Schedulers.LINEAR_WARMUP_DECAY = 'linear_warmup_decay'
Enum value linear_warmup_decay.
7. Callbacks¶
nnx.nn.callbacks.Callback¶
class nnx.nn.callbacks.Callback()
Base class for training callbacks. Override any subset of the hooks.
Details
``distributed`` declares the callback's rank behaviour under
``train(distributed=DDP())``: ``"all"`` — run on every rank
(it must not write files: records and decisions are identical on every
rank); ``"writer"`` — run on the writer rank only (it may write; it must
not change control flow such as ``ctx.should_stop``); ``"unsupported"``
— refused (it writes when constructed: use
``nnx.distributed.writer_only``). ``None`` (undeclared) is refused under
DDP and ignored otherwise.
nnx.nn.callbacks.Callback.on_train_begin¶
nnx.nn.callbacks.Callback.on_train_begin(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.Callback.on_epoch_begin¶
nnx.nn.callbacks.Callback.on_epoch_begin(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.Callback.on_epoch_end¶
nnx.nn.callbacks.Callback.on_epoch_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.Callback.on_train_end¶
nnx.nn.callbacks.Callback.on_train_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.Callback.on_optimizer_update¶
nnx.nn.callbacks.Callback.on_optimizer_update(self, ctx: '_CallbackContext', event: 'Any') -> 'None'
Called once per committed optimizer update of an objective run — never per microbatch, for an all-masked window or a skipped step. event is a detached nnx.objectives.UpdateEvent (optimizer name, its update index, epoch / batch, microbatches, loss values); ctx.update_count is the run's committed-update count.
nnx.nn.callbacks.Callback.checkpoint_transforms¶
nnx.nn.callbacks.Callback.checkpoint_transforms(self) -> 'tuple[NNCheckpointTransform, ...]'
Completed topology transforms to persist on the final checkpoint.
nnx.nn.callbacks.EarlyStopping¶
class nnx.nn.callbacks.EarlyStopping(monitor: 'Optional[Union[str, MonitorSpec]]' = None, patience: 'int' = 10, min_delta: 'float' = 0.0, mode: 'str' = 'min', name: 'Optional[str]' = None)
Stop training when the monitored metric stops improving.
Details
Args:
monitor: which data-point field to track. ``None`` (default) selects
automatically from the validation data point, once per
``train()`` call: ``val_edp.error`` when the first validated
epoch reports a finite error, otherwise ``val_edp.loss`` when it
reports a finite loss (regression and other evaluators that
leave ``error`` as ``None``). The choice then stays fixed for
that run; ``selected_monitor`` reports it. The default never
reads training metrics and only supports ``mode="min"``.
An explicit ``"val_edp.error"``, ``"val_edp.loss"``,
``"train_edp.error"`` or ``"train_edp.loss"`` reads exactly that
field with no fallback. Unlike BEST selection and
ReduceLROnPlateau, there is no validation→training fallback.
When the tracked field (or its whole data point) is absent, the
epoch is not counted toward patience and one ``RuntimeWarning``
per run names the monitor and what is missing, instead of the
callback going silently inactive. A NaN/±inf value never
becomes the best and counts as an epoch without improvement
(also reported once per run).
patience: epochs with no improvement before stopping — a nonnegative
integer count (NumPy integers accepted and normalized; zero
stops on the first non-improving epoch). Fractional, boolean
or string values raise ``ValueError`` at construction.
min_delta: minimum change to qualify as improvement.
mode: improvement direction for the monitored field. ``"min"``
(default): lower is better — the meaning of every accepted
monitor (loss / error), and the direction BEST selection and
ReduceLROnPlateau also assume, so it is almost always right.
``"max"`` only reverses this callback's comparison for one of
the four accepted keys; it does not enable accuracy/F1
monitors, which are rejected, and it does not change how the
rest of NNx ranks ``error`` / ``loss``. ``"max"`` requires an
explicit ``monitor``.
**Named monitors.** ``monitor`` may instead be a
:class:`~nnx.MonitorSpec` — a split and metric (``loss``,
``error`` or a metric declared in ``NNTrainParams.metrics``)
with its own direction, ``min_delta`` and missing / non-finite
policy. It reads the epoch's whole-epoch training summary or
validation record and applies
:meth:`~nnx.MonitorSpec.improved`, the rule BEST selection and
plateau scheduling use when the run declares the same monitor,
so all three agree. ``mode`` and ``min_delta`` must then stay
at their defaults (set them on the spec).
name: component name under which the patience state is checkpointed.
By default ``"early_stopping"``, numbered
(``early_stopping.2`` …) in callback order when a run has
several; an explicit name must be unique within the run.
Checkpointable: the best value, the epochs waited and the selected
monitor are saved with every checkpoint's training state and restored
on a stateful warm resume *after* ``on_train_begin`` resets them, so a
resumed run stops at the same epoch an uninterrupted one would. It is
an optional component: resuming from a checkpoint written without it
starts with fresh patience.
Example::
EarlyStopping(monitor="val_edp.loss", mode="min", patience=5)
nnx.nn.callbacks.EarlyStopping.selected_monitor¶
property nnx.nn.callbacks.EarlyStopping.selected_monitor
Field compared in the current run (None until the automatic default has seen a validated epoch with a finite error or loss).
nnx.nn.callbacks.EarlyStopping.component_spec¶
nnx.nn.callbacks.EarlyStopping.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.nn.callbacks.EarlyStopping.component_state¶
nnx.nn.callbacks.EarlyStopping.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.nn.callbacks.EarlyStopping.check_component_state¶
nnx.nn.callbacks.EarlyStopping.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
Reject, before anything is restored, patience tracked for a different monitor or direction.
nnx.nn.callbacks.EarlyStopping.load_component_state¶
nnx.nn.callbacks.EarlyStopping.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.nn.callbacks.EarlyStopping.on_train_begin¶
nnx.nn.callbacks.EarlyStopping.on_train_begin(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.EarlyStopping.on_epoch_end¶
nnx.nn.callbacks.EarlyStopping.on_epoch_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.LRMonitor¶
class nnx.nn.callbacks.LRMonitor(bounded: 'bool' = True)
Logs the current LR each epoch. History exposed at .history.
Details
`.history` keeps one entry per epoch (the primary optimizer's LR at the
epoch's end) whatever the scheduler's clock. An ``optimizer_update``-clock
scheduler also records `.update_history`: ``(update index, LR
after that update's scheduler step)`` for the primary optimizer, a
separate trace of the steps NNx takes — empty on the epoch clock and
under ``Trainer``'s ``auto_step_schedulers=False``, where the step
function steps the schedule itself.
In a run with a history journal (``nnx.history``) the log is
bounded too: it keeps the LRs of the last ``retention`` epochs (the
journal's bound, applied per epoch), and `.update_history` the updates
of those epochs. ``bounded=False`` keeps every epoch's LR (one float
each), and every update's, in any run.
nnx.nn.callbacks.LRMonitor.on_epoch_end¶
nnx.nn.callbacks.LRMonitor.on_epoch_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.ModelCheckpoint¶
class nnx.nn.callbacks.ModelCheckpoint(epochs: 'Optional[list[int]]' = None, tag: 'str' = 'custom')
Save a custom-tagged checkpoint at user-specified epochs.
Details
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.
Files are **weights-only** and say so (``training_state_present`` is
``False``): pass ``resume_from_checkpoint="<tag>_e<epoch>"`` with
``resume_mode="weights_only"`` (or the default ``"auto"``) to
warm-start from one; ``resume_mode="stateful"`` rejects it before
anything is restored.
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"``.
nnx.nn.callbacks.ModelCheckpoint.on_epoch_end¶
nnx.nn.callbacks.ModelCheckpoint.on_epoch_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.TensorBoardCallback¶
class nnx.nn.callbacks.TensorBoardCallback(log_dir: 'Optional[str]' = None, flush_each_epoch: 'bool' = True)
Stream train/val metrics + LR to a TensorBoard SummaryWriter.
Details
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.
nnx.nn.callbacks.TensorBoardCallback.on_epoch_end¶
nnx.nn.callbacks.TensorBoardCallback.on_epoch_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.TensorBoardCallback.on_train_end¶
nnx.nn.callbacks.TensorBoardCallback.on_train_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.WandbCallback¶
class nnx.nn.callbacks.WandbCallback(project: 'Optional[str]' = None, wandb_run=None, **init_kwargs)
Stream train/val metrics + LR to Weights & Biases.
Details
Requires `wandb` — lazily imported. Pass `project=` to start a new run,
or `wandb_run=` to attach to an externally-managed run.
nnx.nn.callbacks.WandbCallback.on_epoch_end¶
nnx.nn.callbacks.WandbCallback.on_epoch_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
nnx.nn.callbacks.WandbCallback.on_train_end¶
nnx.nn.callbacks.WandbCallback.on_train_end(self, ctx: '_CallbackContext') -> 'None'
No public description is currently available.
8. Fine-tuning (nnx.finetune)¶
nnx.finetune.freezing.freeze¶
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.
Details
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.
nnx.finetune.freezing.unfreeze¶
nnx.finetune.freezing.unfreeze(module: 'nn.Module', *patterns: 'str') -> 'int'
Mirror of freeze — set requires_grad=True on matching parameters. Returns the count newly unfrozen.
nnx.finetune.freezing.frozen¶
nnx.finetune.freezing.frozen(module: 'nn.Module') -> 'list[str]'
List the dotted parameter names currently frozen under module.
Details
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.
nnx.finetune.loading.load_pretrained¶
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.
Details
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). Either way a key both sides hold
with different shapes raises ``RuntimeError`` before anything
is copied.
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.
nnx.finetune.loading.LoadPretrainedResult¶
class nnx.finetune.loading.LoadPretrainedResult(loaded_keys: 'list[str]', missing_keys: 'list[str]', unexpected_keys: 'list[str]') -> 'None'
Outcome of a load_pretrained call.
Details
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.
nnx.finetune.param_groups.NNParamGroupSpec¶
class nnx.finetune.param_groups.NNParamGroupSpec(*, name_pattern: 'str', lr: 'Optional[float]' = None, lr_multiplier: 'Optional[float]' = None, weight_decay: 'Optional[float]' = None) -> 'None'
One row in NNOptimParams.param_groups.
Details
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.
Specs are matched in list order and the **first** matching spec wins —
rules never merge. Put specific rules before broad ones: a parameter
that should get settings from two rules needs its own combined rule.
Example:
# Freeze nothing, but train the encoder at 1/100th the head's LR
# and disable weight_decay on every bias term. `encoder.*bias`
# must come first: with `encoder.*` first, the encoder biases
# would match it and keep weight decay.
NNOptimParams(
name=Optims.ADAM,
max_lr=1e-3,
momentum=(0.9, 0.999),
weight_decay=5e-4,
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),
],
)
# encoder.bias: lr 1e-5, wd 0 encoder.weight: lr 1e-5, wd 5e-4
# head.bias: lr 1e-3, wd 0 head.weight: lr 1e-3, wd 5e-4
nnx.finetune.param_groups.NNParamGroupSpec.state¶
nnx.finetune.param_groups.NNParamGroupSpec.state(self) -> 'dict'
No public description is currently available.
nnx.finetune.param_groups.NNParamGroupSpec.from_state¶
nnx.finetune.param_groups.NNParamGroupSpec.from_state(state: 'dict') -> 'NNParamGroupSpec'
No public description is currently available.
nnx.finetune.param_groups.build_param_groups¶
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.
Details
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; later matches are ignored, never merged, so
list specific patterns (``encoder.*bias``) before broad ones
(``encoder.*``, ``*.bias``).
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.
9. Parameter-efficient fine-tuning (nnx.peft)¶
LoRA + DoRA + IA3 + Prefix-Tuning + Prompt-Tuning + Adapters. The apply_*_to helpers wrap layers in place. LoRA, IA3, Prefix and Prompt ship save_*_weights / load_*_weights, which persist only the trainable delta. DoRA persists with the full state_dict() (the LoRA helpers omit magnitude), and AdapterLayer is inserted by hand.
9.1. LoRA¶
nnx.peft.lora.LoRALinear¶
class nnx.peft.lora.LoRALinear(base: 'nn.Linear', *, r: 'int' = 8, alpha: 'float' = 16.0, dropout: 'float' = 0.0)
Linear layer wrapped with a LoRA low-rank residual.
Details
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.
``lora_A`` / ``lora_B`` are allocated with the base weight's dtype
and device, so convert or move the base *before* wrapping and the
adapter composes immediately (a meta-device base yields meta
adapters). The base is never moved, recast or replaced.
The wrapper and its dropout inherit the base layer's train/eval mode
at construction, so wrapping an eval layer keeps inference
deterministic; a later ``.train()`` / ``.eval()`` switches them as
usual. Modes are runtime state and never serialized.
nnx.peft.lora.LoRALinear.in_features¶
property nnx.peft.lora.LoRALinear.in_features
No public description is currently available.
nnx.peft.lora.LoRALinear.out_features¶
property nnx.peft.lora.LoRALinear.out_features
No public description is currently available.
nnx.peft.lora.LoRALinear.forward¶
nnx.peft.lora.LoRALinear.forward(self, x: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.peft.lora.LoRALinear.extra_repr¶
nnx.peft.lora.LoRALinear.extra_repr(self) -> 'str'
No public description is currently available.
nnx.peft.lora.apply_lora_to¶
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 nn.Linear submodule whose dotted name matches any of name_patterns with a LoRALinear. Returns the number of layers wrapped.
Details
Patterns use shell-style globs (``fnmatch``) against the dotted
submodule name — e.g., ``"layers.0"``, ``"encoder.*"``, ``"*"`` for
every Linear. Every registration path is considered, including the
non-first path through a shared container.
**Shared modules:** a Linear inside a container that is
registered under several names has one registration slot; it is
wrapped once and every path sees the same wrapper. A Linear itself
registered under more than one name (``m.a = m.b = linear``) cannot be
wrapped without splitting it into independent layers, so selecting
any of its names raises ``ValueError`` naming every alias — after the
whole matched set is validated and before any wrapper is built, so a
rejected call modifies nothing. Tied tensors between distinct modules
are not aliases and are outside this check — but wrapping one of those
modules freezes the shared tensor for both (wrapping a tied output
head also freezes the tied token embedding).
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 — unique registration slots, so a
shared container's Linear counts once (may be 0 if no patterns
match).
Raises:
ValueError: if ``name_patterns`` is empty, if a selected Linear
is registered under more than one name, or if its parent reads
its weight directly (``nn.MultiheadAttention.out_proj``, a torch
Transformer layer's ``linear1`` / ``linear2``); nothing is
modified.
**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.
nnx.peft.lora.save_lora_weights¶
nnx.peft.lora.save_lora_weights(module: 'nn.Module', path: 'Union[str, Path]') -> 'str'
Save ONLY the LoRA parameters of module to path.
Details
The output is a plain ``torch.save`` of a dict-subset of the full
state_dict, containing exactly the ``lora_A`` / ``lora_B`` tensors
owned by the :class:`LoRALinear` wrappers in ``module`` — selected by
registered ownership, not by key substring, so a module *named*
``lora_A_projection`` never leaks its frozen base weights. Loadable
via :func:`load_lora_weights`. An adapter-only export records no
transformation recipe (``nnx.transforms``): it cannot rebuild the
wrapped topology alone — apply the same adapters (or materialize the
recipe) first.
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).
nnx.peft.lora.load_lora_weights¶
nnx.peft.lora.load_lora_weights(module: 'nn.Module', source: 'Union[str, Path, dict]') -> 'int'
Load LoRA parameters into module from source.
Details
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. The source is
first restricted to the keys ``module``'s own LoRA wrappers register
(derived from the destination, not from the source's key names), so
a full checkpoint — or one carrying colliding ``*.base.*`` keys —
can never overwrite a frozen base tensor. A wrong-shaped owned key
raises a size-mismatch ``RuntimeError`` before anything is copied.
9.2. DoRA¶
nnx.peft.dora.DoRALinear¶
class nnx.peft.dora.DoRALinear(base: 'nn.Linear', *, r: 'int' = 8, alpha: 'float' = 16.0, dropout: 'float' = 0.0)
Linear layer wrapped with a DoRA weight decomposition.
Details
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.
The row normalization is carried out in FP32 for FP16/BF16 layers
(FP64 stays FP64), so a zero or tiny row of ``V`` — from pruning,
explicit zero initialization or a zero row in a loaded base — yields
a finite, zero output equal to the base instead of NaN, and learned
nonzero updates stay differentiable. The effective weight is cast
back to the layer dtype before the matmul, so the output dtype is
unchanged. Persist a DoRA adapter with the full ``state_dict()``:
the LoRA-only helpers omit ``magnitude``. Like :class:`LoRALinear`,
the wrapper and its update dropout inherit the base layer's
train/eval mode at construction.
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)``.
nnx.peft.dora.DoRALinear.forward¶
nnx.peft.dora.DoRALinear.forward(self, x: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.peft.dora.DoRALinear.extra_repr¶
nnx.peft.dora.DoRALinear.extra_repr(self) -> 'str'
No public description is currently available.
nnx.peft.dora.apply_dora_to¶
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 nn.Linear submodule whose dotted name matches any of name_patterns with a DoRALinear. Returns the number of layers wrapped.
Details
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 a selected Linear
is read directly by its parent (``nn.MultiheadAttention.out_proj``,
a torch Transformer layer's ``linear1`` / ``linear2``), or if it
is registered under more than one name — the whole matched set
is validated before any wrapper is built, so nothing is
modified. A Linear inside a shared container is one slot and is
wrapped once; tied tensors between distinct modules are outside
this check (see :func:`nnx.peft.apply_lora_to`).
TypeError: if a selected Linear is a subclass with its own
``forward`` (a QAT fake-quantized layer), which DoRA's recomposed
matmul would bypass. Every wrapper is built before any slot is
replaced, so a refused call modifies nothing.
9.3. IA3¶
nnx.peft.ia3.IA3Linear¶
class nnx.peft.ia3.IA3Linear(base: 'nn.Linear')
Linear layer wrapped with an IA3 per-output-dim scaling vector.
Details
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).
``scaling`` takes the base weight's dtype and device, so the output
keeps the base dtype (a half base stays half) and a pre-converted or
pre-moved base composes without a second ``.to()``.
The wrapper inherits the base layer's train/eval mode at construction
(modes are runtime state and never serialized).
Args:
base: the :class:`nn.Linear` to wrap.
nnx.peft.ia3.IA3Linear.in_features¶
property nnx.peft.ia3.IA3Linear.in_features
No public description is currently available.
nnx.peft.ia3.IA3Linear.out_features¶
property nnx.peft.ia3.IA3Linear.out_features
No public description is currently available.
nnx.peft.ia3.IA3Linear.forward¶
nnx.peft.ia3.IA3Linear.forward(self, x: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.peft.ia3.IA3Linear.extra_repr¶
nnx.peft.ia3.IA3Linear.extra_repr(self) -> 'str'
No public description is currently available.
nnx.peft.ia3.apply_ia3_to¶
nnx.peft.ia3.apply_ia3_to(module: 'nn.Module', *name_patterns: 'str') -> 'int'
Wrap every nn.Linear submodule whose dotted name matches any of name_patterns with an IA3Linear. Returns the number of layers wrapped.
Details
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 a selected Linear
is read directly by its parent (``nn.MultiheadAttention.out_proj``,
a torch Transformer layer's ``linear1`` / ``linear2``), or if it
is registered under more than one name — the whole matched set
is validated before any wrapper is built, so nothing is
modified. A Linear inside a shared container is one slot and is
wrapped once; tied tensors between distinct modules are outside
this check (see :func:`nnx.peft.apply_lora_to`). Also raised,
before any change, when a selected Linear is the base of a
:class:`DoRALinear` (DoRA reads its base's weight directly).
nnx.peft.ia3.save_ia3_weights¶
nnx.peft.ia3.save_ia3_weights(module: 'nn.Module', path: 'Union[str, Path]') -> 'str'
Save ONLY the IA3 scaling parameters of module to path.
Details
The output is a ``torch.save`` of a dict-subset of the full
state_dict, containing exactly the ``scaling`` tensors owned by the
:class:`IA3Linear` wrappers in ``module`` (selected by registered
ownership, not by key substring). 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).
nnx.peft.ia3.load_ia3_weights¶
nnx.peft.ia3.load_ia3_weights(module: 'nn.Module', source: 'Union[str, Path, dict]') -> 'int'
Load IA3 scaling parameters into module from source.
Details
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.
9.4. Prefix Tuning¶
nnx.peft.prefix.PrefixTuner¶
class nnx.peft.prefix.PrefixTuner(model: 'TransformerNN', *, n_prefix: 'int' = 10, n_layers: 'Optional[int]' = None)
Wrap a TransformerNN with learnable per-layer K/V prefixes.
Details
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.
Each layer's prefix takes the dtype and device of that block's
``attn.w_qkv`` weight. A projection already adapted by
:class:`~nnx.peft.LoRALinear` / :class:`~nnx.peft.DoRALinear` /
:class:`~nnx.peft.IA3Linear` (``apply_*_to`` before wrapping) is
resolved to the :class:`nn.Linear` it wraps; the adapter stays in the
attention path, so its delta or scaling still shapes K/V, and its
parameters are frozen with the rest of the model.
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).
nnx.peft.prefix.PrefixTuner.forward¶
nnx.peft.prefix.PrefixTuner.forward(self, *args, **kwargs)
Delegate to the wrapped model. The prefix injection happens inside each block's monkey-patched MHA forward.
nnx.peft.prefix.PrefixTuner.trainable_parameters¶
nnx.peft.prefix.PrefixTuner.trainable_parameters(self) -> 'Iterator[nn.Parameter]'
Yield only the learned prefix tensors.
Details
The wrapped model's parameters are frozen on construction; this
is the iterator you hand to an optimizer.
nnx.peft.prefix.PrefixTuner.prefix_state_dict¶
nnx.peft.prefix.PrefixTuner.prefix_state_dict(self) -> 'dict'
Return a state-dict containing only the prefix tensors, keyed for round-trip via load_prefix_weights.
Details
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``, …
nnx.peft.prefix.save_prefix_weights¶
nnx.peft.prefix.save_prefix_weights(tuner: 'PrefixTuner', path: 'Union[str, Path]') -> 'str'
Save ONLY the prefix tensors of tuner to path.
Details
Args:
tuner: a :class:`PrefixTuner` instance.
path: destination file path.
Returns:
The path written, so calls can be chained.
nnx.peft.prefix.load_prefix_weights¶
nnx.peft.prefix.load_prefix_weights(tuner: 'PrefixTuner', source: 'Union[str, Path, dict]') -> 'int'
Load prefix tensors into tuner from source.
Details
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.
9.5. Prompt Tuning¶
nnx.peft.prompt.PromptTuner¶
class nnx.peft.prompt.PromptTuner(model: 'TransformerNN', *, n_prompt_tokens: 'int' = 20)
Wrap a TransformerNN with a learnable soft prompt.
Details
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 —
and allocated with the dtype/device of the wrapped model's token
embedding, so convert or move the model *before* wrapping.
nnx.peft.prompt.PromptTuner.effective_max_seq_len¶
property nnx.peft.prompt.PromptTuner.effective_max_seq_len
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.
nnx.peft.prompt.PromptTuner.forward¶
nnx.peft.prompt.PromptTuner.forward(self, tokens: 'torch.Tensor') -> 'torch.Tensor'
Run the wrapped model with the soft prompt prepended.
Details
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.
nnx.peft.prompt.PromptTuner.trainable_parameters¶
nnx.peft.prompt.PromptTuner.trainable_parameters(self) -> 'Iterator[nn.Parameter]'
Yield only the soft-prompt tensor.
Details
The wrapped model's parameters are frozen on construction; this
is the iterator you hand to an optimizer.
nnx.peft.prompt.PromptTuner.prompt_state_dict¶
nnx.peft.prompt.PromptTuner.prompt_state_dict(self) -> 'dict'
Return a state-dict containing only the soft-prompt tensor, keyed for round-trip via load_prompt_weights. Only the tuner-owned soft_prompt entry qualifies — never anything under model.*, even if a nested base module's name happens to contain the marker.
nnx.peft.prompt.save_prompt_weights¶
nnx.peft.prompt.save_prompt_weights(tuner: 'PromptTuner', path: 'Union[str, Path]') -> 'str'
Save ONLY the soft-prompt tensor of tuner to path.
Details
Args:
tuner: a :class:`PromptTuner` instance.
path: destination file path.
Returns:
The path written, so calls can be chained.
nnx.peft.prompt.load_prompt_weights¶
nnx.peft.prompt.load_prompt_weights(tuner: 'PromptTuner', source: 'Union[str, Path, dict]') -> 'int'
Load the soft-prompt tensor into tuner from source.
Details
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.
9.6. Adapters¶
nnx.peft.adapters.AdapterLayer¶
class nnx.peft.adapters.AdapterLayer(dim: 'int', bottleneck: 'int', activation: 'Callable[[], nn.Module]' = torch.nn.modules.activation.GELU)
Bottleneck residual block: y = x + up(act(down(x))).
Details
``up.weight`` and ``up.bias`` are zero-initialized so at step 0
the layer's output equals its input exactly. ``up`` receives a
gradient from the first step; ``down`` receives none until ``up``
moves away from zero (its gradient flows through ``up.weight``).
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.
nnx.peft.adapters.AdapterLayer.forward¶
nnx.peft.adapters.AdapterLayer.forward(self, x: 'torch.Tensor') -> 'torch.Tensor'
No public description is currently available.
nnx.peft.adapters.AdapterLayer.extra_repr¶
nnx.peft.adapters.AdapterLayer.extra_repr(self) -> 'str'
No public description is currently available.
10. Pruning (nnx.prune)¶
nnx.prune.magnitude.magnitude_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.
Details
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. ``NNModel.train``
refuses a registered net that still carries the masks: run
such a schedule on ``NNModel(module=net)`` or a custom loop.
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)``, or if
``bake=False`` selects a Linear whose parent reads its weight
directly (``nn.MultiheadAttention``'s ``out_proj``).
**Idempotency note:** with ``bake=True``, 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. With
``bake=False`` the masks compose: a second call prunes ``sparsity`` of
the remaining weights (0.5 twice leaves 0.75 zero).
nnx.prune.semi_structured.semi_structured_24¶
nnx.prune.semi_structured.semi_structured_24(net: 'nn.Module', *, layer_pattern: 'str' = '*') -> 'int'
Swap each matched nn.Linear's weight with a 2:4 semi-structured sparse tensor via torchao.sparsity.sparsify_.
Details
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, and the conversion encodes each group
of 4 by the *positions* of its nonzeros, not by magnitude: extra
nonzeros are dropped by position. Prune the weight to a valid 2:4
pattern first (a custom 2:4 mask; :func:`magnitude_prune` alone gives
no per-group structure).
**Inference only:** torchao replaces each matched weight with a
``requires_grad=False`` parameter, so the matched Linears stop
training after this call.
11. Model surgery (nnx.surgery)¶
Walkthrough at Model surgery. Every primitive returns a fresh module (expand_embedding also returns a frozen-row mask) and composes with NNModel.train() for the "load checkpoint → surgery → refine" loop.
nnx.surgery.widen.widen¶
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.
Details
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 whose
consumer can be *proven*: the target lives directly inside an
``nn.Sequential`` (the following siblings are walked) or in
``FeedFwdNN.layers`` (the effective per-layer activation /
dropout is resolved), and every op between it and the next
Linear is elementwise — the built-in activations except
Softmax, ``Dropout``, ``Identity``. Anything else raises
before any allocation (see ``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) — in eval mode when
a ``Dropout`` sits between the target and its consumer, since a
stochastic training forward is not identical by construction.
Raises:
KeyError: if ``layer_name`` is not a submodule of ``model``.
TypeError: if the named submodule is not :class:`nn.Linear` (or a
subclass with its own ``forward``).
ValueError: if ``new_width`` is not an integer strictly greater
than the current ``out_features``; if the target or a module on
its path carries a parametrization, a legacy ``weight_norm`` /
``spectral_norm`` hook or an unbaked prune mask; if no plain
downstream Linear exists; if
a width-dependent op (Softmax, LayerNorm, BatchNorm, or any
module outside the elementwise allowlist) sits between the
target and its consumer; if the target lives in a container
other than ``nn.Sequential`` / ``FeedFwdNN.layers`` (module
registration order is not data flow), or in a subclass with
its own forward or a container with hooks; if a module on the
path from the target to its consumer has hooks; or if the
target or its consumer is registered under more than one
path (an alias used in several places). Validation runs on the source
before anything is copied or allocated, so a rejected call
leaves the model, its module identities and the RNG untouched.
nnx.surgery.deepen.deepen¶
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.
Details
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`
(addressed by position in a numeric container or by its
dotted key in a named one) — the primitive splices
``Linear(I) → ReLU`` in after it, keeping every original
key and every repeated slot, and registering the new
modules as ``_nnx_deepen_linear_N`` / ``_nnx_deepen_relu_N``
in a named container.
* an :class:`nn.Linear` inside a parent :class:`nn.ModuleList`
whose grandparent module applies ReLU at that site (the
FeedFwdNN contract; ``params.activation_for(idx)`` is
consulted, so per-layer overrides are honoured) — the
primitive inserts a new identity-init Linear into the
ModuleList right after and rebuilds the grandparent's
``params`` so ``hidden_dims`` / ``activations`` /
``dropout_probs`` describe the deeper topology (the new
site gets dropout ``0.0``). Adopt ``deeper.params`` as the
``net_params`` of a fresh ``NNModel`` when rebuilding.
Returns:
A fresh :class:`nn.Module` whose forward output matches the
original within ``atol=1e-5`` (eval mode when dropout is
configured; a seeded training forward also matches, because the
zero-dropout site draws nothing from the RNG).
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 (a non-ReLU activation in a
Sequential is this case), or is a ReLU subclass with its own
``forward``.
ValueError: if the activation effective at a FeedFwdNN insertion
site is anything other than ReLU (sigmoid / tanh / GELU break
function-preservation); if the ReLU is not fed by an
``nn.Linear`` through elementwise / norm modules only; if the
Sequential is a subclass or carries its own parameters, buffers
or hooks, or the target ReLU has hooks; or if a ModuleList is not
the ``layers`` / ``fcs`` of a stock FeedFwdNN / ConvNN.
nnx.surgery.drop_layer.drop_layer¶
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 nn.Identity.
Details
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``, or a candidate scores NaN or infinite.
nnx.surgery.low_rank.low_rank_factorize¶
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.
Details
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` (or a subclass
with its own ``forward``).
ValueError: if ``rank`` is out of range, ``method`` unknown, or
``linear`` carries a parametrization, a legacy ``weight_norm`` /
``spectral_norm`` hook or an unbaked prune mask.
nnx.surgery.embedding.expand_embedding¶
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.
Details
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` (or a subclass
with its own ``forward``).
ValueError: if ``new_num_embeddings`` is not an integer strictly greater,
if ``init`` is unknown, or if ``emb`` carries a parametrization, a
legacy ``weight_norm`` / ``spectral_norm`` hook or an unbaked prune
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¶
nnx.quantize.ptq.quantize_int8(model: '_M') -> '_M'
Return a new instance of type(model) with int8 weight-only quantized net.
Details
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 wrapper is a shallow
copy of ``model`` — the same class (a :class:`GenerativeNNModel` keeps
its tokenizer and ``generate``; a user subclass keeps its overrides and
custom attributes) sharing every other attribute (``params``,
``net_params``, ``device``, ``loss_fn``, ``tokenizer``) with the
original — only ``net`` is the independent quantized copy. No
constructor is re-run, so no weights are re-initialized and no RNG
is consumed; subclasses that define their own copy hooks are copied
through them.
Args:
model: a trained :class:`NNModel` (or subclass). PTQ has no
training step; this function is a pure post-process.
Returns:
a new instance of ``type(model)`` whose ``net`` is the quantized
deep-copy of ``model.net``. The new model can be used for
``predict`` / ``evaluate`` / ``generate`` / ``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]``.
nnx.quantize.qat.qat_train_step_factory¶
nnx.quantize.qat.qat_train_step_factory(base_step: 'Optional[TrainStepFn]' = None, qat_config: 'str' = '8da4w') -> 'TrainStepFn'
Return a TrainStepFn that runs base_step against a fake-quantized model.
Details
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.
nnx.quantize.qat.QATLifecycleCallback¶
class nnx.quantize.qat.QATLifecycleCallback(qat_config: 'str' = '8da4w', *, groupsize: 'int' = 32)
Manage the torchao prepare / convert lifecycle around training.
Details
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.
A completed conversion also contributes a versioned checkpoint transform
containing ``qat_config`` and ``groupsize``. The final ``LAST`` checkpoint
persists that recipe, so :meth:`NNModel.from_checkpoint` can rebuild the
converted torchao topology before loading its quantized tensors.
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.
nnx.quantize.qat.QATLifecycleCallback.on_train_begin¶
nnx.quantize.qat.QATLifecycleCallback.on_train_begin(self, ctx: '_CallbackContext') -> 'None'
Insert fake-quant ops into ctx.model.net in place.
nnx.quantize.qat.QATLifecycleCallback.on_train_end¶
nnx.quantize.qat.QATLifecycleCallback.on_train_end(self, ctx: '_CallbackContext') -> 'None'
Convert fake-quant ops in ctx.model.net to true int4/int8 modules.
Details
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.
nnx.quantize.qat.QATLifecycleCallback.checkpoint_transforms¶
nnx.quantize.qat.QATLifecycleCallback.checkpoint_transforms(self) -> 'tuple[NNCheckpointTransform, ...]'
Describe the completed conversion so checkpoint loaders can replay it.
13. Diffusion (nnx.diffusion)¶
nnx.diffusion.schedules.NoiseSchedulers¶
class nnx.diffusion.schedules.NoiseSchedulers(Enum)
Diffusion noise-schedule factory. Enum-as-factory pattern (like nnx.Nets, nnx.Optims): each enum variant's __call__ constructs the underlying NoiseSchedule.
nnx.diffusion.schedules.NoiseSchedulers.LINEAR¶
nnx.diffusion.schedules.NoiseSchedulers.LINEAR = 'linear'
Enum value linear.
nnx.diffusion.schedules.NoiseSchedulers.COSINE¶
nnx.diffusion.schedules.NoiseSchedulers.COSINE = 'cosine'
Enum value cosine.
nnx.diffusion.schedules.NoiseSchedule¶
class nnx.diffusion.schedules.NoiseSchedule(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') -> 'None'
Precomputed DDPM noise schedule.
Details
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.
nnx.diffusion.schedules.NoiseSchedule.to¶
nnx.diffusion.schedules.NoiseSchedule.to(self, device) -> 'NoiseSchedule'
Return a copy with every tensor moved to device. The kind and T fields are unchanged.
nnx.diffusion.nets.DiffusionMLP¶
class nnx.diffusion.nets.DiffusionMLP(input_dim: 'int', hidden_dims: 'list[int] | None' = None, time_embed_dim: 'int' = 32)
Conditional MLP for low-dim diffusion: forward(x_t, t) -> ε_pred.
Details
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.
nnx.diffusion.nets.DiffusionMLP.forward¶
nnx.diffusion.nets.DiffusionMLP.forward(self, x: 'torch.Tensor', t: 'torch.Tensor') -> 'torch.Tensor'
Predict noise added to x at timestep t.
Details
Args:
x: ``(B, *)`` clean shape; flattened internally to ``(B, D)``.
t: ``(B,)`` integer timesteps.
Returns:
Tensor of the same shape as ``x``.
nnx.diffusion.nets.DiffusionMLP.unpack_batch¶
nnx.diffusion.nets.DiffusionMLP.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.
nnx.diffusion.nets.sinusoidal_time_embed¶
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.
Details
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)``.
nnx.diffusion.training.diffusion_train_step_factory¶
nnx.diffusion.training.diffusion_train_step_factory(schedule: 'NoiseSchedule') -> 'TrainStepFn'
Build a DDPM noise-prediction TrainStepFn.
Details
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=...)``.
Raises:
TypeError: if ``schedule`` is not a :class:`NoiseSchedule` (the
enum member itself, say, instead of ``NoiseSchedulers.COSINE()``).
nnx.diffusion.objective¶
module nnx.diffusion.objective
DDPM noise prediction as an objective.
Details
:func:`diffusion_train_step_factory` is an imperative step: it draws the
timesteps and noise from the global RNG, takes a mean ``F.mse_loss`` and
hands it to ``finalize_step``, which owns the optimizer step — so it cannot
accumulate gradients over uneven microbatches or share the update engine.
:func:`diffusion_objective` describes the same loss as an objective
(``nnx.objectives``) and leaves backward, accumulation, mixed precision,
clipping and the optimizer step to NNx's shared update engine::
model.train(params=NNTrainParams(...), objective=diffusion_objective(schedule, seed=0))
For each microbatch the objective returns one ``"noise_mse"`` loss term:
the **sum** of squared errors between the predicted and the true noise
over every element, normalized by the element count. Over an update
window the engine divides the summed errors by the summed counts, so
microbatches ``[2, 1]`` take exactly the update of the full batch of 3 for
the same timesteps and noise.
The timesteps and noise come from the objective's **own** generator (a CPU
``torch.Generator``, seeded from ``seed`` or, when ``None``, from one draw
of the global RNG at first use). Its state is checkpointed component state
(``"diffusion.objective"``), so a stateful warm resume continues the
stream, and a :func:`~nnx.diffusion.sample` preview with its own generator
never perturbs it. Each run starts the stream afresh (a stateful resume
then restores the saved one), so an instance reused for a second, equally
seeded run draws the same noise. The draws happen on the CPU — the stream
is then the same on every device and survives a resume onto another — at
the cost of a host-to-device copy per microbatch; pass
``noise_fn(x_0, generator) -> (t, eps)`` to supply the timesteps and noise
yourself (fixed values in tests, another timestep distribution). The
``generator`` it receives is the objective's CPU generator: a ``noise_fn``
that draws on a CUDA device for large image batches must use its own CUDA
generator or the global RNG instead (which a stateful resume restores, but
which a ``sample`` preview without its own generator would advance). Its
timesteps must be int64 (or int32) values in ``[0, T)``.
nnx.diffusion.objective.diffusion_objective¶
nnx.diffusion.objective.diffusion_objective(schedule: 'NoiseSchedule', *, seed: 'Optional[int]' = None, noise_fn: 'Optional[NoiseFn]' = None, nonfinite: 'str' = 'fail') -> 'DiffusionObjective'
DDPM noise prediction as an objective (see DiffusionObjective) — the objective counterpart of diffusion_train_step_factory, the imperative step.
nnx.diffusion.objective.DiffusionObjective¶
class nnx.diffusion.objective.DiffusionObjective(schedule: 'NoiseSchedule', *, seed: 'Optional[int]' = None, noise_fn: 'Optional[NoiseFn]' = None, nonfinite: 'str' = 'fail') -> 'None'
DDPM noise prediction as an objective — see the module docstring.
Details
Callable as an objective and a checkpointable component
named ``"diffusion.objective"``: its state is the objective spec (the
schedule's length and fingerprint, the loss term) and the generator's
state, so a stateful resume draws the timesteps and noise an
uninterrupted run would have drawn. A resume with a different schedule
is rejected before anything is restored.
Args:
schedule: the :class:`NoiseSchedule` (any device; the coefficients
are indexed on the model's device).
seed: seeds the objective's generator; ``None`` draws one seed from
the global RNG at first use (reproducible under ``set_seed``).
noise_fn: optional ``(x_0, generator) -> (t, eps)``: ``t`` a
``LongTensor[B]`` in ``[0, T)``, ``eps`` shaped like ``x_0``.
Defaults to ``t ~ Uniform{0..T-1}``, ``eps ~ N(0, I)`` from the
objective's generator.
nonfinite: the engine's non-finite policy (``"fail"`` / ``"skip"``).
nnx.diffusion.objective.DiffusionObjective.check_run¶
nnx.diffusion.objective.DiffusionObjective.check_run(self, model: 'Any', *, optimizers: 'Mapping[str, torch.optim.Optimizer]', callbacks: 'Any') -> 'None'
A new run starts the noise stream afresh — from seed, or from one draw of the (just seeded) global RNG — whatever an earlier run of this instance drew; a stateful resume then restores the saved stream.
nnx.diffusion.objective.DiffusionObjective.generator¶
property nnx.diffusion.objective.DiffusionObjective.generator
The objective's own CPU generator (created at first use).
nnx.diffusion.objective.DiffusionObjective.draw¶
nnx.diffusion.objective.DiffusionObjective.draw(self, x_0: 'torch.Tensor') -> 'tuple[torch.Tensor, torch.Tensor]'
This microbatch's timesteps and noise, on x_0's device.
nnx.diffusion.objective.DiffusionObjective.component_spec¶
nnx.diffusion.objective.DiffusionObjective.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.diffusion.objective.DiffusionObjective.component_state¶
nnx.diffusion.objective.DiffusionObjective.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.diffusion.objective.DiffusionObjective.check_component_state¶
nnx.diffusion.objective.DiffusionObjective.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.diffusion.objective.DiffusionObjective.load_component_state¶
nnx.diffusion.objective.DiffusionObjective.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.diffusion.sampling.sample¶
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.
Details
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.
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¶
nnx.paradigms.distillation.kd_train_step_factory(teacher: 'NNModel', *, alpha: 'float' = 0.5, temperature: 'float' = 4.0) -> 'TrainStepFn'
Build a knowledge-distillation TrainStepFn.
Details
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.
At the first step: a single-logit head with ``alpha > 0`` (no
class distribution to distill), or a teacher of another shape.
TypeError: if ``teacher`` is not an ``NNModel`` (stored teacher
records train through ``TeacherDataset.objective()``).
nnx.paradigms.distillation.feature_kd_train_step_factory¶
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 TrainStepFn.
Details
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, or if the head has a single logit
and ``alpha > 0`` (no class distribution to distill).
TypeError: if ``teacher`` is not an ``NNModel`` (stored teacher
records train through ``TeacherDataset.objective()``).
nnx.paradigms.born_again.born_again_train¶
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.
Details
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 ``NNModel`` wrapper is reused, but its network weights are reset
to the caller-provided initialization before every student generation.
This follows the original Born-Again Networks procedure while keeping
memory usage to the live student, one frozen teacher, and one initial
state dictionary.
Args:
model: the :class:`NNModel` to train. Its initial weights seed every
fresh student; its final state corresponds to the LAST generation.
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: base configuration for every :meth:`NNModel.train`
call. Generation zero uses it unchanged. Each later generation
records the preceding run as its parent, producing a distinct
content-addressed run with independent history and 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`` and the
optimizer accumulates gradients, ``alpha`` / ``temperature``
are out of range, the head has a single output with
``alpha > 0`` (no class distribution to distill; the width is
read from the net's ``params.output_dim`` or from the last layer
of an ``nn.Sequential``), or the net is a ``TransformerNN``
(sequence output). Nothing is trained then. Another module's
single-output head is refused by the KD step at generation 1,
after generation 0 trained.
TypeError: if the training loader holds stored teacher records
(for any ``generations``), or if ``generations > 1`` and a
keyword is not a :func:`kd_train_step_factory` parameter.
PrecisionUnsupportedError: if ``generations > 1`` and the model's
precision is reduced: the KD generations run in full
precision only, so nothing is trained.
nnx.paradigms.offline_distillation¶
module nnx.paradigms.offline_distillation
Offline teacher distributions: distill a student from stored, permissioned teacher probabilities.
Details
``kd_train_step_factory`` and ``kd_objective`` forward a **live** teacher
every step and compare temperature-softened **logits** (scaled by ``T²``).
Precomputed teacher output is often only a probability distribution — a
decision provider's answers, an export from a model you may no longer run —
with no pre-temperature logits to soften. This module trains on those
records directly, offline:
- :class:`TeacherRecord` — one stored distribution: the sample and question
ids, the ordered candidate ids and their probabilities (finite,
nonnegative, summing to 1 within ``1e-6``), the teacher model and
revision, the schema digest (e.g. a ``nnx.decisions`` question's
``digest()``), the probability ``semantics`` and the record's
``provenance`` — its ``source`` and the **declared** ``training_rights``
under which it may train a student. NNx records the declaration; it does
not verify entitlement. An optional ``label`` (a candidate id) is the
sample's hard target. :func:`write_teacher_records` /
:func:`read_teacher_records` keep them as strict JSONL
(``nnx.teacher-record/1``).
- :class:`TeacherDataset` — records joined to the student's inputs **by
sample id**, every distribution aligned **by candidate id** to the
student's output order (``candidates``). A record missing a candidate,
naming an unknown or duplicate one, another schema digest, a repeated
sample id or a sample with no input is refused when the dataset is built
— before any optimizer exists. :meth:`TeacherDataset.loader` yields
:class:`TeacherBatch`es.
- :class:`OfflineDistillationObjective` (:meth:`TeacherDataset.objective`)
— an objective for NNx's shared update engine: ``alpha ·
KL(teacher ‖ student) + (1 − alpha) · CE(student, label)``, both at
temperature 1 — no logits are invented and no ``T²`` is applied. The two
terms share one mask (every row) and one denominator (the microbatch's
rows), so microbatches accumulate exactly to the full-batch update; a
partly labelled batch is refused, and ``alpha < 1`` needs every record
labelled. The teacher's probabilities are detached, and the objective
never back-propagates, zeroes, steps, scales or schedules. Its identity —
objective version, schema digest, ``alpha`` and the candidate alignment
(``identity()``) — is checkpointed as component state
``"offline_distillation.objective"``, so a stateful resume with any of
them changed is refused before the first resumed update; pass the same
identity to ``ExperimentManifest.for_model(objective=...)`` to record it
in a run's provenance.
- :func:`evaluate_offline` (and :meth:`TeacherDataset.eval_step` for the
training loop's validation) reports **teacher agreement** (the student's
top candidate is the teacher's, over every record), **imitation loss**
(mean ``KL(teacher ‖ student)`` in nats, over every record) and
**labelled quality** (accuracy and NLL on the labelled records only),
each with its own denominator: agreeing with a teacher is not being
right.
Nothing here reaches the network: records, inputs and the student are
local. Live distillation (``kd_train_step_factory``,
``feature_kd_train_step_factory``, ``born_again_train``, ``kd_objective``)
keeps its logit / temperature contract and refuses these records.
nnx.paradigms.offline_distillation.TeacherRecord¶
class nnx.paradigms.offline_distillation.TeacherRecord(sample_id: 'str', question_id: 'str', candidates: 'tuple[str, ...]', probabilities: 'tuple[float, ...]', teacher: 'str', revision: 'str', schema_digest: 'str', semantics: 'str', provenance: 'Mapping[str, str]', label: 'Optional[str]' = None) -> 'None'
One stored teacher distribution over a question's candidates (see the module docstring). Every field but label is required.
nnx.paradigms.offline_distillation.TeacherRecord.distribution¶
nnx.paradigms.offline_distillation.TeacherRecord.distribution(self) -> 'dict[str, float]'
{candidate id: probability}.
nnx.paradigms.offline_distillation.TeacherRecord.state¶
nnx.paradigms.offline_distillation.TeacherRecord.state(self) -> 'dict[str, Any]'
The record as strict JSON (nnx.teacher-record/1).
nnx.paradigms.offline_distillation.TeacherRecord.from_state¶
nnx.paradigms.offline_distillation.TeacherRecord.from_state(state: 'Mapping[str, Any]') -> 'TeacherRecord'
No public description is currently available.
nnx.paradigms.offline_distillation.TeacherDataset¶
class nnx.paradigms.offline_distillation.TeacherDataset(records: 'Iterable[TeacherRecord]', inputs: 'Mapping[str, Any]', *, candidates: 'Sequence[str]', schema_digest: 'str') -> 'None'
TeacherRecords joined to the student's inputs (a mapping from sample id to its input tensor) and aligned to the student's output order candidates (see the module docstring).
Details
Args:
records: the teacher records — one per sample, all with this
``schema_digest`` and exactly these candidates (in any order).
inputs: ``{sample id: input tensor}``; every record's sample needs
one (extra inputs are ignored).
candidates: the student's output order: output ``i`` scores
``candidates[i]``.
schema_digest: the question schema the records answer (e.g. a
``nnx.decisions`` question's ``digest()``).
nnx.paradigms.offline_distillation.TeacherDataset.sample_ids¶
property nnx.paradigms.offline_distillation.TeacherDataset.sample_ids
No public description is currently available.
nnx.paradigms.offline_distillation.TeacherDataset.labelled¶
property nnx.paradigms.offline_distillation.TeacherDataset.labelled
"all", "none" or "partial": which records carry a label.
nnx.paradigms.offline_distillation.TeacherDataset.collate¶
nnx.paradigms.offline_distillation.TeacherDataset.collate(self, items: 'Sequence[tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]]') -> 'TeacherBatch'
No public description is currently available.
nnx.paradigms.offline_distillation.TeacherDataset.loader¶
nnx.paradigms.offline_distillation.TeacherDataset.loader(self, batch_size: 'int', *, shuffle: 'bool' = False, seed: 'Optional[int]' = None) -> 'DataLoader'
A DataLoader of TeacherBatches; shuffle draws its order from seed (a generator a stateful resume restores).
nnx.paradigms.offline_distillation.TeacherDataset.objective¶
nnx.paradigms.offline_distillation.TeacherDataset.objective(self, *, alpha: 'float' = 0.5, nonfinite: 'str' = 'fail') -> 'OfflineDistillationObjective'
This dataset's OfflineDistillationObjective — refused here, before any optimizer exists, for a partly labelled dataset (any alpha: its batches would mix labelled and unlabelled rows) and, with alpha < 1 (a hard term), for one without labels.
nnx.paradigms.offline_distillation.TeacherDataset.eval_step¶
nnx.paradigms.offline_distillation.TeacherDataset.eval_step(self) -> 'OfflineDistillationEval'
An eval_step_fn for NNModel.train over a validation loader of this dataset's kind (see OfflineDistillationEval).
nnx.paradigms.offline_distillation.TeacherBatch¶
class nnx.paradigms.offline_distillation.TeacherBatch(inputs: 'torch.Tensor', teacher: 'torch.Tensor', target: 'torch.Tensor', labelled: 'torch.Tensor', sample_ids: 'tuple[str, ...]', schema_digest: 'str', candidates: 'tuple[str, ...]') -> 'None'
A microbatch of a TeacherDataset: the student's inputs, the teacher probabilities (rows, candidates) in the dataset's candidate order, the hard target indices (-1 where unlabelled), the labelled row mask, the sample_ids, and the dataset's schema_digest and candidates (the order of the probability columns), which an objective or evaluation checks against its own.
nnx.paradigms.offline_distillation.OfflineDistillationObjective¶
class nnx.paradigms.offline_distillation.OfflineDistillationObjective(*, schema_digest: 'str', candidates: 'Sequence[str]', alpha: 'float' = 0.5, nonfinite: 'str' = 'fail') -> 'None'
alpha · KL(teacher ‖ student) + (1 − alpha) · CE(student, label) over stored teacher distributions (see the module docstring). Build it with TeacherDataset.objective.
Details
Terms (both normalized by the microbatch's rows — one shared mask and
denominator): ``"teacher_kl"`` (weight ``alpha``, omitted when
``alpha == 0``) and ``"hard_ce"`` (weight ``1 − alpha``, omitted when
``alpha == 1``).
nnx.paradigms.offline_distillation.OfflineDistillationObjective.identity¶
nnx.paradigms.offline_distillation.OfflineDistillationObjective.identity(self) -> 'dict[str, Any]'
The objective's identity: id, version, schema digest, alpha and the candidate alignment — checkpointed as component state and fit for ExperimentManifest.for_model(objective=...).
nnx.paradigms.offline_distillation.OfflineDistillationObjective.component_spec¶
nnx.paradigms.offline_distillation.OfflineDistillationObjective.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.paradigms.offline_distillation.OfflineDistillationObjective.component_state¶
nnx.paradigms.offline_distillation.OfflineDistillationObjective.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.paradigms.offline_distillation.OfflineDistillationObjective.check_component_state¶
nnx.paradigms.offline_distillation.OfflineDistillationObjective.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.paradigms.offline_distillation.OfflineDistillationObjective.load_component_state¶
nnx.paradigms.offline_distillation.OfflineDistillationObjective.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
Nothing to restore: the checked identity is the whole state.
nnx.paradigms.offline_distillation.OfflineDistillationEval¶
class nnx.paradigms.offline_distillation.OfflineDistillationEval(*, schema_digest: 'str', candidates: 'Sequence[str]') -> 'None'
An eval_step_fn (NNModel.train(eval_step_fn=...)) over a validation loader of TeacherBatches: its metrics are teacher_kl, teacher_agreement and — with labelled records — label_accuracy and label_nll.
Details
The record's ``loss`` is the **imitation loss** (``teacher_kl``) whatever
the objective's ``alpha``: it measures agreement with the teacher, not
quality. A scheduler, early stopping or BEST selection that should
follow labelled quality must monitor ``label_accuracy`` or ``label_nll``
by name (``MonitorSpec``). Use :func:`evaluate_offline` for every
denominator.
nnx.paradigms.offline_distillation.evaluate_offline¶
nnx.paradigms.offline_distillation.evaluate_offline(model: 'NNModel', data: 'Union[TeacherDataset, DataLoader]', *, batch_size: 'int' = 256) -> 'DistillationReport'
Score model against a TeacherDataset (or a DataLoader over one, such as TeacherDataset.loader; its dataset gives the schema and candidates): teacher agreement, imitation loss and labelled quality, each with its own denominator (see DistillationReport). Runs in eval mode under no_grad and restores the model's modes.
nnx.paradigms.offline_distillation.DistillationReport¶
class nnx.paradigms.offline_distillation.DistillationReport(agreement: 'Measure', imitation_loss: 'Measure', label_accuracy: 'Measure', label_nll: 'Measure', candidates: 'tuple[str, ...]' = (), schema_digest: 'str' = '') -> 'None'
A student scored against stored teacher distributions.
Details
- ``agreement``: records whose student top candidate is the teacher's
(ties go to the first candidate in the dataset's order), over every
record;
- ``imitation_loss``: mean ``KL(teacher ‖ student)`` in nats, over every
record;
- ``label_accuracy`` / ``label_nll``: quality on the labelled records
only — unavailable without any.
nnx.paradigms.offline_distillation.DistillationReport.state¶
nnx.paradigms.offline_distillation.DistillationReport.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.paradigms.offline_distillation.DistillationReport.text¶
nnx.paradigms.offline_distillation.DistillationReport.text(self) -> 'str'
No public description is currently available.
nnx.paradigms.offline_distillation.Measure¶
class nnx.paradigms.offline_distillation.Measure(value: 'Optional[float]', denominator: 'int', reason: 'Optional[str]' = None) -> 'None'
One reported quantity with its denominator, or value=None and the reason it is unavailable.
nnx.paradigms.offline_distillation.Measure.state¶
nnx.paradigms.offline_distillation.Measure.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.paradigms.offline_distillation.read_teacher_records¶
nnx.paradigms.offline_distillation.read_teacher_records(path: 'Union[str, os.PathLike[str]]') -> 'list[TeacherRecord]'
The records of a JSONL file written by write_teacher_records; a malformed line — a number too large for a float included — raises TeacherRecordError naming path:line.
nnx.paradigms.offline_distillation.write_teacher_records¶
nnx.paradigms.offline_distillation.write_teacher_records(path: 'Union[str, os.PathLike[str]]', records: 'Iterable[TeacherRecord]') -> 'None'
One record per line (strict JSON, sorted keys), written atomically.
nnx.paradigms.offline_distillation.TeacherRecordError¶
class nnx.paradigms.offline_distillation.TeacherRecordError
A teacher record, dataset or batch that cannot train or be scored.
14.2. Contrastive¶
nnx.paradigms.contrastive.simclr_train_step_factory¶
nnx.paradigms.contrastive.simclr_train_step_factory(*, temperature: 'float' = 0.5) -> 'TrainStepFn'
Build a SimCLR TrainStepFn.
Details
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.
nnx.paradigms.contrastive.nt_xent_loss¶
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.
Details
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).
14.3. Augmentation¶
nnx.paradigms.augmentation.mixup_train_step_factory¶
nnx.paradigms.augmentation.mixup_train_step_factory(*, alpha: 'float' = 0.4) -> 'TrainStepFn'
Build a Mixup TrainStepFn.
Details
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.
nnx.paradigms.augmentation.cutmix_train_step_factory¶
nnx.paradigms.augmentation.cutmix_train_step_factory(*, alpha: 'float' = 1.0) -> 'TrainStepFn'
Build a CutMix TrainStepFn for 4D image batches.
Details
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.
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¶
nnx.paradigms.moe.moe_train_step_factory(*, aux_loss_weight: 'float' = 0.01) -> 'TrainStepFn'
Build an MoE-aware supervised TrainStepFn.
Details
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``.
14.5. I-JEPA¶
Walkthrough at I-JEPA. The ViTNN encoder is documented in §4.
nnx.paradigms.jepa.jepa_train_step_factory¶
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 TrainStepFn.
Details
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). It must be a submodule of
``model.net`` (the ViTNN), registered before training:
``NNModel.train`` builds the optimizer from ``model.net``, and
checkpoints and gradient clipping cover only ``model.net``.
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)``; at the
first step, when the predictor is not a submodule of
``model.net`` (it would never be updated).
nnx.paradigms.jepa.JEPATrainStep¶
class nnx.paradigms.jepa.JEPATrainStep(step: 'TrainStepFn', target_encoder: 'nn.Module', ema_momentum: 'float') -> 'None'
The I-JEPA step returned by jepa_train_step_factory.
Details
Callable as a ``TrainStepFn`` and a checkpointable component
named ``"jepa.target_encoder"``: every checkpoint's
training state carries the EMA target encoder's weights, and a stateful
warm resume restores them, so a run split at an epoch boundary
continues with the same target encoder an uninterrupted run would
have. The component is required — resuming from a stateful checkpoint
that records component state but not this component fails before
anything is restored, while a checkpoint with no component state at all
resumes with a fresh target copy and a warning.
nnx.paradigms.jepa.JEPATrainStep.component_spec¶
nnx.paradigms.jepa.JEPATrainStep.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.paradigms.jepa.JEPATrainStep.component_state¶
nnx.paradigms.jepa.JEPATrainStep.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.paradigms.jepa.JEPATrainStep.load_component_state¶
nnx.paradigms.jepa.JEPATrainStep.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.paradigms.jepa_objective¶
module nnx.paradigms.jepa_objective
I-JEPA latent prediction as an objective.
Details
:func:`~nnx.paradigms.jepa_train_step_factory` is an imperative step: it
takes a masked mean squared error, hands it to ``finalize_step`` (which
owns the optimizer step) and runs :func:`~nnx.paradigms.update_ema` at
once — so it cannot accumulate gradients over uneven microbatches or share
the update engine. :func:`jepa_objective` describes the same loss as an
objective (``nnx.objectives``)::
target = build_target_encoder(model.net)
predictor = JEPAPredictor(embed_dim=..., n_patches=model.net.n_patches)
model.net.add_module("_jepa_predictor", predictor) # optimizer-owned, once
model.train(params=NNTrainParams(...), objective=jepa_objective(target, predictor, mask_fn))
For each microbatch the objective returns one ``"latent_mse"`` loss term:
the **sum** of squared errors between the predicted and the target
embeddings at the masked positions, normalized by the number of those
elements — so microbatches with different target masks combine as
``(sum_1 + sum_2) / (count_1 + count_2)``. The EMA target encoder runs
under ``no_grad`` and is never optimizer-owned; it advances **once per
committed update**, from the updated online weights (the objective's
``after_update`` hook) — never per microbatch, and not at all for a skipped
window.
nnx.paradigms.jepa_objective.jepa_objective¶
nnx.paradigms.jepa_objective.jepa_objective(target_encoder: 'nn.Module', predictor: 'nn.Module', mask_fn: 'MaskFn', *, ema_momentum: 'float' = 0.996, nonfinite: 'str' = 'fail') -> 'JEPAObjective'
I-JEPA latent prediction as an objective (see JEPAObjective) — the objective counterpart of jepa_train_step_factory, the imperative step.
nnx.paradigms.jepa_objective.JEPAObjective¶
class nnx.paradigms.jepa_objective.JEPAObjective(target_encoder: 'nn.Module', predictor: 'nn.Module', mask_fn: 'MaskFn', *, ema_momentum: 'float' = 0.996, nonfinite: 'str' = 'fail') -> 'None'
I-JEPA latent prediction as an objective — see the module docstring.
Details
Callable as an objective and a checkpointable component
named ``"jepa.objective"``: its state is the objective spec (loss term,
EMA momentum), the predictor's **reference** (its name inside
``model.net`` — its weights are saved once, with the net), the EMA
target encoder's weights and the EMA update counter, so a stateful
resume continues exactly where an uninterrupted run would be. A saved
state carrying a second copy of the predictor's weights, a different
predictor reference or a different spec is rejected before anything is
restored.
Before any run is reserved (:meth:`check_run`) the objective refuses:
a net without the ViT patch contract (``n_patches``,
``patch_positions()``); a predictor that is not a submodule of
``model.net`` or whose trainable parameters are not owned exactly once
by the run's optimizers; a target encoder inside ``model.net``,
sharing its parameters or owned by an optimizer; a target parameter
without a same-named, same-shaped online parameter; and callbacks that
change the net's topology (``checkpoint_transforms``, e.g. QAT), which
would break the EMA's name correspondence and the predictor reference.
Args:
target_encoder: the EMA copy of ``model.net``
(:func:`~nnx.paradigms.build_target_encoder`); frozen and put in
eval mode once :meth:`check_run` has validated it (never here,
so a misplaced target cannot freeze the online network).
predictor: the predictor (:class:`~nnx.paradigms.JEPAPredictor` or
the same ``forward(context_embeds, context_positions,
target_positions)`` contract), registered under ``model.net``.
mask_fn: ``(n_patches, device) -> (context_mask, target_mask)``,
complementary ``BoolTensor[n_patches]``, sampled once per
microbatch (shared across its rows). Masks drawn from the global
RNG (:func:`~nnx.paradigms.random_block_mask`'s default) continue
across a stateful resume, which restores that RNG.
ema_momentum: fixed EMA decay in ``[0, 1)``:
``target ← momentum · target + (1 − momentum) · online``.
nonfinite: the engine's non-finite policy (``"fail"`` / ``"skip"``).
nnx.paradigms.jepa_objective.JEPAObjective.check_run¶
nnx.paradigms.jepa_objective.JEPAObjective.check_run(self, model: 'Any', *, optimizers: 'Mapping[str, torch.optim.Optimizer]', callbacks: 'Sequence[Any]') -> 'None'
Refuse a run this objective cannot train — called by NNModel.train and Trainer.train with the built optimizers and the callbacks, before any run is reserved or parameter changes. Raise to refuse; the default accepts every run.
nnx.paradigms.jepa_objective.JEPAObjective.after_update¶
nnx.paradigms.jepa_objective.JEPAObjective.after_update(self, events: 'tuple[UpdateEvent, ...]') -> 'None'
Advance the EMA target from the updated online weights — once per committed update, whatever the number of named optimizers.
nnx.paradigms.jepa_objective.JEPAObjective.component_spec¶
nnx.paradigms.jepa_objective.JEPAObjective.component_spec(self) -> 'ComponentSpec'
No public description is currently available.
nnx.paradigms.jepa_objective.JEPAObjective.component_state¶
nnx.paradigms.jepa_objective.JEPAObjective.component_state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.paradigms.jepa_objective.JEPAObjective.check_component_state¶
nnx.paradigms.jepa_objective.JEPAObjective.check_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'list[str]'
No public description is currently available.
nnx.paradigms.jepa_objective.JEPAObjective.load_component_state¶
nnx.paradigms.jepa_objective.JEPAObjective.load_component_state(self, state: 'Mapping[str, Any]', *, version: 'int') -> 'None'
No public description is currently available.
nnx.paradigms.jepa.JEPAPredictor¶
class nnx.paradigms.jepa.JEPAPredictor(*, embed_dim: 'int', n_patches: 'int', predictor_dim: 'Optional[int]' = None, n_layers: 'int' = 2, n_heads: 'int' = 2, ffn_mult: 'int' = 4)
Tiny ViT-like predictor: (context_embeds, target_positions) -> predicted_target_embeds.
Details
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.
nnx.paradigms.jepa.JEPAPredictor.forward¶
nnx.paradigms.jepa.JEPAPredictor.forward(self, context_embeds: 'torch.Tensor', context_positions: 'torch.Tensor', target_positions: 'torch.Tensor') -> 'torch.Tensor'
Predict embeddings at target_positions from context_embeds.
Details
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.
nnx.paradigms.jepa.build_target_encoder¶
nnx.paradigms.jepa.build_target_encoder(source: 'nn.Module') -> 'nn.Module'
Deep-copy source, freeze every parameter, return the copy.
Details
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.
nnx.paradigms.jepa.update_ema¶
nnx.paradigms.jepa.update_ema(source: 'nn.Module', target: 'nn.Module', momentum: 'float') -> 'None'
In-place EMA update: target ← momentum * target + (1 - momentum) * source.
Details
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).
nnx.paradigms.jepa.random_block_mask¶
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.
Details
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``; when
``block_scale`` is not an ordered pair in ``(0, 1)`` or
``block_aspect`` not an ordered pair of finite positive ratios; or
for a 1x1 grid (no room for a context). A block is never
empty and never the whole grid: each side is at least one
patch, a side longer than the grid is clipped to it, and a
block that would cover the grid loses one column.
14.6. DPO¶
Walkthrough at DPO.
nnx.paradigms.dpo.dpo_train_step_factory¶
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 TrainStepFn.
Details
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: strength of the reference-policy constraint (the scale of
the implicit reward). Larger ``beta`` keeps the policy closer
to the reference; smaller ``beta`` lets it deviate more. 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 be right-padded (content left-aligned) to fixed
lengths, as ``NNPreferenceDataset`` yields them.
Raises:
ValueError: if ``beta`` is not a finite positive number, or if
``pad_token_id`` is not a non-negative integer.
15. Embeddings (nnx.embeddings)¶
End-to-end walkthrough at Embeddings. Opt-in via pip install "thekaveh-nnx[embeddings]".
nnx.embeddings.contrastive_trainer.ContrastiveTextDataset¶
class nnx.embeddings.contrastive_trainer.ContrastiveTextDataset(pairs: 'list[tuple[str, str]]')
Wraps (anchor, positive) string pairs as a torch Dataset.
Details
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.
nnx.embeddings.contrastive_trainer.pair_collate¶
nnx.embeddings.contrastive_trainer.pair_collate(batch: 'Iterable[tuple[str, str]]') -> 'tuple[list[str], list[str]]'
Collate a batch of (anchor, positive) string tuples into two parallel lists (anchors, positives). Pair this with ContrastiveTextDataset when building your own DataLoader.
nnx.embeddings.contrastive_trainer.train_contrastive¶
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.AdamW, verbose: 'bool' = False) -> 'Any'
Train backbone on (anchor, positive) pairs via NT-Xent.
Details
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 ``lr``, 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).
nnx.embeddings.contrastive_trainer.embed_texts¶
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.
Details
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: the device of the returned tensor (the backbone encodes
where it lives and is never moved). ``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.
nnx.embeddings.contrastive_trainer.text_contrastive_train_step_factory¶
nnx.embeddings.contrastive_trainer.text_contrastive_train_step_factory(*, temperature: 'float' = 0.5) -> 'TrainStepFn'
Build a TrainStepFn for text-pair contrastive training.
Details
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.
nnx.embeddings.faiss_export.export_to_faiss¶
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.
Details
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: passed to :func:`embed_texts`, which encodes where the
backbone lives; the embeddings are moved to the CPU for the
index, so this argument has no effect on the result.
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``.
nnx.embeddings.faiss_export.export_to_safetensors¶
nnx.embeddings.faiss_export.export_to_safetensors(backbone: 'Any', out_path: 'Union[str, Path]') -> 'str'
Persist backbone.state_dict() to disk for downstream reload.
Details
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.
16. Interop (nnx.interop)¶
16.1. Experimental GGUF export¶
Walkthrough and stock-runtime limitations at Experimental GGUF export. Opt-in via pip install "thekaveh-nnx[gguf-write]".
nnx.interop.gguf.writer.write_gguf¶
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.
Details
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"``. Stock llama.cpp/Ollama do not
implement this architecture; the artifact is intended for
GGUF inspection or a reader that explicitly supports NNx.
Do not relabel it ``"llama"``: NNx uses interleaved RoPE,
which is not LLaMA's split-half layout.
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.
TypeError: when the network carries a PEFT adapter (LoRA / DoRA /
IA3 / prefix tuning), or a projection (``w_qkv`` / ``w_o`` /
``w1``–``w3`` / ``lm_head``) is not a plain ``nn.Linear`` (a
low-rank factorized one): GGUF stores plain weights. Nothing is
written.
nnx.interop.gguf.tensor_name_map.map_tensors¶
nnx.interop.gguf.tensor_name_map.map_tensors(net: 'TransformerNN') -> 'dict[str, np.ndarray]'
Walk a TransformerNN and return {gguf_name: numpy_array}.
Details
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.
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).
nnx.interop.ollama.export_ollama_modelfile¶
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 an experimental model.gguf + Modelfile bundle.
Details
Stock Ollama does not implement the ``nnx_transformer`` GGUF
architecture. Emission verifies bundle structure only; it does not
establish runtime compatibility.
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
from the documented 0.32.2 set. Each entry becomes a
``PARAMETER <key> <value>`` line; only ``stop`` accepts a
list or tuple, rendered as repeated lines. String values use
an injection-safe subset without quotes or control characters.
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``.
16.2. ONNX export conformance (nnx.export_conformance)¶
Executable profiles that separate structural validity from numerical parity: FeedFwdNN 4-8-2 FP32 under both exporters, run in ONNX Runtime on CPU, with every stage's outcome and the evidence recorded. Walkthrough, failure classes and exit codes at Export conformance. Opt-in via pip install "thekaveh-nnx[onnx-runtime,onnx-dynamo]"; python scripts/check_export_conformance.py --output conformance.json runs both profiles.
nnx.export_conformance¶
module nnx.export_conformance
Executable ONNX export conformance profiles.
Details
``onnx.checker`` proves an exported file is *well formed*; it says nothing
about what a runtime computes from it. A **profile** fixes everything a
parity claim depends on — the model and its weights, the exporter and its
options, the runtime, provider and dtype, the input cases and the
tolerances — runs it, and records each stage's outcome separately::
dependencies ─► export ─► check ─► load ─► parity
- **dependencies** — ``onnx`` and ``onnxruntime`` (and ``onnxscript`` for
the dynamo exporter) are importable. A missing one is a
``missing-dependency`` failure, never a skip.
- **export** — ``NNModel.to_onnx`` writes the graph (and any external data)
into the profile's own directory; the model's parameters, gradients and
(mixed) train / eval modes must be unchanged afterwards.
- **check** — ``onnx.checker.check_model(path, full_check=True)``; records
the opset, IR version and input / output signature. Every file the
exporter wrote is hashed (SHA-256), external tensors included.
- **load** — the hashes are verified first (a mismatch is rejected before
anything is loaded), then an ``onnxruntime.InferenceSession`` is created
on the declared provider.
- **parity** — every valid input case's raw logits must match the native
model in eval mode within ``rtol`` / ``atol``, and an input of the wrong
feature width must be refused by the runtime (``input-contract``).
Failures are classified (:data:`FAILURES`); a record's ``level`` is
``"executed"`` only when the runtime ran the graph and every case passed,
``"structural"`` when only the checker did, and ``"none"`` otherwise.
Tested profiles (:data:`PROFILES`) — the only targets NNx claims executed
parity for: ``FeedFwdNN`` 4-8-2 (ReLU, dropout 0, deterministic nontrivial
FP32 weights), exported at batch 2 with a dynamic batch dimension and run
on ONNX Runtime's CPU provider at batches 1, 2, 3 and 7, under the legacy
TorchScript exporter (``feedfwd-fp32-torchscript``) and the
``torch.export``-based one (``feedfwd-fp32-dynamo``). Everything else —
other architectures, quantized exports, Netron files, GGUF / Ollama
artifacts — stays structural or container-only: a checker pass or a
parser round trip never promotes a target to executed parity.
Run both profiles with ``python scripts/check_export_conformance.py
--output conformance.json`` (needs ``thekaveh-nnx[onnx-runtime,onnx-dynamo]``).
Nothing here imports ``onnx`` or ``onnxruntime`` until a profile runs.
nnx.export_conformance.Profile¶
class nnx.export_conformance.Profile(name: 'str', exporter: 'str', opset_version: 'int' = 17, export_batch: 'int' = 2, dynamic_batch: 'bool' = True, batches: 'tuple[int, ...]' = (1, 2, 3, 7), rtol: 'float' = 0.0001, atol: 'float' = 1e-05, runtime: 'str' = 'onnxruntime', provider: 'str' = 'CPUExecutionProvider', dtype: 'str' = 'float32', model: 'Mapping[str, Any]' = <factory>) -> 'None'
A conformance profile: one model, exporter configuration, runtime, provider, dtype, set of input cases and tolerances.
nnx.export_conformance.Profile.requires¶
property nnx.export_conformance.Profile.requires
No public description is currently available.
nnx.export_conformance.Profile.options¶
nnx.export_conformance.Profile.options(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.export_conformance.Profile.cases¶
nnx.export_conformance.Profile.cases(self) -> 'list[dict[str, Any]]'
No public description is currently available.
nnx.export_conformance.run_profile¶
nnx.export_conformance.run_profile(profile: 'Union[str, Profile]', directory: 'Union[str, os.PathLike[str]]', *, source_revision: 'Optional[str]' = None) -> 'dict[str, Any]'
Run one profile, writing its artifacts into directory/<name>/ (which must not exist or be empty), and return its record (see the module docstring). Never raises for a failing stage: every outcome is classified in the record.
nnx.export_conformance.run_profiles¶
nnx.export_conformance.run_profiles(directory: 'Union[str, os.PathLike[str]]', names: 'Optional[Sequence[str]]' = None, *, source_revision: 'Optional[str]' = None) -> 'dict[str, Any]'
Run names (every tested profile by default) under directory; the report's status is passed only when every profile passed.
nnx.export_conformance.execute¶
nnx.export_conformance.execute(record: 'Mapping[str, Any]', directory: 'Union[str, os.PathLike[str]]') -> 'dict[str, Any]'
Re-run the load and parity stages of a saved record against its artifacts under directory: the hashes are verified before anything is loaded (a mismatch is a hash-mismatch failure), and the native reference is rebuilt from the recorded config, whose weights must hash to the recorded digest. Returns the new record.
nnx.export_conformance.verify_artifacts¶
nnx.export_conformance.verify_artifacts(record: 'Mapping[str, Any]', directory: 'Union[str, os.PathLike[str]]') -> 'None'
Refuse artifacts that differ from the record: a file missing, a SHA-256 or size mismatch (external tensors included), or a file the record does not list. directory is the root the record's artifacts.directory lives under. An unreadable file or folder is a refusal too.
nnx.export_conformance.validate_record¶
nnx.export_conformance.validate_record(record: 'Any') -> 'None'
Refuse a record that does not follow nnx.export-conformance/1: missing or unknown keys (in the record and in its runtime, exporter and opset), anything but strict JSON, unknown statuses or failure classes, a failure in the wrong stage (an input case's included), an input case outcome no run gives, a stage that ran after a failed one, inconsistent status / failure / level, a passed stage its own evidence contradicts (parity without a compared valid case or with a failed case, an export that changed the model, a check without an opset or a listed model file), malformed hashes or paths, non-finite tolerances, or settings no profile produces (a model or batch larger than a profile allows included).
nnx.export_conformance.build_model¶
nnx.export_conformance.build_model(config: 'Mapping[str, Any]' = {'net': 'feed_fwd', 'input_dim': 4, 'hidden_dims': [8], 'output_dim': 2, 'activation': 'relu', 'dropout': 0.0, 'weights': {'generator': 'torch.Generator', 'seed': 0, 'distribution': 'normal', 'std': 0.5}}) -> 'Any'
The profile's NNModel (CPU, FP32) with deterministic weights; the caller's global RNG streams are left untouched.
nnx.export_conformance.exit_code¶
nnx.export_conformance.exit_code(report: 'Mapping[str, Any]') -> 'int'
0 when every profile passed, else the exit code of the first failing profile's failure class (EXIT_CODES).
nnx.export_conformance.save_report¶
nnx.export_conformance.save_report(report: 'Mapping[str, Any]', path: 'Union[str, os.PathLike[str]]') -> 'None'
Write the report as strict, sorted JSON (atomically).
nnx.export_conformance.summary_lines¶
nnx.export_conformance.summary_lines(report: 'Mapping[str, Any]') -> 'list[str]'
One line per profile: its verdict, level, each stage's status and how many input cases passed.
nnx.export_conformance.ConformanceError¶
class nnx.export_conformance.ConformanceError
A malformed conformance record, or artifacts that do not match it.
17. HuggingFace Hub + safetensors¶
Opt-in via pip install "thekaveh-nnx[hub]". Two integration surfaces:
- safetensors checkpoints —
NNCheckpoint.to_file(..., format="safetensors")writes the safetensors format alongside the default pickle path, andNNCheckpoint.from_file(path)reads either (detected from the file; see §3NNCheckpoint). Loadable by outside-Python tools (ComfyUI, vLLM, AutoGPTQ). - Hub publish / load —
NNModelmixes inhuggingface_hub.PyTorchModelHubMixin, sosave_pretrained(local_dir),push_to_hub(repo_id), andNNModel.from_pretrained(repo_id)work directly on a trained model. The mixin methods are inherited and live onNNModelitself. The NNx keywords offrom_pretrained(precision=,exclude_submodules=,map_location=,strict=) are documented in Hub.
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¶
class nnx.generation.LogitsProcessor(*args, **kwargs)
Callable protocol: logits, token_history -> adjusted_logits.
Details
``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.
nnx.generation.logits_chain¶
module nnx.generation.logits_chain
LogitsChain — a typed, ordered chain of LogitsProcessors.
Details
The Builder is the opt-in alternative to
`GenerativeNNModel.generate(temperature=..., top_k=..., top_p=...,
repetition_penalty=...)`: it lets the user compose the chain
declaratively and add custom processors (e.g., logit-bias for
forbidden tokens) without having to drop down to
`apply_chain(processors=[...])` and reconstruct the canonical order
themselves.
Builder ordering: regardless of the method-call order, `.build()`
sorts the processors into NNx's canonical order (the same order
`GenerativeNNModel.generate` builds from its inline kwargs):
RepetitionPenalty → TopKFilter → TopPFilter → TemperatureScaling
Temperature is deliberately LAST — temperature=0 greedy is implemented
via ±inf argmax markers that must not be re-filtered, so this differs
from HF's transformers, which applies temperature before top-k/top-p
(for temperature≠1 the nucleus token set can differ).
Custom processors (added via `.custom(processor)`) are appended after
TemperatureScaling — the "post-default" slot for user extensions.
nnx.generation.LogitsChain¶
class nnx.generation.LogitsChain(*, processors: 'list[LogitsProcessor]' = <factory>) -> 'None'
A typed, ordered sequence of LogitsProcessors.
Details
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.
nnx.generation.LogitsChain.apply¶
nnx.generation.LogitsChain.apply(self, logits: 'torch.Tensor', token_history: 'list[int]') -> 'torch.Tensor'
Run every processor in self.processors in order. Thin wrapper around apply_chain.
nnx.generation.LogitsChain.builder¶
nnx.generation.LogitsChain.builder() -> 'LogitsChainBuilder'
Return a fluent builder. See LogitsChainBuilder.
nnx.generation.LogitsChainBuilder¶
class nnx.generation.LogitsChainBuilder() -> 'None'
Fluent builder for a LogitsChain.
Details
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.
nnx.generation.LogitsChainBuilder.repetition_penalty¶
nnx.generation.LogitsChainBuilder.repetition_penalty(self, penalty: 'float') -> 'LogitsChainBuilder'
Add a RepetitionPenalty processor with the given penalty.
nnx.generation.LogitsChainBuilder.top_k¶
nnx.generation.LogitsChainBuilder.top_k(self, k: 'int') -> 'LogitsChainBuilder'
Add a TopKFilter with the given k.
nnx.generation.LogitsChainBuilder.top_p¶
nnx.generation.LogitsChainBuilder.top_p(self, p: 'float') -> 'LogitsChainBuilder'
Add a TopPFilter (nucleus sampling) with the given p.
nnx.generation.LogitsChainBuilder.temperature¶
nnx.generation.LogitsChainBuilder.temperature(self, t: 'float') -> 'LogitsChainBuilder'
Add a TemperatureScaling processor with the given temperature.
nnx.generation.LogitsChainBuilder.custom¶
nnx.generation.LogitsChainBuilder.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. Not with temperature(0.0) (greedy): its ±inf argmax markers must not be re-filtered, so build() refuses the pair; use OrderedLogitsPipeline with the custom stage before greedy.
nnx.generation.LogitsChainBuilder.build¶
nnx.generation.LogitsChainBuilder.build(self) -> 'LogitsChain'
Construct the LogitsChain with processors in canonical order.
Details
Standard processors that were chained are emitted in the
fixed `_CANONICAL_ORDER`; custom processors come after, in
the order they were added.
Raises:
ValueError: for custom processors after ``temperature(0.0)``:
greedy turns the logits into ±inf argmax markers, which a
logit bias or a token ban would corrupt (the chain would pick
a banned or arbitrary token).
nnx.generation.pipeline¶
module nnx.generation.pipeline
OrderedLogitsPipeline — immutable decoding stages run in declared order.
Details
:class:`~nnx.generation.LogitsChain` (via its builder) sorts the standard
processors into NNx's canonical order. An :class:`OrderedLogitsPipeline`
does the opposite: it runs its stages **exactly in the order they were
declared**, duplicates and custom positions included, so ordered decoding
configuration can be written down as data::
from nnx.generation import LogitsStage, OrderedLogitsPipeline
nucleus_then_cool = OrderedLogitsPipeline.of(LogitsStage.top_p(0.8), LogitsStage.temperature(0.5))
with_penalty = nucleus_then_cool.prepend(LogitsStage.repetition_penalty(1.2)) # a new pipeline
model.generate(prompt, logits_pipeline=with_penalty)
- **Immutable.** A pipeline holds a tuple of stages; :meth:`append` and
:meth:`prepend` return new pipelines. A built-in stage is a frozen spec
(:class:`LogitsStage`) — a processor passed in is read once and copied
into one, so mutating that processor later changes nothing — and every
run compiles fresh processors from the specs.
- **Validated.** Every value is finite and in its domain when the stage is
made. A zero-temperature stage is greedy: it turns the logits into
``±inf`` argmax markers, so it is **terminal** — a stage after it raises
(nothing is ever reordered).
- **Custom stages.** Any ``(logits, token_history) -> logits`` callable can
be a :class:`CustomStage` at any position (a subclass of a built-in
processor stays custom, keeping its own ``__call__``; the terminal rule
sees only built-in stages; a nested pipeline is flattened into its
stages). Its state is the caller's (the
pipeline keeps the reference, never a copy). It is runtime-only unless
its ``tag`` names a :class:`LogitsStageCodec` registered in this process
with :func:`register_logits_stage_codec`; a pipeline is never rebuilt by
importing code named in data, and nothing about it enters a checkpoint
or a run identity.
- **Versioned round trip.** :meth:`OrderedLogitsPipeline.state` writes
``{"version": 1, "stages": [...]}`` with one tagged entry per stage, in
order; :meth:`OrderedLogitsPipeline.from_state` rebuilds it.
nnx.generation.OrderedLogitsPipeline¶
class nnx.generation.OrderedLogitsPipeline(stages: 'tuple[Stage, ...]' = ()) -> 'None'
Decoding stages run exactly in declared order. See the module docstring for immutability, validation, the terminal greedy rule, custom stages and serialization.
nnx.generation.OrderedLogitsPipeline.of¶
nnx.generation.OrderedLogitsPipeline.of(*stages: 'Any') -> 'OrderedLogitsPipeline'
A pipeline of stages in this order.
nnx.generation.OrderedLogitsPipeline.append¶
nnx.generation.OrderedLogitsPipeline.append(self, *stages: 'Any') -> 'OrderedLogitsPipeline'
A new pipeline with stages after this one's.
nnx.generation.OrderedLogitsPipeline.prepend¶
nnx.generation.OrderedLogitsPipeline.prepend(self, *stages: 'Any') -> 'OrderedLogitsPipeline'
A new pipeline with stages before this one's.
nnx.generation.OrderedLogitsPipeline.processors¶
nnx.generation.OrderedLogitsPipeline.processors(self) -> 'list[LogitsProcessor]'
Fresh processors for every stage, in order (custom stages are the caller's callables).
nnx.generation.OrderedLogitsPipeline.state¶
nnx.generation.OrderedLogitsPipeline.state(self) -> 'dict[str, Any]'
{"version": 1, "stages": [...]}, one tagged entry per stage in order. A custom stage needs a registered codec for its tag.
nnx.generation.OrderedLogitsPipeline.from_state¶
nnx.generation.OrderedLogitsPipeline.from_state(state: 'Mapping[str, Any]') -> 'OrderedLogitsPipeline'
Rebuild a pipeline from state. A custom entry is decoded only by a codec registered in this process for its tag.
nnx.generation.LogitsStage¶
class nnx.generation.LogitsStage(kind: 'str', value: 'Union[int, float]') -> 'None'
One built-in decoding stage: kind is "repetition_penalty" (>= 1), "top_k" (an integer >= 1), "top_p" (in (0, 1]) or "temperature" (>= 0; 0 is greedy and terminal). Values are validated here, before any generation.
nnx.generation.LogitsStage.temperature¶
nnx.generation.LogitsStage.temperature(value: 'float') -> 'LogitsStage'
No public description is currently available.
nnx.generation.LogitsStage.top_k¶
nnx.generation.LogitsStage.top_k(value: 'int') -> 'LogitsStage'
No public description is currently available.
nnx.generation.LogitsStage.top_p¶
nnx.generation.LogitsStage.top_p(value: 'float') -> 'LogitsStage'
No public description is currently available.
nnx.generation.LogitsStage.repetition_penalty¶
nnx.generation.LogitsStage.repetition_penalty(value: 'float') -> 'LogitsStage'
No public description is currently available.
nnx.generation.LogitsStage.terminal¶
property nnx.generation.LogitsStage.terminal
Whether this is a greedy (zero-temperature) stage: it emits ±inf argmax markers, so nothing may follow it.
nnx.generation.LogitsStage.processor¶
nnx.generation.LogitsStage.processor(self) -> 'LogitsProcessor'
A fresh processor for this stage.
nnx.generation.LogitsStage.state¶
nnx.generation.LogitsStage.state(self) -> 'dict[str, Any]'
No public description is currently available.
nnx.generation.CustomStage¶
class nnx.generation.CustomStage(function: 'Callable[[torch.Tensor, list[int]], torch.Tensor]', tag: 'Optional[str]' = None) -> 'None'
A user callable (logits, token_history) -> logits at a declared position. Its state belongs to the caller: the pipeline keeps this reference and never copies it. tag names the LogitsStageCodec that can serialize it; without a registered codec the stage is runtime-only.
nnx.generation.CustomStage.processor¶
nnx.generation.CustomStage.processor(self) -> 'LogitsProcessor'
No public description is currently available.
nnx.generation.LogitsStageCodec¶
class nnx.generation.LogitsStageCodec(tag: 'str', encode: 'Callable[[Callable[..., Any]], Mapping[str, Any]]', decode: 'Callable[[Mapping[str, Any]], Callable[[torch.Tensor, list[int]], torch.Tensor]]') -> 'None'
How to serialize custom stages tagged tag: encode(function) returns a JSON-like mapping, decode(mapping) returns the callable. Registered in-process with register_logits_stage_codec; data never names code to import.
nnx.generation.register_logits_stage_codec¶
nnx.generation.register_logits_stage_codec(codec: 'LogitsStageCodec') -> 'None'
Register codec for its tag (re-registering a tag replaces it).
nnx.generation.unregister_logits_stage_codec¶
nnx.generation.unregister_logits_stage_codec(tag: 'str') -> 'None'
Remove the codec registered for tag (a missing tag is a no-op).
nnx.generation.registered_logits_stage_codecs¶
nnx.generation.registered_logits_stage_codecs() -> 'tuple[str, ...]'
The registered codec tags, in registration order.
nnx.generation.TemperatureScaling¶
class nnx.generation.TemperatureScaling(temperature: 'float')
Divide logits by temperature before sampling.
Details
``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.
nnx.generation.TopKFilter¶
class nnx.generation.TopKFilter(top_k: 'int')
Keep only the top-k logits per row; set the rest to -inf.
Details
-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.
nnx.generation.TopPFilter¶
class nnx.generation.TopPFilter(top_p: 'float')
Nucleus (top-p) sampling: keep the smallest set of tokens whose cumulative probability exceeds top_p.
Details
Edge case: if a single token's probability exceeds ``top_p``, only
that token is retained. ``top_p=1.0`` keeps every token (no
floating-point cumulative sum can exceed it by rounding).
nnx.generation.RepetitionPenalty¶
class nnx.generation.RepetitionPenalty(penalty: 'float')
Penalize previously-seen tokens (HF-style).
Details
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).
nnx.generation.apply_chain¶
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.
nnx.generation.sample_next_token¶
nnx.generation.sample_next_token(logits: 'torch.Tensor', *, generator: 'Optional[torch.Generator]' = None) -> 'int'
Draw one token id from softmax(logits).
Details
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.
19. Visualization¶
19.1. Run-output viz (nnx.vis_utils)¶
nnx.vis_utils¶
module nnx.vis_utils
Run-output visualization helpers.
Details
`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 plotting method is a `@staticmethod` returning either a
`plotly.graph_objects.Figure` (for plots) or a `pandas.DataFrame` (for
tables); the helpers `generate_colors` (a color list) and
`get_scatter_plot_vm` (a scatter view-model) feed them. 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(...)`.
nnx.vis_utils.confusion_matrix¶
nnx.vis_utils.confusion_matrix(Y_true, Y_pred, class_names=None, title: 'str' = 'Confusion matrix', normalize: 'bool' = False)
Render a confusion matrix heatmap. Y_true and Y_pred are 1-D arrays of integer class labels. If class_names is provided, axis labels use the named classes, and every label must be a class index 0..len(class_names) - 1 (others raise ValueError); otherwise the labels present, in sorted order.
nnx.vis_utils.classification_report¶
nnx.vis_utils.classification_report(Y_true, Y_pred, class_names=None) -> 'pd.DataFrame'
Per-class precision / recall / f1 / support as a DataFrame. Use the return value for tabular display (print(df.to_string()) or notebook auto-display) or to feed back into downstream analysis.
nnx.vis_utils.multi_line_plot¶
nnx.vis_utils.multi_line_plot(x, yss, title, yss_legend, x_axis_label, y_axis_label, x_ticks_inc=20, label_size=12, title_size=14, fig_size: 'tuple' = (1000, 600), margin_size={'l': 15, 'r': 15, 't': 30, 'b': 15, 'pad': 0}, renderer=None)
Render a multi-group line chart and return the Plotly Figure.
Details
Each group in `yss` is drawn with a distinct color; each line within
a group uses a distinct dash style. `yss_legend` is a (group_labels,
line_labels) tuple — group_labels name the colored groups (one per
entry in `yss`), line_labels name the dash styles shared across
groups. Both legends are added as no-trace markers so the legend
reads cleanly; data traces carry "group (line)" hover names.
Returns the Figure. If `renderer` is non-None, also calls
`fig.show(renderer=renderer)` so notebook callers see the chart
inline; pass `renderer=None` (the default) for headless usage.
nnx.vis_utils.scatter_plot¶
nnx.vis_utils.scatter_plot(vm, renderer=None, fig_size: 'tuple' = (1000, 600), label_size: 'int' = 12, title_size: 'int' = 14, margin_size={'l': 15, 'r': 15, 't': 30, 'b': 15, 'pad': 0})
Render a colored scatter plot from a view-model dict and return the Plotly Figure.
Details
`vm` is the structure produced by `get_scatter_plot_vm`: title, xs/ys
column views, plus a `ts` group axis carrying labels + colors per
category. Honors `renderer` the same way as `multi_line_plot`.
nnx.vis_utils.two_dim_tsne_checkpoint_logits¶
nnx.vis_utils.two_dim_tsne_checkpoint_logits(checkpoint: 'NNCheckpoint', ds: 'NNDataset', n_samples: 'int', random_state: 'int | None' = 0, renderer: 'str | None' = None, fig_size: 'tuple' = (1000, 600), title_size: 'int' = 14, label_size: 'int' = 12, margin_size={'l': 15, 'r': 15, 't': 30, 'b': 15, 'pad': 0}, *, module: 'Any' = None, batch_adapter: 'Any' = None)
Project the first n_samples test logits of checkpoint to 2D via t-SNE and render them colored by ground-truth class.
Details
Useful for eyeballing class separability of an intermediate
checkpoint — pass the BEST checkpoint to see how well-trained the
decision space ended up. Returns the Plotly Figure.
A registered module is rebuilt through its factory and
batches are split by its batch adapter (``batch_adapter=`` overrides
it). A runtime-only checkpoint needs ``module=`` — the weights are
loaded into it (its identity, state keys and training modes are
kept) — and fails with ``MissingModelFactoryError`` before any data
is read otherwise.
nnx.vis_utils.generate_colors¶
nnx.vis_utils.generate_colors(n)
Return n distinct hex-color strings spaced evenly around the HSV hue circle.
Details
Uses 60% saturation and 95% brightness for legend readability. `endpoint=False`
keeps the first and last hues maximally distinct — the hue circle wraps (0.0 ==
1.0 == red), so a closed interval would hand the first and last class identical
colors. Used by every other VisUtils method that draws categorical groups.
nnx.vis_utils.get_scatter_plot_vm¶
nnx.vis_utils.get_scatter_plot_vm(data, title, col_xs, label_xs, col_ys, label_ys, col_ts, labels_ts, colors_ts, uni_ts)
Build the view-model dict consumed by scatter_plot.
Details
Splits the input dataframe by the categorical column `col_ts`, attaches
per-category colors and labels, and precomputes the per-category
x/y slices so the plotting function stays simple.
Each function above is also a VisUtils static method of the same name (VisUtils.confusion_matrix is nnx.vis_utils.confusion_matrix).
nnx.vis_utils.VisUtils¶
class nnx.vis_utils.VisUtils()
Plotly-based visualizations for NNx run outputs (the plots extra: pip install "thekaveh-nnx[plots]"; classification_report needs no extra).
Details
The static-method facade for run-output visualization. Companion to
`nnx.viz` (model internals — weight / activation / gradient inspection)
and `nnx.vis_utils`'s module-level aliases (which re-export the most
common methods of this class for ergonomic top-level access).
Class attributes are the plotting methods' styling defaults.
``confusion_matrix`` reads them at each call. ``multi_line_plot``,
``scatter_plot`` and ``two_dim_tsne_checkpoint_logits`` take them as
keyword defaults fixed when the class is defined: pass ``renderer=`` /
``fig_size=`` (and so on) per call to change them there.
- `TITLE_SIZE` / `LABEL_SIZE`: Plotly title and axis-label font sizes.
- `FIG_SIZE`: default `(width, height)` in pixels.
- `MARGIN_SIZE`: default margin spec (passed to `Figure.update_layout`).
- `RENDERER`: Plotly renderer override. `None` lets Plotly auto-detect
(notebook / browser / etc.); set to a string like ``"png"`` or
``"browser"`` to pin behavior in non-default environments.
See README §4.6 "Visualization" and `docs/concepts.md`
for usage examples.
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¶
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.
Details
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`,
or if `max_channels` or `cols` is not a positive integer, or
if the layer returns a mapping that holds no tensor.
RuntimeError: If the forward hook on `layer_name` never fires
(the layer is not reached by this input's forward path).
nnx.viz.attribute.attribute¶
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.
Details
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` of the attribution: a 4-D input shows its first
sample mean-pooled over channels, a 3-D input its first sample,
and a 2-D ``(B, D)`` input one row per sample.
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.
nnx.viz.gradient_flow.gradient_flow¶
nnx.viz.gradient_flow.gradient_flow(model: 'Union[nn.Module, NNModel]') -> 'go.Figure'
Return a Plotly bar chart of per-parameter L2 gradient norms.
Details
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.
nnx.viz.netron.netron_export¶
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).
Details
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 (checked before anything is written). The export
itself needs the `onnx` package (the `onnx` extra); without it
`torch.onnx` raises its own exporter error.
NotImplementedError: When the module declares
`onnx_export_unsupported` (graph classifiers, link
predictors); nothing is written.
nnx.viz.summary.summary¶
nnx.viz.summary.summary(model: 'Union[nn.Module, NNModel]', *, input_size: 'tuple[int, ...] | None' = None, input_data: 'Union[torch.Tensor, tuple, list, dict, None]' = None, batch: 'Any' = None, depth: 'int' = 4, col_names: 'tuple[str, ...]' = ('output_size', 'num_params', 'mult_adds')) -> 'ModelStatistics'
Return a torchinfo.ModelStatistics summary for model.
Details
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, or a
dict of keyword inputs. Useful when the model takes multiple positional
arguments or a non-tensor input (graphs, dicts) that `input_size` can't describe.
batch: A training batch, split into the forward's inputs by the model's batch
adapter (an `NNModel`'s adapter, else the module's `unpack_batch`,
else one positional input). Mutually exclusive with `input_size` / `input_data`.
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]`.
ValueError: Before the module is touched, when `batch` is combined with
`input_size` / `input_data`, or when `input_size` is given for a model whose
adapter takes keyword inputs (a synthetic tensor cannot describe them).
The module's identity, state and every submodule's training mode are left as they
were (torchinfo itself only restores the top-level mode).
nnx.viz.weight_histogram.weight_histogram¶
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.
Details
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),
or if `bins` or `cols` is not a positive integer.
20. Utilities¶
nnx.utils¶
module nnx.utils
Pretty-printing helpers used throughout nnx.
Details
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.
nnx.utils.print_tree¶
nnx.utils.print_tree(tree, level: 'int' = 0, *, file=None) -> 'None'
Pretty-print a nested dict as an indented tree.
Details
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``.
nnx.utils.print_table¶
nnx.utils.print_table(data: 'dict', header: 'bool' = True, title: 'Optional[str]' = None, *, file=None) -> 'None'
Print data as a 2-column key/value table.
Details
Pass ``file=`` to redirect output. Defaults to ``sys.stdout``.
nnx.utils.flatten_dict¶
nnx.utils.flatten_dict(data: 'dict', parent_key: 'str' = '', sep: 'str' = '.') -> 'dict'
Flatten a nested dict so nested keys become parent.child style.
Details
>>> flatten_dict({"a": 1, "b": {"c": 2}})
{'a': 1, 'b.c': 2}
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.