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.

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.

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

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.

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.

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.

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.

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.

``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.

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.

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).

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.

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.

`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).

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.

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).

``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.

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.

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.

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.

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.

``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.

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``.

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.

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.

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).

`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).

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`.

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.

`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.

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.

Disable this when the custom step function owns scheduler timing.

`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.

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).

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).

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``).

``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.

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.

``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()``.

``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.

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.

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.

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`).

``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.

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.

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.

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.

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.

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"``.

``status`` is ``"ok"``, ``"missing"`` (no value; no decision) or
``"nonfinite"`` (an epoch without improvement); ``improved`` says
whether this epoch became the best.

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.

``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``.

``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.

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.

``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).

``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"``.

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.

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.

``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()``.

``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.

``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.

: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.

``id`` is bookkeeping — how results are keyed, never shown to a model;
``description`` is the model-facing text.

: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.

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`.

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.

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.

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).

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.

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.

: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``.

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.

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.

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.

``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).

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.

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).

``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).

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()``.

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.

``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"``).