23. Changelog

All notable changes to NNx are documented here. Format follows Keep a Changelog; versioning is roughly SemVer — pre-1.0, we allow behavior changes (typically bug fixes) without renaming public APIs.

This file intentionally keeps the standard Keep a Changelog heading format rather than the hierarchical numbering used by the rest of the documentation.

[Unreleased]

Added

• CUDA rows on the feature-composition matrix. The matrix now has a cuda evidence profile beside the CPU one. FP16 AMP (TC-06), BF16 AMP (TC-10) and two-GPU NCCL DDP (TC-11) have tests, and a report from a GPU run verifies them. Until one runs, they stay unverified. CONTRIBUTING describes the opt-in GPU runner and a manual run (#420).

Fixed

• A DDP fit releases its wrapper when it ends. After train(distributed=DDP()), including one that raised (a non-finite loss, say), the DistributedDataParallel wrapper and its hold on the process group are dropped. They no longer live on in an exception's traceback until an arbitrary garbage collection, possibly after the group is destroyed. nnx.distributed.shutdown() also collects such garbage before leaving the group (#478).
• Invalid training settings fail before training starts. NNTrainParams and NNTrainerParams refuse a seed outside [0, 2**32), a non-string data_id and an unknown resume_from_checkpoint tag. A Checkpoints member there is stored as its tag string, so it no longer crashes the first run.yaml write after an epoch has trained.
• Checkpoint reads and writes refuse a path-shaped type. NNCheckpoint.load(run=..., type="../...") and the other checkpoint entry points raise ValueError instead of reading or writing outside the run directory.
• A resumed linear_warmup_decay schedule on the epoch clock needs an explicit warmup_steps. Its default (a tenth of each session's epochs) is not in the restored scheduler state, so a resume silently warmed up over a different length.
• Continuous targets under a regression loss without a task fail with a hint. The default steps score classification. An MSE (or L1 / Smooth-L1 / Huber) model without NNModelParams(task=TaskSpec.regression(...)) whose targets are not whole numbers now gets that hint before its first update, and from evaluate(), not a scikit-learn metrics error.
• Imperative paradigm steps are refused as objective= before a run is reserved. This covers KD, feature KD, Mixup, CutMix, MoE, SimCLR, DPO and text contrastive steps, as diffusion and JEPA already were.
• A stopped epoch never reaches a full-history callback. With a history journal, a stop at an update boundary also discards the epoch's chunks that were already written.
• More late failures move earlier. born_again_train refuses gradient accumulation before generation 0. Trainer.train and NNModel.train without DDP() refuse a multi-rank process group. Trainer.train refuses update listeners and ctx.stop_at_update, which it never delivers. set_seed refuses a seed outside [0, 2**32) or a non-integer one. apply_ia3_to refuses to wrap a DoRA layer's base, and GGUF export refuses PEFT-adapted or prefix-tuned networks with a hint. decide_result(policy=...) checks the policy before calling the provider.
• Smaller fixes:
  • widen updates FeedFwdNN.params.hidden_dims.
  • DiffusionMLP and sample() run in the net's dtype.
  • bf16 modules work in the benchmark weights digest, weight_histogram and activation_map.
  • VisUtils.classification_report keeps a named class absent from the data.
  • load_provenance raises ValueError for JSON that is not an object.
  • lr_finder stops at a non-finite loss.
  • init_process_group validates timeout_seconds like shutdown().
  • A dynamo-skipped forward under on_failure="error" restores the model's buffers.
  • random_block_mask validates its ranges.
  • Safetensors checkpoints are flushed to disk before the rename.
  • nnx.export_conformance is bound on nnx.
  • A decision-benchmark metric with a value carries a note (rows it does not count) instead of an "unavailable" reason; the CSV gains a note column.
  • decide_result(policy=...) refuses a policy that is not an AbstentionPolicy.
  • A history-journal run stopped in its first epoch loads.
  • A damaged BEST checkpoint of the runs/best incumbent no longer fails a later run's final step; the pointer is re-elected.
  • OneCycleLR keeps each parameter group's own peak learning rate.
  • Mixup and CutMix draw from the torch RNG, so seed controls them and a stateful resume continues the uninterrupted run.
  • Contrastive, KD, MoE and augmentation settings refuse NaN and infinity.
  • deepen refuses an nn.Sequential subclass or a container with hooks, parameters or buffers; on a FeedFwdNN it keeps eval mode.
  • expand_embedding keeps the source's requires_grad and mode; PromptTuner refuses a prompt that fills max_seq_len.
  • GGUF export works for bf16 networks.
  • Trainer.train builds writer_only(...) callbacks before its checks, so their monitors bind.
  • weight_histogram and activation_map render large grids and validate their grid arguments; VisUtils.confusion_matrix refuses labels that are not class indices.
  • nnx.viz.attribute's figure shows the first sample of a 3-D attribution.
  • A prediction with zero rows accepts an empty sample_ids list.
  • NumPy scalars in params fields are stored as plain floats, so run.yaml serializes them.
  • A first run save that fails releases its reservation, so a retry is not refused.
  • Trainer.train refuses grad_clip_norm / accumulate_grad_batches that a trainer_step_fn never applies, and an objective run's optimizers that share a parameter.
  • A schedule whose optimizer was frozen in its first window advances without torch's misleading warning.
  • compare() of provenance manifests never equates a dotted key with a nested one.
  • A deeply nested provenance.json is reported as malformed, not a RecursionError.
  • split_links refuses non-integer edge categories; LinkPredictor refuses zero width / outputs; link_metrics refuses targets outside the declared domain.
  • validate_decision_request_result refuses options on a Boolean question.
  • The imperative JEPA step refuses an empty context or target mask by name.
  • NNOptimParams with an integer momentum (momentum=0) or betas trains; NumPy integers keep the Python-int run id.
  • NNTrainParams / NNTrainerParams flags must be bools.
  • lr_finder sweeps a negative loss and refuses non-finite arguments.
  • FixedHeadProvider refuses a regression loss without a task.
  • LogitsChainBuilder refuses custom processors after temperature(0.0), which would corrupt the greedy choice.
  • link_metrics refuses probabilities outside [0, 1] (pass logits with from_logits=True).
  • A fully masked task batch with infinite outputs no longer turns the evaluated loss into NaN.
  • OneCycle keeps a single scoped group's multiplier (a Trainer's strict optimizer).
  • A Trainer objective run refuses auto_step_schedulers=False with schedulers (nothing would step them); NNTrainerParams checks its optims / schedulers containers.
  • A custom eval step's extra-only validation record survives NNRun.load.
  • NNTransformerParams.tie_embeddings must be a bool; recipes accept NumPy integer r / rank; LoRA, prefix and prompt sizes must be integers, checked before the base is frozen.
  • NLIProvider.check refuses a bad request before any model call, so jobs and the benchmark spend no call on it.
  • The imperative diffusion and JEPA steps honour the model's batch adapter.
  • NoiseSchedulers validate T and the cosine s; save_pretrained refuses quantized weights before writing.
  • link_metrics refuses NaN logits and non-finite categorical logits; a 1-D attention_mask is an LMTaskError.
  • An objective run's whole-epoch training loss keeps each window's term weight and reduction (an annealed weight used the epoch's first).
  • CategoricalParam refuses choices that compare equal (True, 1, 1.0), which optuna would collapse.
  • NNParams.activation(s) and NNOptimParams.name must be enum members (a string failed far from where it was set).
  • link_metrics scores a -inf category mask in a row with a finite maximum again.
  • A LoRA / low-rank recipe bundle resumes; only a train-end conversion (QAT) is refused.
  • LoRA alpha / dropout accept NumPy reals; ExperimentManifest.for_model refuses train= beside config["train"].
  • A run's notebook table escapes every value.
  • Decision providers and Capabilities check their settings at construction (max_batch, ordinal, name, primitives).
  • The imperative JEPA step refuses a predictor outside model.net (it was never updated).
  • Re-applying an adapter in a mixed LoRA / IA3 stack wraps nothing again.
  • FixedHeadProvider reports non-finite logits as a ProviderFailure (new nnx.prediction.NonFiniteLogitsError) and a too-wide Bernoulli head as an InvalidDecisionRequest.
  • NNSchedulerParams refuses warmup_steps > total_steps; train_bpe always keeps an <unk> token, so unseen characters encode.
  • An optimizer built with a momentum of the wrong shape names the problem instead of a bare AssertionError.
  • A window dropped as non-finite (nonfinite="skip") no longer turns the whole-epoch training loss into NaN.
  • A warm-up longer than the run's default schedule length is refused when the scheduler is built.
  • predict / predict_proba accept NumPy float64 input; evaluation records print their metric values.
  • widen / deepen refuse an adapted or factorized neighbour layer before changing anything; GGUF judges the quantization label before requiring gguf.
  • A resume_epochs="planned" resume keeps the plan's provenance fingerprint.
  • Link tasks refuse a message graph with a repeated edge, and training negatives that cannot fit at LinkTask construction.
  • FixedHeadProvider checks label strings and refuses a non-2-D categorical output as a request error.
  • Library warnings (EarlyStopping, predict over a shuffling loader) no longer crash in a REPL, python -c or a notebook.
  • An objective epoch whose every window is dropped as non-finite reports a NaN training loss, so EarlyStopping counts it; an objective whose terms are all detached is refused instead of training nothing.
  • load_lora_weights / load_ia3_weights / load_prefix_weights change nothing when they refuse a wrong-shaped key.
  • The diffusion objective refuses a denoiser output of the wrong shape; bf16 / fp16 denoisers train in their own dtype.
  • decide_result returns an invalid_decision_response error for an answer of the wrong count or question; question_from_state refuses keys its kind does not take.
  • A numeric setting given as an integer too large for a float raises a ValueError naming the field.
  • A failed eager re-run under compile= restores the model's RNG and buffers before its error propagates.
  • The diffusion step and both JEPA paths refuse a prediction of the wrong shape instead of broadcasting it into the loss.
  • ExperimentManifest(config=None) reloads with its fingerprint, so the run keeps its provenance.
  • Construction refuses a plateau scheduler factor >= 1, ModelCheckpoint epochs that can never match, and non-string NNTrainerParams optimizer / scheduler names; lr_finder refuses a non-integer num_iter before the sweep.
  • netron_export(launch=True) checks for netron before writing the ONNX file.
  • low_rank_factorize (and low_rank transform recipes) and offline-distillation scoring work on MPS models.
  • A refused task evaluation leaves the network's training modes as it found them.
  • MoELinear, ViTNN and ViTBlock refuse float or bool sizes and dropouts at construction.
  • Regression tasks, causal-LM training, Standardizer, the NLI provider and the diffusion objective work on MPS (no float64 there); the compile snapshot restores the MPS generator.
  • Task and probability labels given as NumPy strings are stored as plain strings, so the run saves its run.yaml.
  • widen refuses a container with its own forward or hooks instead of silently changing the model's output.
  • generate(max_new_tokens=...), TopKFilter, Bootstrap and decision results validate their inputs at construction.
  • Near-greedy sampling (a small temperature) on half-precision logits picks the argmax instead of the first overflowed token.
  • Offline distillation scores an fp64 student in its own dtype; to_onnx takes a float64 NumPy example.
  • NumPy strings and bools in model, metric, optimizer-factory and provenance configs are stored as plain values; decision-result ids too.
  • deepen refuses a hooked target ReLU; widen looks only at hooks on the path to its consumer.
  • Trainer refuses a stop_at_update set in on_optimizer_update at that update, with advice that applies.
  • NNPreferenceDataset validates its lengths and pad id before encoding; NNCheckpoint.to_file refuses an unknown format before creating a directory; split_links refuses a bare-string categories.
  • Knowledge distillation refuses single-logit heads (the soft loss was always 0) and a teacher of another shape; distillation, Mixup / CutMix and MoE steps threshold a BCE head's predictions instead of taking an argmax over one logit.
  • Run comparison reads idps.csv floats exactly (the default parser was off by one ULP on about one value in five).
  • A refused or failing on_train_begin after a resume hands the model back unchanged (NNModel and Trainer).
  • deepen refuses a FeedFwdNN / ConvNN subclass with its own forward; wildcard LoRA / DoRA / IA3 refuse a MultiheadAttention projection; ExperimentManifest keeps labels given as a one-shot iterable.
  • generate() validates stop (a list of non-empty strings) and seed.
  • A weights-only resume, and a refusal at the first on_epoch_begin, also hand the model back unchanged.
  • A FeedFwdNN on graph split batches scores only the split's seed rows, as the graph nets do.
  • Distillation with alpha=0 trains a binary head again; born_again_train refuses a single-logit head before generation 0; distillation, MoE and Mixup / CutMix records read probability, soft and ignored targets as the default step does.
  • Wildcard PEFT refuses TransformerEncoderLayer / TransformerDecoderLayer projections; benchmark Record.execution must be strict JSON; a registered run's summary shows its factory.
  • NNModel.train also hands a resumed model back when the first on_epoch_begin raises; the progress bar names the value it shows (val.loss, train.error).
  • Transform recipes refuse a Linear its parent reads directly, as the PEFT helpers do; distillation and MoE steps refuse regression targets before the update.
  • A ModelCheckpoint tag shaped like a training-state sidecar is no longer deleted by the next save; SGD_NESTEROV needs momentum > 0 and NNSchedulerParams.kind must be a Schedulers member, both at construction.
  • export_bundle refuses the best pointer (export the run by its own id); forward benchmarks digest modules with non-tensor state; benchmark records' execution must survive the JSON round trip.
  • The loops warn when a parameter requires gradients but no optimizer owns it (unfrozen after a param_groups optimizer left it out), instead of silently never training it.
  • Mixup / CutMix refuse regression targets before the update; categorical tasks score int32 / int16 targets; GGUF export refuses a factorized projection with a clear TypeError.
  • Forward benchmarks digest torchao-quantized and packed-quantized weights exactly; validation messages show integer bounds exactly; a DDP step-count refusal on a resume's first epoch also rolls back.
  • PrefixTuner places prefixes by the projection's first floating tensor (a factorized w_qkv works) and refuses before freezing the model; embed_texts / export_to_faiss return embeddings on the requested device without moving the backbone.
  • QAT refuses a non-int groupsize and a run whose groupsize divides no Linear; a tiny temperature samples the argmax; KD / Mixup / CutMix / MoE refuse a sequence model's output.
  • The unowned-parameter warning names both causes and points at the caller's train(); t-SNE projection refuses a non-categorical checkpoint before inference; benchmark digests cover sparse state; an empty held-out benchmark slice reports "no eligible rows".
  • Re-applying IA3 over an IA3 → DoRA stack is a no-op; an Ollama float parameter too large for a float is a TypeError.
  • NNTransformerParams dropout fields refuse non-numbers like dropout_prob.
  • Training-loop and history warnings (unreported optimizer updates, an unfed monitor, a skipped plateau step, a resume without training state, worker RNG, component rollback and cleanup failures, lineage gaps) point at the caller's train(), so a second run in a notebook warns again; kd_objective refuses a sequence model's output as the KD step does.
  • NNModel.train refuses a registered net that still carries an unbaked pruning mask (its checkpoints could not rebuild it); magnitude_prune(bake=False) refuses a MultiheadAttention projection; KD with alpha=1 ignores the hard labels (an all-ignored batch no longer fails).
  • A tiny temperature on all-negative logits samples the argmax, not token 0; the QAT "prepared nothing" error names a root Linear; the DPO beta docstring states the right direction (larger beta keeps the policy closer to the reference).
  • CausalLMTask refuses an attention mask with padding before a real token (left padding): the model never receives the mask, so the NLL was not independent of padding. Benchmark weight digests include each tensor's dtype and shape (and a per-channel quantization axis).
  • A lazy-layer model's final LAST records the materialized optimizer topology, so a materialized module resumes from last; eval_step_fn runs with the net in eval mode (modes restored after); to_onnx refuses a mapping example with a clear TypeError.
  • NNModel.train refuses a registered net whose weights a parametrization replaced (weight_norm, spectral_norm), since no checkpoint could reload it; distillation refuses a flat (N,) single-logit head (it was softmaxed across the batch); train_contrastive refuses lr <= 0; a very negative mask value no longer makes a normal temperature greedy; predict, to_onnx and netron_export on a read-only (pandas) array no longer warn.
  • from_pretrained refuses an explicit params=, net_params=, tokenizer= or transforms= instead of silently ignoring it; the parametrization refusal also covers a registered module's top-level parameters.
  • KD, Mixup / CutMix and MoE steps on an MLP over graph split batches read only the split's seed rows (they scored other splits' nodes); born_again_train refuses a single-output registered or runtime nn.Sequential before generation 0; a Jev provider with only an async client is refused before a synchronous job or benchmark call; declared metrics named like a record field log as metrics/<name>; a failure to record a failed attempt never replaces the training error, even under -W error.
  • from_pretrained also refuses a dict or list override that differs from the repo's config.json; a Jev provider with only an async client is refused before a synchronous job or benchmark counts a call, even inside an event loop (a notebook cell); NNDataset checks the training transform's output is a tensor too; an explicit REDUCE_LR_ON_PLATEAU scheduler prints its plateau settings; VisUtils.multi_line_plot places its ticks on the x values.
  • A synchronous decision job or benchmark no longer blocks a provider that runs its own event loop (asyncio.run inside decide); building an NNDataset with a random training transform no longer moves the global random streams; born_again_train refuses a TransformerNN before generation 0; NNPreferenceDataset refuses a pad id outside the vocabulary and dpo_train_step_factory a bool or negative one; widen, low_rank_factorize and expand_embedding refuse non-integer sizes.
  • ExperimentPlan.validate() reports an unregistered optimizer factory and an update-clock scheduler budget a custom step cannot plan, so a search stops before its first trial instead of failing every trial; LinkPredictor casts float64 node features to its parameters' dtype, as GraphClassifier does.
  • inspect_bundle reports a built-in net's net_params (its dimensions), so an untrusted bundle's architecture can be read before reconstruction; NNPreferenceDataset checks a pad id against the tokenizer's full length (added tokens included); a warning attributed to a plain script is no longer replaced by a DeprecationWarning under -W error; a seeded NNDataset with a random base transform keeps the global random streams; FeedFwdMoENN resolves its activations once, so torch.compile can trace it on older torch.
  • An explicit train_step_fn=default_train_step (what qat_train_step_factory returns) records the declared training metrics and supports a training-split monitor; under DDP a batch one rank cannot score stops every rank instead of hanging the others; VisUtils.confusion_matrix without names labels the classes present; calling a Standardizer on a DataFrame checks its columns; activation_map reads a layer that returns a mapping; a bootstrap interval whose bounds meet while resamples vary says so.
  • ExperimentPlan also reports an update-clock scheduler budget over a training loader without a length (a factory's loader when fit() calls it); TeacherDataset takes read-only and reversed NumPy inputs.
  • predict / predict_proba take a full-graph loader list (NNGraphDataset(sampler="full")) as evaluate and iter_predict do; dpo_train_step_factory refuses a bool or non-numeric beta; NNPreferenceDataset reports the tokenizer's full length (added tokens included) as its input_dim / output_dim.
  • RankingTask.eval_step() scores an MPS model (its scores widen to float64 on the CPU).
  • An explicit default_train_step also gets the checkpoint-reconstructibility guard and a search's updates_per_trial cap, as the implicit default step does.
  • widen, low_rank_factorize and expand_embedding refuse a weight_norm / spectral_norm module (parametrization or legacy hook) or an unbaked prune mask instead of silently dropping it, reading a stale weight and unfreezing it.
  • Surgery refuses a subclass of nn.Linear / nn.Embedding / nn.ReLU with its own forward (a fake-quantized QAT Linear, say) instead of rebuilding it as a plain module.
  • A resume that fails after restoring its state (a warnings-as-errors filter, say) hands the model and RNG back unchanged in a single-process fit too, through NNModel.train and Trainer.train.
  • TopPFilter(1.0) (and generate(top_p=1.0)) keeps every token; float32 rounding had dropped some.
  • A resume whose first epoch fails before any batch (a loader set_epoch, or the unowned-trainables warning raised as an error) hands the model and RNG back in Trainer.train too.
  • validate_response refuses a vendor_score for a Choice or Boolean instead of dropping it; ChoiceResult / ScoreResult refuse a distribution that is not (id, p) pairs with InvalidDecisionResponse.
  • FeedFwdNN, FeedFwdMoENN and ConvNN accept non-contiguous inputs (a transposed NumPy array, a channels_last tensor or net); MoELinear names a malformed count with ValueError.
  • load_provenance refuses JSON nested past 512 levels as "nested too deeply" on every Python version (3.14 parsed it and reported another error).
  • drop_layer refuses a NaN or infinite importance score; diffusion_train_step_factory refuses a non-NoiseSchedule before training.
  • two_dim_tsne_checkpoint_logits refuses multi-output targets before the projection; multi_line_plot legends every drawn dash style and checks label counts; weight_histogram counts NaN / Inf values in each subplot title.
  • DoRA refuses an nn.Linear subclass with its own forward (a QAT fake-quantized layer), at wrap time and when QAT prepare swaps the base later, instead of bypassing it; a PEFT wrap that fails part-way leaves every layer and its trainability unchanged.
  • benchmark_forward / compare_compile no longer recurse without end on a tensor-subclass wrapper buffer; random_block_mask refuses a non-finite block_aspect with ValueError; the docs builders' link validators read prose after an inline code span as prose too.
  • A Trainer run whose callbacks convert modules at train end (QAT) resumes statefully from its final LAST: the checkpoint records the optimizer topology from before the conversion, as NNModel.train does.
  • nnx.viz.summary(model, input_size=...) builds its dummy input on the model's device instead of moving an MPS (or, on a CUDA host, a CPU) model; validate_record refuses a malformed weights_sha256.
  • netron_export casts a NumPy float64 example to the net's dtype, as to_onnx does; summary(input_size=) builds its dummy in the model's floating dtype (a .double() model) and on its device even without parameters; benchmark coverage reasons count one fault across samples (messages no longer carry the sample id); the wiki and site builds no longer read prose after an inline code span (`x` [a]: ...) as a link reference definition.

0.4.1 (2026-10-11)

Bug Fixes

• release: build the artifact smoke venvs on the runner's work disk (76b6706), closes #445

Documentation

• fix the site banner path and the re-audit's follow-ups (#472) (cb36cd1)
• repair audited documentation across README, site and wiki (#472) (812d6d1)
• repair the audited documentation across README, site and wiki (#472) (ff082fc)
• version-aware install guidance and the audit's follow-ups (#472) (b8e12f6)

0.4.0 (2026-10-09)

Added

• JEPA-trained encoders rebuild without the predictor from bundles and the Hub. reconstruct_bundle(..., exclude_submodules=...) and NNModel.from_pretrained(..., exclude_submodules=...) leave out the named top-level submodules (the _jepa_predictor the JEPA recipe registers under model.net) as from_checkpoint does since #395, every other weight still strict (unlike strict=False, a missing weight is refused). When the weights hold submodules the rebuilt net lacks, the error (a strict load's, or a registered module's schema check) names them and the option of the entry point called, with any names already passed — a bundle's no longer suggests from_checkpoint, and the Hub's no longer surfaces torch's bare error. A bundle rebuilt without a submodule serves inference and refuses resume() (#441).
• ExperimentPlan.resuming(..., epochs=). A plan can ask for resume_epochs: epochs="planned" continues the source run's plan up to n_epochs in total, "additional" trains n_epochs more, and None (default, unchanged) keeps the training parameters' own value. validate() checks it as NNTrainParams does (resume.epochs) and reports a "planned" horizon with a "weights_only" mode or with no run to resume (including one set on the plan's training parameters) before any factory runs or run is reserved; at "additional" it stays out of state(), so run ids are unchanged (#442).
• Exact resume to the planned horizon (resume_epochs="planned") and verified resume points. NNTrainParams.resume_epochs chooses what n_epochs means on a resume: "additional" (default, unchanged) trains n_epochs more; "planned" continues the run's own plan up to n_epochs in total, carrying on the logical step counter, ctx.committed_updates and the phase-checkpoint cadence, so on deterministic single-process training with in-process loaders a run split at epoch boundaries (and interrupted any number of times) equals the uninterrupted run bit for bit — with Adam or SGD, ReduceLROnPlateau acting; it needs a stateful source that records its plan — a weights-only resume, a pre-#394 checkpoint, another planned total or a finished plan is refused before anything is restored (a Trainer always resumes "additional"). Every checkpoint is now a resume point stamped by a manifest (<tag>.pt.manifest.json: generation ordinal, checkpoint id, logical position, SHA-256 per file), staged before the checkpoint is published and made live right after, with files fsynced and a Ctrl-C held for that window — a writer that dies inside it leaves a point verification reads from its staged manifest (verifying never writes); a resume verifies it first (NNCheckpoint.verify, NNCheckpoint.resume_point) and refuses a torn, mixed or corrupted point with ResumePointError naming why. Checkpoints written before manifests load as before. run.resume_status (in metadata.yaml, never run.yaml) adds source_generation and source_checkpoint_id, and callbacks read the scheduler through ctx.scheduler (#394).
• A public ViT network type: Nets.VIT and NNViTParams. NNViTParams (a NNParams subclass) carries every ViTNN setting — image_size, patch_size, in_channels, d_model, n_layers, n_heads, ffn_mult, attn_dropout, resid_dropout — validated at construction (positive integer sizes, image_size % patch_size == 0, d_model % n_heads == 0, and the encoder's base-field convention: input_dim = in_channels * image_size ** 2, output_dim = d_model, no hidden layers, dropout_prob = 0.0). Nets.VIT(params) builds ViTNN; NNRun.load returns the stored NNViTParams and NNModel.from_checkpoint rebuilds the encoder with its weights (through NNParams.resolve_from_state, no caller-side swap). NNModel.from_checkpoint(..., exclude_submodules=("_jepa_predictor",)) rebuilds a JEPA-trained encoder without the predictor the JEPA recipe registers under model.net (every other weight still loads strictly; an unknown name is refused). Task adapters refuse Nets.VIT up front (an encoder, not an (N, C) head). The architecture is always in state(); in_channels, ffn_mult and the dropouts are omitted at their defaults, so existing run ids are unchanged. docs/jepa.md and example 16 build the encoder this way (#395).

• A bounded DDP teardown: nnx.distributed.shutdown(timeout_seconds=60). The counterpart of init_process_group(): every rank meets at a teardown barrier, then the process group is destroyed, so no rank tears down its Gloo connections while a peer still uses them (a peer could otherwise abort with "terminate called without an active exception"). With Gloo, a rank that does not arrive in time makes the others raise RuntimeError within the timeout instead of blocking in the call; the group is then left for process exit, which may still wait for the missing rank or the group's own timeout (with NCCL on older torch the group's timeout bounds the wait). Without a group, or after a successful call, it does nothing. The torchrun recipes (docs/quickstart.md, docs/architecture.md §3) and examples/ddp_supervised.py end with it (#419).

• A published feature-composition matrix. docs/feature-composition.yaml (FEAT-039) registers which feature combinations are verified, unsupported or unverified. Each scenario across training controls, checkpoint/resume, model transformations, inference/export, decisions and typed results, and packaging (FEAT-010, -023, -025, -029, -030, -031, -033 and -037 each appear in at least one combination) records its device / dtype / dependency profile, expected boundary, pytest node ids and rationale. scripts/check_feature_composition.py validates the schema and unique IDs, collects every node id (a stale reference fails naming its scenario), runs the scenarios with a phase-recording plugin (an import, a skip, an xfail or a setup / teardown failure never verifies), and renders docs/feature-composition.md from that evidence report. A report from other sources or another profile renders unverified. A new CI feature-composition job gates verified rows and keeps the report; the docs site and wiki render the same report (identical evidence digest). CUDA AMP and ONNX Runtime rows stay unverified on CPU (#256).

• Decision-model research pilot (experimental, no-go). examples/decision_model_pilot.py (FEAT-023) scores (state, question, candidate description) triples with a shared encoder and scalar head registered through register_model_factory (no new Nets member): Choice decisions softmax the candidates (cross-entropy) and Boolean ones take one logit (binary cross-entropy), both trained; Score is refused. A validated pre-execution manifest (licences, disjoint fit/select/held-out family splits, schemas, seeds, baselines, compute cap, stop rule, model identity, predeclared trade-off) comes first. The offline mechanics mode enforces the compute cap and produces a checkpoint reloaded through the registry, candidate-aligned replay records, trial and stop records, and a no-go report. The empirical mode writes its manifest and reports blocked: the resources are not supplied, and the empirical runs are outside this repository. See docs/decision-model-pilot.md (#262).

• Budgeted experiment search through Optuna (nnx.search, extra optuna). search(plan, space, apply=..., monitor=..., budget=SearchBudget(trials=, updates_per_trial=, deadline_seconds=), study_name=, storage=, sampler=, pruner=) runs a sequential Optuna study in which each trial's parameters become a fresh ExperimentPlan via apply (the base plan is unchanged; fresh model, optimizer, loaders and callbacks; its own attempt id and run), scored by one validation monitor equal to the plan's own. The space (FloatParam, IntParam, CategoricalParam), budget, monitor and plan are validated before any trial, and a stored study with a different identity is refused. Completed, failed, pruned and orphaned trials all count against the budget. The update cap stops right after the N-th committed update, a deadline starts no trial once expired and stops a running one at its next boundary (recording the overrun), and a trial stopped mid-epoch keeps only its committed epochs. Pruning sees each validated epoch's value at step = epoch; missing or non-finite values fail. A trial's value is its run's monitor decision (the epoch its BEST checkpoint holds). A failed fit is recorded with its message and the search goes on; a configuration error in apply stops it. Every TrialOutcome keeps parameters, terminal reason, resources, monitor value and attempt / run / checkpoint links, and SearchResult.best is the best completed trial. The training loop gains ctx.update_listeners and ctx.stop_at_update (stop at an update boundary without committing the epoch in progress). import nnx never imports Optuna. See docs/concepts.md §29 and examples/search_offline.py (#266).

• Single-node data parallelism under torchrun (nnx.distributed). model.train(params, distributed=DDP()) trains over the process group torchrun started (nnx.distributed.init_process_group(); Gloo on CPU, NCCL with one GPU per rank) through PyTorch's DistributedDataParallel, with every committed update equal to one process training on the ranks' union batch: gradients synchronize only at committed updates and are normalized by the global loss denominator (unequal valid counts included), an all-masked global window is skipped everywhere, and non-finite losses, step-count mismatches, preflight refusals and writer-side commit failures fail on every rank together instead of hanging. train_loader(dataset, batch_size, policy="pad"|"drop", seed=) partitions each epoch's seeded permutation evenly; validation_loader shards rows unpadded so every id is scored exactly once. Records are global and identical on every rank (monitors, early stopping, BEST and schedulers agree) and every rank returns the same run; only the writer rank holds the lease and writes history, checkpoints (canonical keys; world size, partitions and every rank's RNG for a same-world resume) and provenance. Callbacks declare distributed = "all" or "writer"; writer_only(factory) builds a TensorBoardCallback / WandbCallback on the writer alone. FP32, default steps, one optimizer, one node, no batch normalization. See docs/architecture.md §3, docs/quickstart.md §3.10 and examples/ddp_supervised.py (#265).

• Opt-in torch.compile and bounded benchmarks (nnx.compilation, nnx.benchmarking). model.train(params, compile=CompileSpec(backend="inductor", mode=None, fullgraph=False, dynamic=None, on_failure="error")) runs the built-in step's FP32 forward (and the fit's validation) through a per-call compiled wrapper beside — never replacing — model.net, so checkpoints keep eager keys, optimizers own the eager parameters and export, summary, tying and later inference stay eager. Refused up front with a custom step or objective, a graph or custom net, reduced precision or a topology-changing callback (QAT); not for DDP or quantization, and no speedup is guaranteed. on_failure="eager" re-runs only the failing forward eagerly (never a partly applied update). CompileRecord (requested backend/options/policy/torch version; effective compiled with full or partial capture, eager with its restart, or failed) is on NNRun.compile, in metadata.yaml, checkpoint training state and attempt.json, never in the run id. benchmark_forward / compare_compile separate first-call compile cost from warmed latency, throughput and dispersion from identical weights; profile_forward runs a finite torch.profiler schedule into an explicit directory. See docs/concepts.md §28 and examples/benchmark_compile.py (#264).

• An opt-in Result at fallible boundaries (nnx.result). Frozen Ok / Err with map (only an Ok), map_error (only an Err), bind and recover (each flattening one Result its callback returns; anything else raises TypeError) and unwrap (raising UnwrapError with the typed error). Combinators never catch a callback's exception, and a Result has no truth value — Ok(0), Ok(None) and Ok(False) are successes. Three boundary wrappers convert exactly the exceptions they declare into Err(BoundaryError(code, where, context, cause)): validate_decision_request_result builds a typed question from plain data (InvalidDecisionRequest → invalid_decision_request, no provider call), decide_result separates a provider failure (provider_failure, or invalid_decision_response for an answer that does not fit the question — a malformed Jev answer included — each with its request id; unsupported) from a successful abstention (an Ok of SelectiveDecisions, probabilities and reasons kept), and inspect_bundle_result delegates to the safe inspect_bundle (artifact_missing, bundle_invalid; never unpickles). Every exception-style API is unchanged. See docs/decisions.md §10.1, docs/concepts.md §27 and examples/result_boundaries.py (#263).

• Ordered logits pipelines (nnx.generation.OrderedLogitsPipeline). Decoding stages that run exactly in their declared order, beside the canonically sorted LogitsChain (unchanged): OrderedLogitsPipeline.of(LogitsStage.temperature(0.5), LogitsStage.top_p(0.8)) keeps duplicates (two temperature-0.5 stages equal one 0.25) and custom positions, and append / prepend return new pipelines over immutable LogitsStage specs — a processor passed in is copied by value, and every run compiles fresh processors. Values are validated (finite, in domain) when a stage is made; a zero-temperature stage is greedy and terminal, so a stage after it raises instead of being reordered. generate(logits_pipeline=...) compiles the pipeline once per call on the cached and uncached paths, with token history, on_token, stop and seeded output equal to the explicit processor list; passing it with logits_chain raises before the model is touched. Built-in pipelines round-trip through state() / from_state() ({"version": 1, "stages": [...]}); a CustomStage is runtime-only (its state stays the caller's) unless its tag has a LogitsStageCodec registered in-process — data never names code to import, and a pipeline never enters a checkpoint or a run identity. See docs/lm.md §5.1 and examples/ordered_logits_pipeline.py (#254).

• Typed decisions on Jev models through the TypeSafe SDK (nnx.decisions.JevProvider). A new jev extra (pip install "thekaveh-nnx[jev]", pinned and tested at typesafe-sdk>=0.7,<0.8) adapts Choice, Boolean and Score questions to the SDK's system_one call: one request per text holding every question (a Choice keyed by its options' descriptions, never their ids; a Score by its levels in order; a Boolean as a noul), answers realigned to each question's own order and validated, a Score's expected score in vendor_score. Each result's raw is plain JSON — the model the service resolved (even for an alias request), the requested model, the request id and reported token usage, absent fields left absent, never clients, headers or credentials. A Choice's or Score's confidence is recorded as provider_confidence, Jev's own certainty rather than a probability of being correct; a Boolean gets none. The SDK's RetryPolicy owns retries and the adapter sends each request once; what still fails becomes JevTimeout, JevAuthenticationError, JevRateLimited (with retry_after_ms), JevMalformedResponse or JevError, all ProviderFailures keeping the SDK error as __cause__ and the service's request_id, through a decision job (JobFailed.__cause__) or the benchmark (a record's reason) too. Injected sync / async clients stay the caller's, and the provider never builds a second client beside one (a sync-only provider has no async methods, so a decision job runs it in its own worker thread); clients it builds itself close on with (sync) / async with (both) exit, after an error or cancellation too. Each text is one request, so call budgets count provider calls, not requests — bound them with max_batch. import nnx never imports the SDK, and a provider that would need it raises ImportError naming the extra. scripts/smoke_jev_live.py is an opt-in live smoke pinned to jev-1.13.0. See docs/decisions.md §6 and examples/decision_jev.py (#220).

Changed

• A DDP resume verifies its resume point once. Under distributed=DDP(), the writer rank verifies the point and shares the manifest or its refusal with the other ranks, which no longer hash it one after another under the tag's lock; a torn, mixed or corrupted point is still refused on every rank, with the same ResumePointError reason. Single-process resumes are unchanged (#440).
• Resume points are hashed while they are written. Each checkpoint save computes its SHA-256 digests from the bytes it writes (large chunks on a helper thread overlapping the write) instead of reading the checkpoint and its training-state sidecar back; the manifest format and NNCheckpoint.verify are unchanged (#439).
• The vision, graph and plotting stacks are extras (thekaveh-nnx[vision], [graph], [plots]; [domains] for all three). torchvision, torch_geometric and Plotly are no longer core dependencies, and import nnx loads none of them: the core install trains, predicts, saves and reloads models with torch, NumPy, pandas and scikit-learn alone. Features that need one import it on use and, when it is missing, raise an ImportError naming the extra — NNDataset (vision); the graph nets, NNGraphDataset, Nets.GRAPH_*, nnx.graph_tasks and nnx.link_tasks (graph); VisUtils figures, nnx.viz figures and nnx.lr_finder (plots, which refuses before touching data, model or RNG state). An ImportError from inside an installed package keeps its own cause. With the extras installed, flat and deep names, class identity, signatures and Nets values are unchanged; nnx.__all__ lists a domain name only when its extra is installed, so from nnx import * works on a core install. Without Plotly, NNRun's notebook repr keeps its table and drops the chart, and VisUtils.classification_report still returns a DataFrame. The viz extra now pulls Plotly for its figures. Upgrading: install thekaveh-nnx[domains] to keep the previous dependency set. scripts/smoke_core_install.py checks an installed wheel per profile, run by new CI installed-profiles and release artifact-profiles jobs over the minimum and current torch pairs; the floor-deps lane installs [domains] (#249).

Fixed

• Resume points reach stable storage on macOS. Every resume-point flush (checkpoint, training state, manifests, directory) uses fcntl(F_FULLFSYNC) where the platform has it — macOS's fsync only reaches the drive's cache, so a power loss could lose a point its manifest called complete — and os.fsync elsewhere, as far as the drive honours the flush. A filesystem that does not support the flush does not fail a save; a real I/O error still does (#443).
• NNRun.load reads a run's history back exactly. idps.csv is now parsed round-trip: pandas' default float parser misread about one value in five by one ULP, so a loaded run's losses, learning rates and metrics could differ from the trained run's in the last digit.
• nnx.search.search() and ExperimentPlan.fit() refuse a multi-rank process group. Both train in a single process (a plan never passes distributed=); launched under torchrun, every rank ran its own study or fit into the same working directory. Inside an initialized torch.distributed group of world size > 1 they now raise RuntimeError before loading or creating a study, calling a factory, seeding or creating a run; a one-rank group behaves as before. Train one model across ranks with model.train(..., distributed=DDP()) (#417).

0.3.0 (2026-10-04)

Not published to PyPI; its changes first ship in 0.4.0.

Upgrading from 0.2.3

Default training is unchanged: the same configuration gives the same run id, losses and weights, and runs, checkpoints and Hub folders written by 0.2.3 load and resume (unless their configuration is one this release refuses, below). Check these before upgrading:

• PyTorch 2.4 or newer. The declared floor is now torch>=2.4 / torchvision>=0.19 (see Changed).
• Losses.NEGATIVE_LOG_LIKELIHOOD trains a different objective under the same run id. NLL now applies log_softmax before NLLLoss (see Fixed), but a run id is a hash of the configuration, which did not change: re-running a 0.2.3 NLL configuration resolves to the 0.2.3 run's id, so it is refused as an existing run (or replaces it with overwrite_existing=True), and resuming a 0.2.3 NLL run continues it with the corrected loss. Give the new run a distinct data_id to keep both.
• EarlyStopping() picks its monitor. The default monitor=None follows the validation error, or the validation loss when the evaluator reports no error, so regression runs, which never stopped early before, now can; mode="max" needs an explicit monitor, and a monitor with nothing to read (for example no validation loader) warns once per run (see Fixed).
• Stricter configuration. Configurations 0.2.3 accepted now raise ValueError at construction: NaN, ±infinity or non-real hyperparameters, counts that are not int (an integral float such as patience=8.0 included), batch_sizes entries other than a positive integer or None, and tabular features that are ±infinity or do not fit feature_dtype (see Fixed). The same check applies when a saved run is loaded, so a 0.2.3 run whose run.yaml recorded such a value (patience: 8.0) no longer loads and is not a runs/best candidate. A few operations that went ahead silently now refuse with ValueError too: widen across a Softmax or normalization boundary, and apply_lora_to / apply_dora_to / apply_ia3_to on a Linear registered under more than one name.
• A resume keeps the checkpoint's plateau rule or stops. Resuming with a ReduceLROnPlateau scheduler whose direction, threshold mode or threshold differs from the checkpoint's (for example a changed NNSchedulerParams.threshold) now raises ValueError before anything is restored, where 0.2.3 silently kept the saved rule; resume with the original settings or pass resume_mode="weights_only".
• No stateful resume back on 0.2.3. Runs trained by this release write training-state format 4, which 0.2.3 refuses to resume (it reads 1–3). Checkpoints of runs that use only features 0.2.3 has still load there. A pickle checkpoint of a run using a new feature (a named monitor, a TaskSpec, a registered model, a precision policy) refers to modules 0.2.3 lacks and does not load there; its safetensors checkpoint or Hub folder may (a named-monitor or TaskSpec run's does).

Added

• Offline teacher-distribution datasets (nnx.paradigms.offline_distillation). Distill a student from stored, permissioned teacher probabilities with no live teacher. A TeacherRecord carries the sample and question ids, the ordered candidate ids and their probabilities (finite, nonnegative, summing to 1 within 1e-6), the teacher and its revision, the schema digest, the probability semantics ("predictive", "calibrated" or "empirical") and its provenance — the source and the declared training_rights (recorded, not verified); a missing field fails validation, and write_teacher_records / read_teacher_records keep them as strict JSONL (nnx.teacher-record/1). TeacherDataset(records, inputs, candidates=..., schema_digest=...) joins records to inputs by sample id and aligns every distribution by candidate id, refusing a missing, unknown or duplicate candidate, another schema digest, a repeated sample or a sample without input when it is built — before any optimizer exists. dataset.objective(alpha=...) is an objective for the shared update engine: alpha · KL(teacher ‖ student) + (1 − alpha) · CE(student, label) at temperature 1 (no invented logits, no T²), both terms over one shared mask and denominator per microbatch so accumulation matches the full batch, partly labelled datasets and batches refused, each batch's candidate order checked against the objective's, teacher probabilities detached, and no backward, zero_grad, step, scaler or scheduler call of its own. Its identity (objective version, schema digest, alpha, candidate alignment) is checkpointed component state, so a stateful resume with any of them changed is refused before the first resumed update. evaluate_offline and dataset.eval_step() report teacher agreement and imitation loss over every record and labelled accuracy / NLL over the labelled records, each with its own denominator (the validation record's loss is the imitation loss; monitor label_accuracy / label_nll by name for quality). Nothing reaches the network. kd_train_step_factory, feature_kd_train_step_factory, kd_objective and born_again_train keep their logit / temperature contract and refuse these records, as do the default supervised steps. See Concepts §10.1, Typed decisions §9 and examples/offline_teacher_distillation.py (#244).
• Executable ONNX export conformance (nnx.export_conformance). Structural validity and numerical parity are now separate, recorded claims. Two profiles — feedfwd-fp32-torchscript and feedfwd-fp32-dynamo — export FeedFwdNN 4-8-2 (ReLU, dropout 0, deterministic nontrivial FP32 weights) at batch 2 with dynamic_batch=True and execute it in ONNX Runtime on CPU: every raw logit at batches 1, 2, 3 and 7 must match native eval at rtol=1e-4, atol=1e-5, and a wrong feature width must be refused. Dependencies, export, check, load and parity have separate outcomes, with failures classified as missing-dependency, export-error, state-changed (parameters, gradients or mixed train / eval modes changed by the export), checker-error, hash-mismatch, load-error, mismatch, input-contract or runtime-error. Each record (nnx.export-conformance/1, checked by validate_record) carries the source revision, config and weights digest, artifact SHA-256s (external data included), exporter options, requested and actual opset, dependency and runtime versions, provider, dtype, input cases and tolerances. verify_artifacts and execute reject a changed, missing or unlisted file before anything is loaded. scripts/check_export_conformance.py --output conformance.json runs both profiles, with an exit code per failure class (10–18; --force clears earlier artifacts). A new required CI job, export-conformance, does four things: it checks that import nnx works without the runtime, and that the check then fails with exit code 10; it runs the profiles and their tests with zero skips allowed; and it uploads the JSON with its ONNX and external-data files. A new optional extra, onnx-runtime, has been added, and uv.lock has been regenerated. examples/04_onnx_export.py labels results checker-only (no parity claimed) or executed, while examples/12_quantize_int8.py, Netron and GGUF / Ollama stay structural or container-only. A new docs page, docs/export-conformance.md, renumbers the later pages (#255).
• Leakage-aware link and edge prediction (nnx.link_tasks). split_links(edge_index, num_nodes, val=, test=, seed=, directed=False, self_loops="bar", negatives=1) splits a static homogeneous graph's edges into a versioned manifest (LinkSplit, nnx.link-split/1): canonical edges (an undirected edge is (min, max), so duplicates and reverses cannot cross splits), fixed validation / test negatives drawn from the static complement (never a positive, a reverse, a duplicate or a barred self-loop; a request beyond capacity fails before sampling), the seed and the topology policy; replay(edge_index) reproduces membership. LinkTask(split, mode="binary" | "edge_label") builds candidate batches whose message graph is the training edges only — for every split — and checks each batch before any forward pass (a held-out positive in the message graph is named as a leak; a candidate's target and id must match the split, and training negatives never reuse fixed held-out ones). The objective trains on train batches only, the training evaluator reads the val split with every candidate exactly once, and the default train / evaluate / predict paths refuse link batches. Every training pass re-draws its negatives and shuffles its rows from (seed, epoch) — NNModel.train and Trainer.train now call train_loader.set_epoch(epoch) before each epoch when the loader defines it (PyTorch's DistributedSampler convention) — so a stateful resume draws what the uninterrupted run would have; a wrapper must forward set_epoch to the task's loader (materialised batches replay as they are); the negatives' seed is checkpointed and a resumed loader with another seed is refused. Edge-label mode keeps category order, declares whether labelled edges' existence is visible context and samples no negatives. Binary candidates give (K,) logits, categorical (K, C), aligned with ids, pairs and targets through predict(); evaluation materialises every candidate for exact AUROC / AP, scored from the float64 logits so a confident model's candidates are never tied by rounded probabilities (BCE -(ln .8 + ln .7 + ln .6 + ln .9) / 4, AUROC = AP = 1 for [.8, .3, .6, .1] / [1, 0, 1, 0]), and a one-class set reports them unavailable with the reason, so MonitorSpec("auroc") never elects it. link_predictor_spec(...) (GCN / GraphSAGE / GAT encoder with no activation after its last layer, so dot logits can be negative; dot or MLP decoder) is rebuilt on reload, and the manifest is component state "link.task", so a resume with another split fails before any update. Node-level nets and the graph-pooling adapter refuse link batches. See examples/link_prediction_offline.py (#247).
• Graph-level classification (nnx.graph_tasks). GraphCollection(graphs, ids, targets=..., unlabeled=[...]) holds whole PyG graphs with stable integer ids and graph-level targets, refusing duplicate ids, zero-node graphs, edges outside a graph and an unflagged missing target, and batches them in order (or shuffled with a seed) without ever splitting itself. GraphClassifier(encoder, pool="mean" | "sum", head) pools node rows into one row per graph — graph-local and node-order invariant (an identity encoder on constant features 1/2/3 over 2-, 5- and 9-node graphs gives mean [1, 2, 3], sum [2, 10, 27]) — and validates each batch (graph ids, ptr / batch-vector consistency, cross-graph edges, one target per graph) before the forward pass. graph_classifier_spec(...) is a registered recipe (GCN / GraphSAGE / GAT encoder, pool, linear head) rebuilt on reload. With TaskSpec.categorical(n, ignore_index=IGNORE), loss and metrics are averaged over labeled graphs (NLLs [1, 2, 3] in batches of 2 then 1 give 2), so records, BEST selection and the run chart (now count-weighted for per-batch task records) use labeled-graph counts; predict_proba reports graph ids as sample ids through shuffling and concatenation. Node-level nets, NNGraphDataset over a multi-graph dataset and two_dim_tsne_checkpoint_logits refuse collection batches, and to_onnx / netron_export refuse a graph classifier before writing. See examples/graph_classification_offline.py (#246).
• Query-grouped ranking evaluation (nnx.ranking). RankingTask(k=(10,), max_relevance=1, relevance_threshold=1, weighting="query" | "pair", candidate_sets="sampled" | "exhaustive", max_buffered=1_000_000, group_size=None) checks (features, query_ids, candidate_ids, relevance[, mask]) batches (or mappings) before any update — mismatched lengths, mixed or duplicate candidate ids within a query, grades outside 0..max_relevance and bad masks are refused — with integer or string ids kept as host values, so candidate identity survives sorting, device moves and reloads. task.objective() is a FEAT-004 objective: the pairwise logistic loss over unequal-relevance pairs within a query (relevance [1, 0], scores [2, 0] → 0.126928), weighted per query or per pair; ties and pairless queries contribute nothing. task.eval_step() buffers the validation stream under max_buffered, joins each query's candidates whatever the batching or order, refuses a duplicate candidate, ranks by score with ties broken by candidate id and averages MRR@k, Recall@k and NDCG@k (gain 2^rel - 1, discount log2(rank + 1)) over queries; a query with no relevant candidate is excluded and counted, and k above a query's size is refused. The record (kind="ranking", count = scored queries) carries the metrics plus queries, excluded_queries, candidates, candidates_per_query and exhaustive_candidates, with no classification placeholder; task.metric_specs() registers ranking.mrr / ranking.recall / ranking.ndcg so MonitorSpec("ndcg_at_10") selects BEST, and the task configuration is checkpointed component state "ranking.task". A validation task record with no value (every target masked or every query excluded, or no error and no loss) no longer falls back to the training metrics for BEST selection or plateau scheduling. examples/ranking_offline.py maps FAISS positions to stable document IDs for an ID-aligned evaluation; export_to_faiss and the embeddings scope are unchanged (#252).
• A causal language-model task with token-weighted evaluation (nnx.lm_tasks). CausalLMTask(vocab_size, alignment="shift_inputs" | "pre_shifted", ignore_id=-100, pad_id=None, smoothing=0.0, tokenizer=None) declares how a batch becomes next-token targets — shifting token ids exactly once, or taking aligned (inputs, targets) windows unchanged — with an optional per-position loss mask (never an attention mask), and rejects bad shapes, dtypes, lengths and out-of-range non-ignored ids before any forward pass. The ignore id, the padding id and the loss mask select one set of valid positions for the objective's denominator, the NLL and the token accuracy. task.objective() is a FEAT-004 objective: the (smoothed) cross-entropy sum over valid positions with the valid count as its denominator, so gradient accumulation normalizes by the window's valid tokens; an all-masked microbatch contributes nothing and an all-masked window takes no optimizer or scaler step and clears its gradients (schedulers still advance once per epoch). task.eval_step() reports the epoch's unsmoothed NLL — total valid-token NLL over the valid-token count, independent of batching and padding — as the record's loss and metrics["nll"], with perplexity (exp(NLL), +inf on overflow) and token_accuracy; an all-masked epoch is unavailable — never elected BEST in place of a training loss, and skipped by a plateau scheduler (for any validation task record with no value: every target masked, or no error and no loss) — no classification placeholder is written, and MonitorSpec("nll") selects BEST on validation NLL. HuggingFace-style {"input_ids", "labels"} batches shift their labels with the ids. The task configuration (and an optional tokenizer identity) is checkpointed component state "lm.causal_task" with the valid tokens the run scored, and a resume with another configuration is refused before the first resumed update. examples/11_tinystories_lm.py now trains with the task (no placeholder metrics, no hand-written accumulation) and ships the bounded offline causal_lm_task_workflow() (inline corpus, pre_shifted windows, train → validation → reload → generate). See docs/lm.md §6 (#251).
• A reproducible, offline-first decision-provider benchmark (nnx.decisions.benchmark). Samples carry a sample id, a typed question (its digest is the schema identity), the provider input, the true label, the task family (held-out families report apart), a grouping unit and the perturbation that produced them (permute_options, redescribe, add_distractors, add_none_of_the_above, add_context, rewrite_input). collect(provider, samples, provider_id=..., budget=Budget(max_calls, max_samples)) is the live collector: explicit provider and budget, one attempt per batch, unsupported records with the provider's own reason, and batches the budget cut short marked partial_batch. Records keep the provider, model revision, prompt identity and execution metadata across a strict-JSONL round trip (write_records / read_records, nnx.decision-record/1). evaluate(samples, records, split=...) replays with no provider call and fits nothing: records join by sample id and question digest (never row position), duplicate / missing / extra / mismatched records become coverage counts, and each slice (in-family, held-out, family, perturbation) reports accuracy, macro-F1, NLL (exact by default, +inf at a zero true-class probability), Brier, ECE with reliability bins and, given an abstention policy applied as given, selective coverage and risk — each with its denominator, or unavailable with a reason. Resources stay null unless declared (measured numbers need hardware). bootstrap_interval resamples grouping units with a recorded seed and flags degenerate samples. JSON, CSV and text exports agree; compare_reports refuses reports of another split or metric identity. See docs/decisions.md §8 and examples/decision_benchmark_offline.py (#243).
• Multi-seed summaries and paired comparisons (nnx.comparison). A read-only utility for repeated runs: it loads no model, reads no checkpoint, elects no runs/best pointer and writes nothing into a run. Each Observation declares its Metric (name, direction, unit), evaluation split, checkpoint-selection rule, configuration, data and split identity, replicate key, completion status and the committed evaluation behind the value, with unknown facts left None. summarize pools only observations that agree on all of those and otherwise reports a stratified Summary naming the differing fields (Summary.pooled() refuses with IncompatibleObservations); each GroupSummary gives the count, mean and sample standard deviation (n - 1; unavailable for n = 1, both for n = 0) and keeps failed and non-finite attempts listed and counted. Duplicate attempts are rejected. compare(a, b, pairing=...) pairs two configurations by unique replicate keys whose meaning you declare, reports signed B - A deltas (never flipped for a metric to minimize) and lists unmatched replicates; an optional Bootstrap(seed=...) records its unit, seed and method, and its interval is labelled seed variability. ComparisonReport exports a machine-readable table() and a concise text() carrying every run and attempt id, status, selection rule, unmatched pair and signed delta, unchanged by input order, and saves / reloads as strict JSON (nnx.comparison/1) with its results re-derived and checked. observations_from_runs reads saved runs (run.yaml, idps.csv or a history journal's committed records, and the provenance files, once each); a run with a parent carries the parent's identity (every ancestor's, with their declared data and split identities), the checkpoint tag, the resume mode and, when it is part of the procedure, the recorded start epoch in its configuration identity, so it pools only with runs descended the same way from identically configured parents; only recorded facts decide which parent attempt a run started from, and an unknown lineage (no record and no known order — say, a deleted parent with no recorded checkpoint) never pools; cycles and over-deep lineages of followed parents are refused in any listing order. See docs/concepts.md §23 and examples/compare_seeds.py (#250).
• Deferred decision jobs: batch independent questions, chain dependent ones (nnx.decisions). DecisionJob.ask(question, id=...), DecisionJob.collect({key: job}), job.map(fn) and job.then(fn) build an immutable job that calls no provider, runs no callback and draws no RNG; run(provider, state=..., limits=Limits(...)) and await arun(...) are the only effect boundaries, and they borrow the provider without closing it. Independent questions over one state share calls — ceil(count / cap) in collection order, keeping question ids and the collect key order — where cap is the provider's max_questions for a decide_many provider (1 for a provider with only decide, such as FixedHeadProvider), lowered by Limits.max_questions. Duplicate ids, unsupported questions and a max_tokens cap the provider cannot count are refused with InvalidJob before any call. A then continuation runs once, after its prerequisite succeeded, and returns Follow(job, state=...): the answer mapped explicitly into follow-up state; max_depth and max_requests bound it (JobLimitExceeded). Runs are fail-fast: JobFailed carries the completed outcomes and the failed and skipped question ids, nothing completed is re-run and the job never retries. arun runs up to max_concurrency calls at once; a cancel event stops scheduling and returns a "cancelled" JobResult, task cancellation cleans up only the job's own tasks, timeout raises JobTimeout, and requests already sent are reported as sent, never rolled back. Outcomes distinguish answered (rows answered / abstained under an abstention policy), failed, skipped and cancelled. A job of ask / collect pickles as plain data; one holding map / then refuses. FixedHeadProvider.check(question, inputs) exposes its pre-call checks. LogitsChain is unchanged. See docs/decisions.md §7 and examples/decision_jobs.py (#245).
• A local label-conditioned baseline: nnx.decisions.NLIProvider. It scores texts against candidate descriptions supplied at inference — none exists when the provider is built, and no output head is resized — with a natural-language-inference model and tokenizer the caller supplies. NNx imports no NLI library, downloads nothing and never calls from_pretrained. A Choice softmaxes the entailment logits across its candidates (choice_scoring="entailment_softmax"; logits log(3) / 0 → (0.75, 0.25)), with results keyed by the request's option ids in its order. A Boolean softmaxes each text's own contradiction and entailment logits (boolean_scoring="entailment_vs_contradiction"; 0 / log(4) → 0.8), never normalizing texts together; Score is refused before any model call. The entailment / contradiction ids (ints, or label names resolved through model.config.label2id), the hypothesis and Boolean templates, the scoring methods and the truncation policy are explicit and validated. Pairs go to the model's device in chunks of pair_batch_size, in eval mode under no-grad, with every submodule's training flag restored on success and failure. Over-long pairs are measured first and either cut (truncation="only_first", reported in raw["truncated"]) or refused before any model call ("error"). Each result's raw records the pair logits and provider.record() (templates, label ids, scoring, truncation, model revision, calibrated: False): the scores are zero-shot, not calibrated, and their transfer to a task stays empirical. embed_texts and export_to_faiss keep their signatures. See examples/decision_nli.py, which reports accuracy, macro-F1 and categorical NLL / Brier on supplied records with the split, revision and settings (#234).
• Diffusion and I-JEPA objective adapters. diffusion_objective(schedule, *, seed=None, noise_fn=None, nonfinite="fail") and jepa_objective(target_encoder, predictor, mask_fn, *, ema_momentum=0.996, nonfinite="fail") (also DiffusionObjective / JEPAObjective) describe DDPM noise prediction and I-JEPA latent prediction as objectives over the shared update engine, so both gain gradient accumulation (exact for uneven microbatches), mixed precision, clipping and the engine's non-finite policy; diffusion_train_step_factory and jepa_train_step_factory are unchanged. Each returns one normalized term ("noise_mse" / "latent_mse") — the squared-error sum over the valid elements over their count — and detached metrics with no classification fields, and never zeroes, back-propagates, scales or steps, so unequal target masks combine as (sum_1 + sum_2) / (count_1 + count_2). Records, update events, train monitors and BEST see that denominator-weighted loss. The JEPA EMA target advances once per committed update, from the updated online weights, never per microbatch or after a skipped window, through a new Objective.after_update(events) hook; Objective.check_run(model, *, optimizers, callbacks) refuses, before any run is reserved, a predictor that is not optimizer-owned exactly once, a target that is part of model.net or optimizer-owned, a target without matching online parameters and topology-changing callbacks; and the imperative diffusion / JEPA steps passed as objective= (or an Objective passed as the step function) are refused. Both adapters are checkpointable components — the diffusion objective's own generator (restarted for each run, so an equally seeded rerun draws the same noise; a sample preview with its own generator never perturbs it) and spec; the JEPA spec, predictor reference, EMA target and counter — so two epochs equal one plus a stateful resume, and a duplicate predictor payload, another predictor reference or another spec is rejected before anything is restored. See examples/08_diffusion_2d_mixture.py::objective_mode and examples/16_ijepa_image_plumbing.py::objective_mode (#257).
• Explicit FP32 / FP16 / BF16 precision policy (nnx.precision). NNModelParams(precision=PrecisionPolicy("fp32" | "fp16" | "bf16", fallback="error" | "fp32")) replaces the Boolean-only mixed_precision (CUDA FP16 or nothing): BF16 autocasts without a scaler on CPU and on CUDA devices with native bf16 support (compute capability 8.0+, not emulation), FP16 keeps the GradScaler on CUDA, and FP32 is the default. The policy resolves into a ResolvedPrecision against the model's device when the model is built and again before any run is reserved; a request the device cannot run raises PrecisionUnsupportedError there (no runs/ artifact) unless the policy names fallback="fp32", which trains in full precision and records why. Autocast wraps only the forward pass — the backward runs outside it, parameters stay float32 (never model.half()) — and the update keeps its unscale → normalize → finite-check → clip → step order, scaler-free BF16 refusing a non-finite gradient. NNModel.train (the default step and objectives) and Trainer.train(objective=...) resolve one policy before callbacks and iteration and reuse it in evaluate() and predict() / predict_proba(), whose outputs stay float32 (BF16 converts to a NumPy dtype). The built-in imperative paradigm steps (finalize_step) and Trainer step functions run in full precision and refuse a reduced policy before any work is done (as does born_again_train with more than one generation, before generation 0 trains); a custom train_step_fn reads TrainStepContext.precision (a run with one records training as applied by the step, outside what the policy covers; an FP16 context without a GradScaler is refused rather than run unscaled), and an objective ObjectiveContext.precision (with ctx.full_precision(output) for float32 outputs before task preparation and records). kd_objective runs its frozen teacher in the teacher's own inference precision (its policy, else full precision), never in the student's — under the legacy flag on CUDA the teacher's forward therefore now runs in full precision rather than sharing the student's float16 autocast. ExperimentPlan.validate() reports an unsupported policy (model.precision) and a full-precision-only step under a reduced one (train_step_fn), and Trainer.train resolves the policy before seeding or building optimizers. Under FP16 an update window is all-or-nothing across a Trainer's named optimizers (none steps when any one's gradients overflow, and an optimizer without gradients that window is left out of the scaler), and named optimizers sharing a parameter are refused before the run. An explicit PrecisionPolicy("fp32") normalizes to no policy, keeping the default's run id. The policy and the FP16 scaler state round-trip checkpoints; a stateful resume into a different effective precision fails before the model is touched, while a resume_mode="weights_only" warm start may switch. from_checkpoint and Hub loading re-resolve on the destination device, never from saved metadata, and accept precision= to replace the policy. run.precision (also read back by NNRun.load from metadata.yaml) records the requested and effective precision, the fallback reason, TF32 (reported separately, never set) and which NNx surfaces apply it (training where NNx runs the default step or an objective, evaluate() and predict(); custom step code applies ctx.precision itself) — not nnx.lr_finder, nnx.diffusion.sampling or generation. nnx.precision_support() reports each device's cells as verified, supported, unverified (hardware absent) or unsupported; CPU BF16 is verified by seeded classification and KD fixtures with accumulation within the predeclared FP32-reference tolerances (REFERENCE_TOLERANCES: FP16 5e-3 / 5e-4, BF16 5e-2 / 5e-3 for losses / weights). The legacy mixed_precision=True keeps its meaning (FP16 in NNModel.train on CUDA, plain FP32 elsewhere, no effect on evaluation, prediction or Trainer) and its run id; a policy contradicting it is refused, one agreeing with it (FP16) normalizes the flag away and keeps the policy's run id, and the precision field is omitted from state() when unset. AMP is decided by the policy, not by NNModel._build_grad_scaler: an FP16 run whose hook returns no scaler (or a disabled one), or a CUDA run that is not FP16 whose hook returns one, is refused before training. NNModelParams pickles written before task / precision existed still load. See examples/01_synthetic_classification.py::precision_workflow (#248).
• Explicit scheduler clock units (NNSchedulerParams.clock). clock="epoch" (the default, omitted from state() so existing configurations keep their run id) steps a scheduler once per completed epoch, as before; clock="optimizer_update" (builder: .clock("optimizer_update")) steps it once per committed update of its optimizer — never per microbatch, an all-masked window or a skipped AMP step — so its horizons count updates (two epochs of five microbatches with accumulate_grad_batches=2 give six steps). Updates are reported explicitly: by the update engine for objectives, by default_train_step after each optimizer step it takes, by finalize_step (so every built-in paradigm step), and by custom steps through the new TrainStepContext.report_update() / TrainerStepContext.report_update(name) (an epoch with no reported update warns); whoever calls optimizer.step() reports it (default_train_step and finalize_step report their own steps, never one the AMP scaler skipped), and every report counts. Where NNx owns the windows and the loader has a length, the default horizon is the planned updates; otherwise total_steps / T_max / step_size must be set, and a one-cycle or warmup-decay schedule stops the run with an error at the first update past total_steps (a cosine schedule that passes a T_max the run did not configure — a defaulted or checkpoint-restored one — warns). One-cycle and warm-up traces match native per-update stepping, accumulation and a short final window included. In Trainer each optimizer keeps its own count and schedule, and auto_step_schedulers=False detaches update clocks too (the step function owns every scheduler step), including after a resume. Each update clock is a required checkpointed component (nnx.scheduler_clock[.<name>]: clock, owner, budget, and the schedule's position), so a stateful resume continues the schedule — reproducing the remaining LR sequence — and refuses a mismatched clock or budget (or a checkpoint from before component state, whose scheduler position counts epochs) before anything steps. A plateau scheduler refuses the update clock. LRMonitor.history keeps one entry per epoch; the new LRMonitor.update_history records (update, LR) per update. See examples/scheduler_clocks.py and the auto_step_schedulers=False scenario in examples/09_gan_with_trainer.py (#236).
• Replayable model transformation recipes (nnx.transforms). TransformRecipe([lora("layers.0", r=4, alpha=8.0), low_rank("layers.1", rank=4)], materialization="in_place" | "fresh") records each operation's id, version, explicit nn.Linear targets and config in order (copied from the caller's list, read-only, hashable and picklable; integer LoRA alpha / dropout are stored as floats, so alpha=4 and alpha=4.0 share a run id). validate(model, optimizers=...) checks exactly what materialize would, mutates nothing and raises one RecipeError naming, by operation index and target, a missing target, a non-nn.Linear or aliased target, overlapping targets, a target inside another operation's subtree, an out-of-range rank, an unknown operation or version, an in-place model that unrecorded surgery already changed, and an optimizer already holding a target's parameters (build the optimizer after the recipe; check_optimizer refuses one holding replaced parameters, the source of a fresh materialization or a frozen LoRA base without its adapter, and accepts one built afterwards over the trainable parameters, a subset, parameters outside the net, or unfrozen base tensors BitFit-style). materialize changes the model in place (the default, keeping its weights) or transforms a fresh, randomly initialized registered base, transactionally (a failure leaves the model and the global random streams as they were), and records the operations; a runtime-only module= is refused. Every run checkpoint tag, ModelCheckpoint snapshots, both NNCheckpoint formats, save_pretrained and run bundles (nnx.bundles) carry them — export_bundle refuses any other recorded transform, such as a callback's own train-end operation, naming its index and id — and from_checkpoint / from_pretrained replay them on a fresh base before loading tensors — low-rank factors are allocated and loaded, never re-factorized — matching keys, trainability and predictions. A recipe model trains through NNModel.train and Trainer.train (any step function), its live topology checked against base-plus-recipe by name, shape and target configuration first; the recipe joins the run id and ExperimentManifest.for_model (runs and manifests without one are unchanged), the final LAST stores the weights once, and a resume from any tag needs the same recipe materialized. An unknown or malformed recorded transform names its index when replay fails; transformed weights that record no recipe (a raw state dict) get an error saying so from from_checkpoint and a strict from_pretrained (strict=False keeps its partial load). See examples/20_low_rank_surgery_ffn.py::recipe_reconstruction_workflow and Surgery §8 (#238).
• Bounded training history with an append journal (nnx.history). NNModel.train(..., history=HistoryJournal(retention, chunk_size)) and Trainer.train(..., history=...) keep only the last retention records in memory (ctx.idps — a live, read-only window — and the returned NNRun.idps) and append every record once to runs/<id>/history/ — immutable JSONL chunks, a hash-chained index, per-epoch summary rows and a small manifest published atomically before LAST — instead of rewriting idps.csv each epoch, so memory stays bounded and bytes written grow linearly with the records. LAST stays the epoch commit marker: readers never show records past its epoch (an uncommitted tail left by a crash is ignored), a failed LAST republishes the previous manifest, and a corrupt, missing or altered committed chunk raises HistoryCorruptionError. NNRun.load reads only the committed tail; the notebook chart reads the per-epoch rows; BEST, the runs/best election and a resume read no records. iter_history(run_id, lineage=False) streams the committed history, export_history_csv(run_id, path, lineage=False) writes the legacy idps.csv layout, and migrate_history(run_id, spec=..., discard_uncommitted=False) explicitly moves a legacy CSV into a journal keeping the run id (read back before the CSV is removed; records past LAST are never dropped silently). A resumed run owns its own files and leaves its source byte-for-byte unchanged; lineage=True follows each run's recorded resume and yields the parent's committed records up to the checkpoint the child resumed from, then the child's, each epoch once (a run that started fresh, such as a born-again generation, keeps its own epochs). chunk_size defaults to min(250, retention). LRMonitor.history keeps the LRs of the last retention epochs (LRMonitor(bounded=False) keeps all); a Callback needing every record declares history_access = "full" (materialised at on_epoch_end and on_train_end), and a plain function callback is refused before training. A journal run is not saved under another runs root (its history would not follow). Opt-in: without history=, runs, run ids and idps.csv are unchanged, and legacy and journal runs load side by side (NNRun.all). NNRun gains a runtime-only history field (the journal directory, never in state() or the run id), and ResumeStatus gains source_epoch — the epoch of the checkpoint a session resumed from, written to metadata.yaml only when known. See examples/history_journal.py and Concepts §4.5.
• Portable, data-only run bundles (nnx.bundles). export_bundle(run_id, destination, checkpoint="last", calibrators=...) publishes one of your run's checkpoints as a directory of data: bundle.json (format, version, generation id, every payload's SHA-256 and size) and a g-<generation>/ directory with state.json (schema-validated JSON with typed $tuple, integer-keyed $intdict, non-finite $float and $tensor encodings), model.safetensors, training.safetensors (optimizer, scheduler, GradScaler, RNG and component state — a "resume" bundle; weights-only checkpoints such as ModelCheckpoint snapshots give an "inference" bundle) and TemperatureCalibrator records, whose labels and weight fingerprints are checked against the model. NNx writes and reads the safetensors format itself, so bundles need no optional dependency. inspect_bundle and validate_bundle never unpickle, call a factory or download; validation rejects missing, unlisted, swapped, corrupted, symlinked and out-of-root payloads, duplicate JSON keys, generation mismatches and unknown versions before any tensor is read. reconstruct_bundle(path, factories=..., components=...) rebuilds the model from caller-supplied factories (a ModelSpec) and names every missing factory or component before any model is allocated; the rebuilt model predicts bit-for-bit like the source, and ReconstructedBundle.resume(params) continues training from the bundle's state exactly like the original run (an "inference" bundle refuses before anything changes). A new generation is published by an atomic manifest swap, so an interrupted export leaves the previous bundle usable. Module extra state that is not a tensor, custom objects in training state and runtime-only modules are refused with no pickle fallback. NNCheckpoint.from_file and NNModel.from_pretrained now name a bundle (or bundle payload) instead of failing obscurely. See docs/hub.md §3, concepts §22, SECURITY.md and examples/run_bundle.py (#237).
• Streaming prediction and mergeable metrics (nnx.streaming). NNModel.iter_predict(loader, spec=None, *, rich=False) returns a context-managed PredictionStream that yields one PredictionBatch(logits, classes, sample_ids) per loader batch — or, with a ProbabilitySpec or rich=True for a model with a task, a PredictionResult — in loader order. Concatenated (concatenate_predictions), the batches are exactly predict() and predict_proba(), including sample ids, graph seed-row slicing and categorical / multilabel / continuous decoding. Each batch runs in eval mode under no_grad with every submodule's training mode restored before it is yielded or its error raised; closing the stream drops its references to the loader's iterator and the model, a closed stream ends iteration like a closed generator and cannot be iterated again (StreamClosedError), and the loader stays the caller's. The stream holds only the batch in flight. StreamingMetrics(metrics, semantics, labels=..., threshold=...) (or StreamingMetrics.for_task) accumulates declared MetricSpecs with update / update_logits, merge and finalize: built-ins keep sums, counts and confusion counts, so memory does not grow with the samples and a merge is order-independent; a merge refuses different metric declarations (id, version, config, name), probability semantics, task labels or thresholds (MetricMergeError); finalize() returns a read-only, repeatable MetricSnapshot (empty and fully masked streams: count=0, every metric unavailable). A metric that needs every stored score (a rank metric such as AUROC) is refused in bounded mode unless materialize=True asks for O(N) storage. streaming_eval_step is an opt-in eval_step_fn that builds the default validation record from counts and sums — the same task counts, weighted loss, metrics, monitor values and BEST choice — and refuses extra_metrics or a stored-score metric before any run is reserved (ExperimentPlan.validate() reports both). The built-in f1 metric now accumulates confusion counts — the same values, with memory bounded by the number of classes. EvalStepContext gains metrics (the run's declared metrics, default ()). predict(), predict_proba() and evaluate() are unchanged and eager. See examples/prediction_stream.py (#242).
• Immutable fluent experiment plans (nnx.plans). ExperimentPlan holds NNParams, NNModelParams, NNTrainParams, data, seed and callbacks; every with_* method (with_epochs, with_seed, with_data, with_callbacks, with_callback_factories, resuming, ...) returns a new plan and leaves its source and siblings unchanged, nested collections included. validate() aggregates field-path Diagnostics (PlanValidation, PlanError) while consuming no loader, calling no factory, building no model, reading no weights and writing no directory, and reports a resume that conflicts with parent lineage statically. probe(example_batch) seeds, builds a temporary model and runs a no-grad forward pass (and the training loss when the batch has a target), then restores the ambient RNG streams, cuDNN / deterministic-algorithm settings and seeding environment variables, even on failure. fit() returns FitResult(model, run, metrics, attempt_id): it compiles to the existing NNModel(...).train(...) loop (seed, then the loader and callback factories once, then the model), so it matches an equally seeded imperative run on weights, history and callback order. Each fit is a distinct attempt whose id is folded into the run's salt, so repeated fits get distinct run directories and a plan never sets overwrite_existing; resuming(run_id) records parent lineage and leaves the parent's artifacts intact. FitResult.metrics names each split and its availability from the run's final epoch, with no second loader pass, and a train-only plan invents no validation. Loaders, callback instances, metric callables, step functions and components are borrowed and shared by every fit, while factories give fresh ones per fit. NNModel.train, set_seed and every params class are unchanged. See examples/experiment_plan.py (#221).
• Abstention policies and risk-coverage evaluation (nnx.abstention). AbstentionPolicy(kind="max_probability" | "margin", threshold=..., labels=..., model_id=..., tuning_split_id=...) accepts a row when its top probability (or top-two margin; a tied top two gives 0) is at or above the threshold and abstains otherwise, with reason below_probability_threshold / small_margin. Ties break by label order. Accepted and abstained outcomes keep the prediction, distribution, score and sample id; malformed probabilities raise AbstentionError rather than abstaining, and inputs are never modified. A policy records its probability field (probabilities or calibrated_probabilities with the calibrator id), threshold, label order, model id and tuning split, round-trips as JSON (nnx.abstention/1), and refuses input of another schema with AbstentionSchemaError before reading it. apply takes a CalibratedPrediction, a categorical PredictionResult or arrays. decide(result, policy, model_id=...) keeps a typed decision a success when it abstains, preserving its distribution and raw-versus-calibrated identity. CoverageReport gives coverage = accepted/total and risk = incorrect/accepted (risk unavailable with status no_accepted when nothing is accepted, never 0), keeping every original sample id; accepted_rows feeds hard-label reports only accepted rows; CoverageAccumulator aggregates chunks to exactly the eager report. select_threshold runs on a named validation split (never the test split), reports every candidate threshold with its counts, returns no_feasible_policy when nothing meets the empirical risk ceiling, and risk_coverage_curve ends in a curve-only accept-none endpoint that no policy can hold or load (a threshold of 1.0 still accepts a score of 1). Validation of probabilities, targets and sample ids, and JSON artifact handling, are now shared with nnx.calibration, and TemperatureCalibrator.transform(..., sample_ids=...) can name an array's rows (so calibrated chunks carry their own ids). See examples/abstention_offline.py (#218).
• Fitted classifier calibration (nnx.calibration). fit_temperature(logits, targets, labels=..., model_id=..., split_id=..., train_split_id=..., test_split_id=...) fits one scalar temperature on a calibration split by minimizing its NLL in float64 over a copy of the categorical (N, C) logits (arrays, tensors or a categorical PredictionResult). It rejects non-finite logits, C < 2, out-of-range targets and a calibration split id equal to the training or test id, and records that ids and loaders cannot prove the rows are disjoint (disjointness: "unverified"). A fit with no interior optimum (a separable split, constant logits, logits no better than uniform) or no convergence returns a failed CalibrationFit with the reason, never a NaN calibrator. The TemperatureCalibrator stores the temperature, the ordered labels, the model id, the fit config and the split ids, and serializes as JSON (nnx.calibration/1, to_json / save / load) that reloads with identical probabilities. transform refuses logits with a reordered label list or another model id before reading them (CalibrationMismatchError), unless a named override records the mismatch. It returns a CalibratedPrediction that keeps the raw logits, raw and calibrated probabilities, decoded argmax, labels, sample ids and calibrator id as separate fields. report(..., split_id="test") shows held-out NLL, Brier, ECE and reliability bins before and after calibration, and its outcome is "improved" only when neither NLL nor Brier got worse; the report also records any override. negative_log_likelihood (floor 1e-12 like the named nll metric; epsilon=None is exact and gives +inf at a zero true-class probability), brier_score (mean class-summed squared error) and reliability_bins (equal-width [lo, hi), the last closed at 1, empty bins with count 0) match the named nll / brier metrics. model_fingerprint(model) hashes the weights for use as a model id. The model, predict() and the run are never modified; the fitted temperature is unrelated to the sampling temperature of generate. The module imports none of its consumers, and it shares the NLL / Brier terms with the named metrics and the softmax with predict_proba. See examples/calibration_offline.py (#217).
• Train-fitted preprocessing and split transforms (nnx.preprocessing). Standardizer.fit(source, rows=..., columns=...) fits per-column (x - mean) / scale on an explicit training membership and reads no other row (DataFrames, tensors, arrays or any row-indexable source). The population statistics are float64, and a constant column gets scale 1. The fitted standardizer is frozen: transform reuses it for validation, test and inference under a declared schema. Named columns must be present and in the fitted order (extra columns are ignored), arrays must have the fitted width, and values must be numeric and finite; anything else raises PreprocessingError before a model sees the data. Statistics, schema, output dtype and the fit-membership identity serialize as primitive JSON (to_json / save / load, format nnx.preprocessing/1) and reload float64-exact; loading rejects bad lengths, non-finite values and non-positive scales. NNTabularDataset(standardize=True) fits on its training split only (random or split= manifest, same membership) and exposes ds.standardizer; standardize=<fitted Standardizer> reuses one without refitting (its columns and dtype are checked before the split). A constant column keeps its exact value as the mean, and the fit copies only the training rows of the feature columns. NNModel.predict takes model-ready features, so apply the reloaded standardizer once. SplitView and NNDataset(train_transform=..., eval_transform=...) give training and evaluation independent input transforms over an untouched base dataset: targets, the base transform, labels and order never change, and views are safe to read alternately and from DataLoader workers. Runtime callables are recorded by name as not reconstructible. ExperimentManifest.for_model(..., preprocessing=standardizer) records the fitted schema and fit membership. Defaults are unchanged, and dataset state() gains a key only when these are used. See examples/preprocessing_offline.py (#240).
• Reproducible group, time and stratified splits (nnx.data_splits). plan_split(ids, strategy=...) returns a SplitManifest of disjoint train / validation / test memberships by sample id, with the strategy, its parameters, the seed and the source identity. strategy="group" keeps every group in one split; "chronological" cuts at cutoffs=(validation_start, test_start) (a timestamp on a cutoff starts the later split, so ties never separate) and excludes the rows within gap before each cutoff; "stratified" keeps every class in every active split and, for a class with too few rows, raises (default) or keeps it in training (insufficient="train"). Plans depend only on ids and their groups / times / labels — never on row order or any global RNG (ids are ordered by SHA-256 of the seed; seed=None draws and records one). A manifest is plain JSON (to_json / save / load, format nnx.split/1) and digest() is its split identity; ExperimentManifest(splits={...}) accepts a manifest and records that digest. SplitManifest.resolve(ids, source=...) maps a plan back onto rows and raises SplitError for duplicate, missing or unexpected ids, ids of another type, a changed or uncheckable source identity, or a positional-only plan (plan_split(ids=None) requires source=). NNTabularDataset(split=manifest, id_col=...) builds its loaders from a manifest instead of random_split (id_col is required; only the planned rows are admitted and counted; empty validation / test → None loaders, empty train → error before any loader); without split= nothing changes. See examples/split_replay.py (#239).
• Versioned experiment provenance (nnx.provenance). NNModel.train(..., provenance=manifest) and Trainer.train(..., provenance=manifest) record an ExperimentManifest — the declared task, label order, model descriptor, data and split identities, objective version and config (ExperimentManifest.for_model(model, train=..., data=..., splits=..., objective=...) collects them without building a model or reading a loader). It is stored in runs/<id>/provenance.json with its fingerprint: SHA-256 of canonical UTF-8 JSON (sorted keys, ordered arrays, format version nnx.provenance/1 inside the hashed document). A fresh attempt is recorded in attempt.json: attempt id, running → completed / failed / cancelled, last committed checkpoint (tag, epoch, generation), error, and on resume the parent attempt and checkpoint generation. Serialization rejects non-finite floats, sets, bytes and callables by path. Identities are declared (never reported as verified), digest (from the explicit hash_file / hash_bytes) or unknown. compare(a, b) reads only provenance files and reports stable field paths as equal / different / missing / unknown / declared / verified, with absent provenance never equal. Both files are written atomically and never rewritten by NNRun.save; NNRun.load exposes them as run.provenance. A failed status write never masks the training error, and run ids are unchanged. See examples/02_resume_training.py::provenance_mode (#241).
• Provider-neutral typed decisions (nnx.decisions). A new inference branch asks typed questions — Choice (2+ Option(id, description)s with unique ids), Boolean and Score (2+ ordered levels) — and returns validated results: ChoiceResult / ScoreResult distributions in the question's order (summing to 1 within 1e-6) and BooleanResult.p_true, with the provider's own output kept in raw. Option ids are bookkeeping and never reach model_view(); digest() changes when options are reordered or reworded and binds each result to its question. validate_response reorders keyed provider output and rejects missing, duplicate, unknown or unlabeled ids, never renormalizing a malformed distribution. ScoreResult.expected_index (sum(i * p_i)) is documented as ordinal, with a vendor's own score kept in vendor_score. Providers implement the DecisionProvider protocol and declare Capabilities (primitives, modalities, dynamic labels, batch limits, inference / training / export), checked before any model call. FixedHeadProvider serves a trained classifier's head over its exact label space or a bijection onto it (Choice, ordinal Score, or, from one logit, the one Boolean the head was trained for — question=Boolean(...), any other Boolean is unsupported), reorders columns, rejects unseen labels before its model-call counter advances and restores model modes. The typed errors UnsupportedCapability, InvalidDecisionRequest, InvalidDecisionResponse and ProviderFailure are public. Importing nnx.decisions starts no backend and needs no hosted-SDK extra. New guide docs/decisions.md (docs numbering shifts by one from API reference on); see examples/decision_fixed_head.py (#219).
• Arbitrary modules and registered model factories (nnx.models). NNModelParams.net now also takes a ModelSpec(id, version, config, seed) naming a factory registered with register_model_factory(id, version, factory) (factory(config) -> nn.Module; plus unregister_model_factory / registered_model_factories / resolve_model_factory), and NNModel(module=m, params=NNModelParams(...)) wraps a caller-owned module as is (model.net is m, never cloned or re-initialized), describing it with a RuntimeModule descriptor (class name plus a fingerprint of the parameter layout and the module's repr, reconstructible=False). Registered modules are built with Python, NumPy and torch seeded from seed (the ambient RNG states restored) and rebuild through the registry from checkpoints (pickle and safetensors), from_pretrained and the checkpoint t-SNE; an unregistered factory (MissingModelFactoryError) or a changed topology fails before any weight loads. Runtime modules train, checkpoint and resume normally; their checkpoints report NNCheckpoint.reconstructible == False, reload with NNModel.from_checkpoint(ckpt, module=...), and portable saves (safetensors checkpoints, save_pretrained) fail before anything is written. Non-built-in modules see batches through a runtime-only batch_adapter — PositionalInputs(n_inputs), KeywordInputs(inputs, target), or the module's own unpack_batch — and predict() accepts a mapping of keyword inputs; the Mixup / CutMix, KD, feature-KD and MoE steps and kd_objective split batches through the same adapter; nnx.viz.summary(model, batch=...) takes the same adapters and keeps every submodule's training mode. run.yaml omits net for these descriptors and str(run) / the notebook table show the descriptor instead of built-in dimensions. Built-in nets, their serialization and run ids are unchanged. The dev extra adds onnxruntime for the ONNX parity test. See examples/custom_module.py and examples/04_onnx_export.py::registered_module_variant (#216).
• Objectives over one shared update engine (nnx.objectives). NNModel.train(objective=...) and Trainer.train(objective=...) accept an objective in place of a step function. An objective returns LossTerms per microbatch: a differentiable numerator plus an explicit denominator (reduction="mean") or none ("sum"); mixing the two within a window fails. The shared engine back-propagates each term into its own buffer and, per update window (accumulate_grad_batches, cut short at the epoch's end), divides every normalized term by its window-total denominator, so uneven microbatches, short windows and masks give exactly the full-batch update. It then runs autocast → unscale → clip → step and applies a declared non-finite policy (nonfinite="fail" raises before stepping, "skip" drops the window). It fires Callback.on_optimizer_update(ctx, event) once per committed optimizer update, never per microbatch, masked window or skipped step. Records carry a distinct update_count, also on ctx.update_count. supervised_objective() and kd_objective(teacher, alpha, temperature) ship; the latter gives distillation gradient accumulation, AMP and clipping. One-microbatch windows take a single backward pass, zero-weight terms are never back-propagated, and parameters frozen or unfrozen between epochs train as with the default step. The update counters are checkpointed component state (nnx.update_engine) and continue across a stateful resume; an objective run's whole-epoch training loss uses the terms' own denominators. In Trainer the engine steps every named optimizer once per committed update, each with its own grad_clip_norm and one event per optimizer; Trainer has no mixed-precision setting, so objectives run there in full precision (with a warning when the model asks for AMP on CUDA). Passing both a step function and an objective fails before anything runs. Imperative factories and finalize_step are unchanged. See examples/10_knowledge_distillation.py::objective_mode (#214).
• Named metric and monitor contracts (nnx.monitors). NNTrainParams(metrics=[MetricSpec(...)], monitor=MonitorSpec(...)) (also on NNTrainerParams and its builder) is opt-in and serialized only when set, so existing run ids are unchanged. A MetricSpec(id, version, config, name) names a registered metric: built-in accuracy / f1 (decoded labels), nll / brier (probabilities: softmax for categorical, element-wise sigmoid for multilabel) and mae / mse (continuous outputs), or custom ones via register_metric(id, version, factory, input=, mode=). Values are accumulated over the full sample, never as batch means, and reported in the records' metrics, from evaluate(loader, metrics=...) and from the default training step. Specs round-trip through run.yaml without code; an unregistered id or an input the model cannot provide fails before any run is reserved or metric code runs. A MonitorSpec(metric, split, mode, min_delta, on_missing, on_nonfinite) is the one improvement rule for BEST selection, a monitor-aligned ReduceLROnPlateau (its direction and an absolute min_delta threshold) and EarlyStopping(monitor=MonitorSpec(...)): for NLL [0.8, 0.7, 0.695, 0.71] with min_delta=0.01, all three see exactly two improvements. Runs with metrics or a monitor record a whole-epoch train_summary (loss and error with full-epoch denominators, declared metrics over the full epoch) and the epoch's selection (MonitorRecord: identity, direction, configuration, value, status, improved) on each epoch's last idp. Both persist through idps.csv and pickle / safetensors checkpoint metadata, and are shown by the progress bar, TensorBoard / W&B (train_epoch/*, monitor/*), injected callbacks and the run chart. Named-monitor runs are never elected into runs/best. See examples/03_custom_metrics.py (#213).
• Checkpointable component state and Trainer warm resume (nnx.components). Anything that carries state across epochs registers as a named, versioned component (ComponentSpec(name, version, required) plus component_state() / load_component_state(state, *, version), the StatefulComponent protocol). NNModel.train and Trainer.train discover components among callbacks and the step function (or take them through components=[...]), write their state into the same training-state sidecar generation as the checkpoint (format v4), and on resume validate every saved component before mutating anything — a missing, unknown or incompatible required component yields one ComponentRestoreError listing every problem — then restore them transactionally after the callbacks' reset hooks, rolling back every restored component, the model and the RNG streams on failure. A component may add a check_component_state hook whose problems join that report. EarlyStopping (early_stopping, optional; several default-named instances are numbered early_stopping.2, … in callback order, and a changed monitor or mode is reported before restoring) resumes its best value, patience and selected monitor, so a split run stops at the uninterrupted epoch; jepa_train_step_factory now returns JEPATrainStep, whose required jepa.target_encoder component persists the EMA target encoder. Trainer checkpoints now carry every named optimizer's and scheduler's state (including the final post-callback LAST) and resume validates each optimizer like NNModel.train does (type, registered factory, parameter topology, scheduler type and one-cycle horizon), and NNTrainerParams gains resume_from_run_id / resume_from_checkpoint / parent_run_id / resume_mode (NNTrainerParamsBuilder.resume_from(run_id, checkpoint, mode)), serialized only as parent lineage. resume_mode ("auto" default, "stateful", "weights_only"; also on NNTrainParams) chooses what is restored — a weights-only warm start needs no shared one-cycle horizon; run.resume_status (ResumeStatus, stored in metadata.yaml) reports it. ModelCheckpoint files declare themselves weights-only; resume_from_checkpoint accepts their <tag>_e<epoch> stem, and a stateful resume from one fails before mutating anything. Sidecars from earlier versions still resume, with components started fresh and a warning; default serialization and run ids are unchanged (#215).
• Task adapters (nnx.tasks, NNModelParams(task=...)). An opt-in TaskSpec — TaskSpec.categorical(n, ignore_index=...), TaskSpec.multilabel(n, threshold=...) or TaskSpec.regression(n) — makes the default training step, NNModel.evaluate(), predict() and predict_proba() task-aware. The adapter validates each batch's output/target shape, dtype and range before any backward pass or optimizer update (TaskValidationError), masks ignore_index / NaN targets out of the loss and every metric alike (user extra_metrics see only the valid targets; masking needs an elementwise loss, evaluated at full shape so pos_weight still applies), averages the loss over the valid targets with the default step's accumulation denominators (uneven batches and accumulated windows reduce exactly like one full batch; an all-masked window takes no update), and computes whole-loader metrics. Regression records carry mse / mae and leave accuracy / f1 / recall / precision and error as None; multilabel records report subset accuracy plus subset_accuracy / element_accuracy; categorical mode matches the legacy loss and metrics. predict() returns continuous values for regression and threshold-decoded indicators for multilabel, and predict_proba(X) may omit its spec (a regression result has probabilities=None, spec=None). The model's loss, output width and net type are checked at construction, and NNModel.train / Trainer.train re-check the live loss_fn before touching a loader. NNEvaluationDataPoint gains optional kind / count / status / metrics fields (an all-masked record has status="empty"), which round-trip through idps.csv and checkpoint metadata; TensorBoard / W&B emit the task metrics. The spec persists as a versioned task: mapping in NNModelParams.state() — omitted without a task, so legacy serialization and run ids are unchanged. See examples/regression_task.py.
• Branching params builders (copy() / from_params()). NNOptimParamsBuilder, NNSchedulerParamsBuilder, NNTransformerParamsBuilder and NNTrainerParamsBuilder gain copy(), which branches a complete or partial builder so setters on either never affect the other (nor values already built), and a from_params(params) classmethod whose build() reproduces params exactly — every init field, the same state() key order and omitted defaults, so the same run id. Configuration containers (param_groups, the per-layer lists, the optims / schedulers maps, extra_metrics) are copied one level deep while params values, loaders (any iterable), metric callables and registered-factory specs keep their identity; neither method iterates data, builds a model, advances an RNG or does I/O. Setters, last-variant-wins, modifier retention and .vocab()'s dimension reset are unchanged. from_params rejects subclasses and foreign values with a TypeError naming the type (an NNOptimFactoryParams goes straight to NNTrainerParamsBuilder.optimizer) and a field the builder cannot carry with a ValueError. NNTrainerParamsBuilder also gains data_id(...) and overwrite_existing(...) setters, so sibling branches can train on distinct data_ids. See examples/builder_branching.py.
• Probability-aware prediction (NNModel.predict_proba, nnx.prediction). An opt-in predict_proba(X, ProbabilitySpec(kind, class_axis, labels)) returns a PredictionResult with the raw logits (identical to predict()), probabilities (categorical: softmax over the declared class axis, rows sum to 1; bernoulli: independent element-wise sigmoid, not normalized), decoded values (argmax classes, or logit >= 0 indicators), the spec's class axis and ordered labels, and per-row sample_ids (input row index for arrays and tensors, iteration position for ordinary loaders — a shuffling loader warns — and the global node index, from n_id / input_id, for graph seed rows). The task kind is declared, never inferred; a wrong label count or an invalid class axis (checked on a loader's first batch) or non-finite logits raise PredictionValidationError, and Bernoulli indicators refuse class_indices. It accepts the same inputs as predict() (arrays, tensors, tuples, loaders with graph seed-row slicing), restores every submodule's training mode, never touches parameters or gradients, and ProbabilitySpec.state() / from_state() round-trip label order and class axis. prediction_from_logits applies the same rules to existing logits. predict() and PredictResult are unchanged; examples/01_synthetic_classification.py compares both APIs on the reloaded BEST model (#211).
• Optims.ADAMW and an Adam-family eps. Optims.ADAMW / NNOptimParams.builder().adamw(...) build torch.optim.AdamW (decoupled weight decay: the weights are scaled by 1 - lr * weight_decay before the Adam update, where ADAM / ADAM_AMSGRAD keep their coupled L2 decay), matching native AdamW through NNParamGroupSpec overrides. NNOptimParams.eps (default 1e-8, torch's) is validated as finite and positive, forwarded to the Adam family, rejected on SGD variants, and omitted from state() at its default, so every existing run id is unchanged (#235).
• Registered optimizer factories (nnx.optimizers). register_optimizer_factory(id, version, factory) plus NNOptimFactoryParams(factory=OptimizerFactorySpec(id, version, config), ...) let a run use an optimizer NNx does not ship without a fake built-in name. NNModel.train() and Trainer.train() now build every optimizer through one shared hook, nnx.build_optimizer, before the run is reserved: the factory is resolved in the registry (an unknown id/version fails up front, importing nothing), called once with exactly the parameter groups a built-in optimizer would receive (one shared ownership rule, resolve_param_groups; each group with explicit lr / weight_decay, strict Trainer ownership applied) and a read-only JSON-like config, and must return a torch.optim.Optimizer over exactly those parameters. run.yaml records only the spec (callables are rejected), so NNRun.load, str(run) and the notebook view reload offline without running factory code and name the optimizer by id@vN; warm resume additionally validates the recorded factory id, version, config (type-strict: 1 and 1.0 differ) and ordered named-parameter topology. New example examples/optimizer_factories.py (#235).

Changed

• NNEvaluationDataPoint's f1 / recall / accuracy / precision are now optional (default None), so regression and other non-classification records no longer have to fabricate them; records built with all four are unchanged. PredictionResult.probabilities and .spec are Optional for continuous (regression-task) results, whose kind is "continuous". examples/26_custom_eval_step.py now leaves the classification fields None in its regression records and points plain regression at TaskSpec.
• The declared PyTorch floor is now torch>=2.4 with torchvision>=0.19 (was torch>=2.0 / torchvision>=0.15). The advertised floor was never honest: the CUDA mixed-precision path constructs torch.amp.GradScaler("cuda"), which does not exist before 2.3, and on 2.3 the full test suite still fails in three library paths (torch.is_autocast_enabled(device_type) in the transformer forward, SDPA float masks for half/bf16/double queries, and the weights-only unpickling allowlist for the training-state sidecar). 2.4.1 / 0.19.1 is the oldest pair on which the whole suite passes; CI now runs a floor-deps lane on exactly that pair (CPU, core dependencies only) next to the frozen all-extras lane, and docs/external-contracts.md publishes the tested matrix. No fallback factory was added: the scaler protocol, checkpoint serialization and resume validation are unchanged. CUDA execution of AMP remains unverified in this change (no CUDA hardware in the local or CI environments) and is recorded as such.

Fixed

• prediction_from_logits(..., sample_ids=...) refuses uint64 ids above 2**63 - 1 with PredictionValidationError instead of silently wrapping them to negative int64 ids, as nnx.calibration and nnx.abstention already did (#218).
• NNTabularDataset accepts a DataFrame whose rows are reversed or strided (df.iloc[::-1]): it previously failed with torch's "At least one stride in the given numpy array is negative" during tensor conversion (#239).
• NNTransformerParams.from_state now restores the inherited per-layer activations and dropout_probs lists that state() records — those that differ from the net-wide scalar, e.g. [RELU, TANH] / [0.1, 0.3]. It dropped them to None, so such a transformer config no longer equalled itself after a run-YAML, checkpoint or safetensors round trip. Lists that merely repeat the scalar are still omitted from state() and reload as None with an identical state() and run id (as for base NNParams); states without those keys load unchanged. NNTransformerParams.builder().build() now supplies its dropout_prob=0.0 default under the set fields instead of beside them (with hidden_dims / activation left to the NNParams defaults), so from_params can carry the inherited fields.
• NNModel.train() and Trainer.train() no longer print the absolute working directory in their completion line: Run saved to <cwd>/runs/<id> is now the portable display path Run saved to runs/<id> (relative to the working directory), derived from the persistence root in one shared helper, so captured notebook output no longer embeds the executing worktree, temporary directory or container mount and is identical across machines apart from the run id. The line is documented as a display path, not artifact provenance; where and how runs are saved is unchanged (#190).
• RMSNorm no longer truncates FP64 inputs to FP32: the mean of squares is computed in promote_types(x.dtype, float32), so FP16/BF16 still reduce in FP32 while FP64 stays FP64. A double TransformerNN, ViTNN or JEPAPredictor no longer collapses large finite activations to zero (a 1e30 input used to normalise to 0) and no longer carries FP32 rounding (~1e-8) through every norm. The cast back to the input dtype still precedes the weight multiply, so output dtypes, parameters, eps and checkpoint keys are unchanged. (#261).
• Double-precision attention is correct again at ordinary sequence lengths: for FP64 queries the dropout-free SDPA path of multi_head_causal_attention now upcasts its FP32 additive causal / prefix mask to FP64. CPU scaled_dot_product_attention silently mis-applies an FP32 mask to FP64 queries once the sequence reaches the kernel's vector width (8 tokens on AVX2, 16 on AVX512), so every attention layer of a double TransformerNN or PrefixTuner returned wrong outputs (errors of order 1) from that length on; the full forward and incremental forward_with_cache decoding could therefore disagree, and greedy generate with and without the KV cache diverged on AVX2 machines. FP16/BF16/FP32 queries receive the mask unchanged (#261).
• Run text files (run.yaml, metadata.yaml, idps.csv, the history marker and the best/POINTER.txt fallback) no longer leak .<name>.XXXXXX temporaries: the atomic writer now owns its temp until the rename, so a failure while opening, writing, flushing or closing the temp, or in the rename, releases the descriptor, unlinks that temp and re-raises the original exception (a cleanup error, including a failing unlink, never masks it) while the destination keeps its previous bytes; a successful rename never touches the temp name again, and a bare filename keeps its temp in the destination directory instead of the system temp dir. A failed history-marker write now also releases the empty runs/<id>/ reservation, so the run can be retried instead of failing with "already exists" (#260).
• Documentation now matches runtime behaviour for custom metrics, EarlyStopping and parameter groups: every custom-metric doc and docstring states the (y_true, y_pred) argument order (the Trainer builder said (y_pred, y_true)) and when metrics run (per batch in the default training step, once on the aggregate in evaluate(), at the custom step's discretion otherwise); the quickstart metric example no longer feeds class predictions to roc_auc_score; EarlyStopping's mode is described as reversing the improvement direction for one of the four accepted loss/error monitors (accuracy/F1 monitors are not accepted); and the parameter-group examples in NNParamGroupSpec, the concepts guide and the quickstart put the specific encoder.*bias rule before encoder.* and *.bias, stating that the first matching rule wins and rules never merge (the old order left encoder biases with weight decay). The Trainer builder docs and the GAN example state that the step function owns optimizer updates while the Trainer steps schedulers per epoch unless auto_step_schedulers=False. New executable doc-snippet and Trainer tests pin each contract, and examples/03_custom_metrics.py (now executed by the smoke tests) checks an asymmetric metric against its validation labels and after NNRun.load (#259).
• NNTabularDataset now rejects feature data it used to admit silently: ±inf in a selected feature column is rejected in source precision before the feature_dtype cast (an int64 cast turned it into an INT64 extreme and a bool cast into True); finite values outside an integer feature_dtype's range (e.g. 300 → int8, 1e30 → int64), which the cast would wrap, are rejected in source precision too; and a finite value that overflows a narrower floating/complex dtype during conversion (e.g. 1e5 → float16) is rejected for features and for a floating target_dtype target. Errors name the columns and the requested dtype (the NaN and non-finite-target diagnostics now name the dtypes too). Every admission check, including the classification contiguity check that used to run after the split, now runs before random_split, so a rejection leaves the DataFrame and the global RNG untouched; unselected columns are still never inspected (#233).
• apply_lora_to, apply_dora_to and apply_ia3_to no longer split a shared layer: patterns now match every registration path (not only the first one named_modules() reports) and targets are grouped by registration slot. A Linear inside a shared container is wrapped once and every path sees the same wrapper (selecting it through its non-first path used to return 0); a Linear registered under more than one name — which a wildcard used to wrap at one path only, leaving the other alias on the raw layer — now raises ValueError naming every alias. The whole matched set is validated before any wrapper is built, so a rejected call changes no module, weight, requires_grad flag or parameter-group membership, and an unselected alias does not block independent targets. Tied tensors between distinct modules are outside this check (#224).
• PEFT wrappers now inherit the train/eval mode of what they wrap at construction: LoRALinear / DoRALinear (directly or via apply_lora_to / apply_dora_to) copy the wrapped Linear's training flag onto the wrapper and its lora_dropout, IA3Linear onto the wrapper, and PromptTuner / PrefixTuner the wrapped TransformerNN's flag onto the tuner. Injecting a nonzero-dropout adapter into an eval() model used to produce a train-mode wrapper whose dropout made inference nondeterministic once the adapter was nonzero, and a mixed-mode parent had every replaced child flattened to train mode. The surrounding model is never switched, modes are not serialized, and a later model.train() still activates the adapter dropout (#223).
• EarlyStopping()'s default monitor no longer goes silently inactive on regression runs: monitor now defaults to None, which selects val_edp.error when the first validated epoch reports a finite error and otherwise val_edp.loss (evaluators that leave error as None), then keeps that field for the whole train() call and re-selects on the next call (exposed as selected_monitor). Classification runs keep tracking the validation error, and the default still never reads training metrics; mode="max" now requires an explicit monitor. An explicit monitor still compares exactly the named field with no fallback, but a missing field or data point (for example no validation loader) is no longer silent: the epoch is not counted and one RuntimeWarning per run, attributed to the calling code and repeated for every run, names the monitor and what is absent. A NaN/±inf value no longer becomes the best (which made every later comparison fail); it counts as an epoch without improvement (#189).
• NNGraphDataset now represents an empty val_mask / test_mask as an absent split in both sampler modes: that loader is None (the optional-loader contract NNDatasetBase, NNModel.train and Trainer.train already honour, so validation is skipped rather than run over zero seed rows), the resolved size is 0 and state() reports "0"; an explicit positive request for an empty neighbor split also yields None. Previously full mode built a zero-node batch that made training fail with evaluate() loader produced zero samples, and neighbor mode failed inside NeighborLoader with batch_size=0. An empty train_mask now raises a ValueError naming it at construction instead of producing a zero-size training loader. Nonempty splits, the full-graph permutation / input_id seed-row scoring and sampler="full"'s rejection of explicit sizes are unchanged.
• Dataset batch_sizes tuples now treat None as the only automatic full-split sentinel in NNDataset, NNTabularDataset, NNPreferenceDataset and NNGraphDataset: zero and False — which truthiness resolution silently turned into a full-split request, changing memory demand and optimizer cadence — as well as negatives, floats (2.0 included), strings and NaN raise a ValueError naming batch_sizes[i] (split), and a list, a bare value or a tuple of the wrong length raises with the 3-tuple contract instead of an IndexError after the dataset was built. Validation runs before the dataset factory, tokenizer.encode or any RNG-consuming split; positive sizes are kept verbatim (a size larger than the split yields one smaller batch, never a clipped metadata value); NumPy integers are normalized to int (a raw numpy.int64 used to be rejected by DataLoader). The empty-split conventions (placeholder 1 plus a None loader) and sampler="full"'s rejection of explicit sizes are unchanged. The four wrappers' docstrings, the quickstart and the DPO guide now state that an omitted / None slot is one batch holding the complete split — one optimizer step per epoch for the default train loader — and show an explicit mini-batch configuration; the runtime default itself is unchanged (#188).
• Public configuration counts are now validated as integers at their boundary: NNParams.input_dim / output_dim / hidden_dims entries / n_heads, NNTransformerParams.vocab_size / n_layers / d_model / max_seq_len / ffn_mult (plus an even d_model / n_heads head width, which RoPE needs, checked at construction and in .builder().layers(...) before the divisibility arithmetic), NNConvParams.conv_channels / in_channels / kernel_size / stride / pool_size (positive) and padding (≥ 0), NNMoEParams.num_experts / top_k, NNOptimParams.accumulate_grad_batches, NNSchedulerParams.patience / cooldown (≥ 0) and step_size / T_max / total_steps / warmup_steps (> 0 when set), NNTrainParams.n_epochs, NNTrainerParams.n_epochs, EarlyStopping(patience) and direct RoPE(dim, max_seq_len). Any non-boolean numbers.Integral — NumPy integers included — is accepted and normalized to a plain int, so state() stays YAML-portable and hashes like the plain-int configuration; booleans, every float (2.0 included), numeric strings and NaN / ±infinity raise a ValueError naming the field before any range, modulo, isqrt or layer allocation. Previously accumulate_grad_batches=2.5 reached the modulo and silently stepped every five batches, and a NumPy-integer count made run.yaml serialization fail. None sentinels, zero-valued padding / patience / cooldown and real-valued hyperparameters are unchanged.
• Public real-valued hyperparameters now reject NaN and ±infinity (and non-real values such as booleans or numeric strings) at their configuration boundary with a ValueError naming the field and its domain: NNOptimParams (max_lr, weight_decay, grad_clip_norm, SGD momentum, Adam betas), NNSchedulerParams (factor, min_lr, threshold, max_lr), NNParamGroupSpec (lr, lr_multiplier, weight_decay, plus the resolved per-group learning rate in build_param_groups), NNTransformerParams.rope_base and direct RoPE(base=...) (which must be positive and now cannot produce non-finite tables), RMSNorm(eps), EarlyStopping(min_delta) and LoRA / DoRA alpha. Meaningful zeros and None sentinels stay valid and serialized state is unchanged.
• deepen now honours dotted names for ReLU insertion in named nn.Sequential containers (root or nested): the site is located by registered key instead of int() coercion, the container is rebuilt from its ordered _modules so a module registered in several slots keeps every slot (a positional rebuild through children() used to drop one and silently corrupt the forward), the inserted modules get collision-free _nnx_deepen_linear_N / _nnx_deepen_relu_N keys, the container's training mode is preserved, and the surgery path helpers address a registered key exactly so retained numeric keys stay reachable after an insertion.
• deepen on a FeedFwdNN now updates the network's immutable params together with its layers: one more hidden_dims entry at the insertion site, a ReLU inserted into an explicit activations list, and dropout 0.0 at the new identity site with every existing site's probability preserved (a nonzero scalar dropout is materialized into a per-layer list). The ReLU check reads the effective site activation (activation_for(i)), so a per-layer override is honoured over the scalar default in both directions. Previously the forward of a deepened net with per-layer overrides raised IndexError, and the inserted layer inherited the scalar dropout.
• widen now validates the target→consumer path before allocating anything: the consumer is inferred only inside nn.Sequential or FeedFwdNN.layers (resolving the effective per-layer activation/dropout, not the scalar defaults) and only through elementwise ops; Softmax / LayerNorm / BatchNorm boundaries, other containers whose registration order is not data flow, missing consumers and modules registered under more than one path raise a ValueError naming the boundary with the source model, module identities and RNG untouched. Previously a Linear → Softmax → Linear net was silently widened with a 0.12 output drift.
• widen and low_rank_factorize now preserve trainability and module modes: the widened target and its downstream consumer, and both SVD factors, inherit the requires_grad flag of the weight/bias they replace (bias flags independently; up.bias carries the source bias role) and the train/eval mode of the module they replace, instead of defaulting every new tensor back to trainable/train mode — an optimizer built after surgery no longer silently updates previously frozen content.
• Manual attention (the path taken when attn_dropout > 0 in train mode) now accumulates scores, additive mask, softmax and dropout in FP32 for FP16/BF16 inputs (FP32 and FP64 keep their own dtype) and casts the result back to the value dtype, so a half or bf16 TransformerNN / PrefixTuner trains with attention dropout instead of failing with expected scalar type Float but found Half; the attn_dropout=0 SDPA path, seeded determinism, causality and eval-mode behaviour are unchanged.
• Adapter-only PEFT artifacts are now selected by registered ownership instead of key substrings: save_lora_weights / load_lora_weights (the lora_A/lora_B of each LoRALinear, DoRA included), save_ia3_weights / load_ia3_weights (each IA3Linear.scaling), PromptTuner.prompt_state_dict / load_prompt_weights (soft_prompt) and PrefixTuner.prefix_state_dict / load_prefix_weights (prefix_keys.i / prefix_values.i). A module merely named lora_A_projection or scaling_projection no longer leaks its frozen base.* tensors into the file, a full or colliding source dict can no longer overwrite a frozen base tensor on load (the allowlist is built from the destination), and loaded counts report only accepted adapter tensors. File format and key names are unchanged.
• PEFT adapters now allocate their parameters from the tensor they compose with, inheriting its dtype and device: LoRALinear / DoRALinear (lora_A, lora_B) and IA3Linear (scaling) from the wrapped weight, PromptTuner (soft_prompt) from the token embedding, and PrefixTuner (per-layer K/V) from each block's attention projection. Wrapping a base that was already converted to half/double or moved to a device — or a meta-device skeleton — composes immediately without a second .to(), and LoRA, IA3, PromptTuner and PrefixTuner still compose with a torchao int8-quantized base (quantize_int8), as in 0.2.3; IA3 no longer silently promotes a half pipeline to float32; the base is never moved or recast.
• DoRALinear now normalizes rows in FP32 for FP16/BF16 layers (FP64 stays FP64) and casts the effective weight back to the layer dtype, so a zero or tiny row of the composed weight — after pruning, explicit zero initialization or a zero row in a loaded base — yields a finite output equal to the base instead of NaN (FP16 rounded the 1e-8 guard to zero), while learned nonzero adapter/magnitude updates stay differentiable and output dtype/device are unchanged.
• quantize_int8 now returns an instance of the same class it was given (typed with a bounded TypeVar): a GenerativeNNModel keeps its tokenizer and generate, and user subclasses keep their overrides and custom attributes, instead of being downgraded to a bare NNModel. The wrapper is a shallow copy (no constructor re-run, no RNG consumed) sharing params/loss/tokenizer with the source; only net is the independent quantized deep copy.
• Warm resume now accepts the same batch sources ordinary training accepts — a DataLoader, a plain re-iterable list of (X, Y) batches, or NNGraphDataset(sampler="full")'s one-element list — instead of failing with AttributeError: 'list' object has no attribute 'num_workers' after restoring state; worker capability is read safely before restoration, and the worker-RNG warning still fires for a real loader with num_workers > 0. NNTrainParams.train_loader / val_loader and NNModel.evaluate are annotated as Iterable accordingly.
• Losses.NEGATIVE_LOG_LIKELIHOOD now optimizes and reports the normalized negative log-likelihood: the supervised training step, evaluate() and the NNx-owned classifier step factories (KD, feature-KD, MoE, Mixup, CutMix, Born-Again) apply log_softmax over the class axis before the native torch.nn.NLLLoss, so NLL and cross-entropy produce identical losses and parameter updates from the same weights (previously raw logits were passed in, yielding negative losses and a gradient without the competing-class term). predict().logits stays raw; custom loss modules, NLLLoss subclasses with their own forward, custom train_step_fns and lr_finder callables keep receiving the raw output.
• Overwriting the run that owns runs/best (overwrite_existing=True) now re-elects the best surviving committed run before the old artifacts are deleted, instead of comparing the replacement with itself and leaving the pointer on a worse or emptied run; only runs with a committed BEST checkpoint can claim the pointer (reservations, .leases, malformed artifacts and speculative history beyond LAST are ignored), and the pointer is removed when no eligible run remains.
• NNCheckpoint.load_training_state, load_optimizer_state and load_with_training_state no longer create runs/<id>/checkpoints/ when probing a run that does not exist: they return their absent result without touching the filesystem, so a preflight probe no longer makes the first train() fail with "run already exists". Existing runs keep the checkpoint lock and generation validation; run IDs are still validated first.
• NNTrainerParams now snapshots optims, schedulers and extra_metrics into read-only mappings at construction (direct constructor, builder().build(), from_state, dataclasses.replace): reusing a builder or mutating the dict you passed in no longer alters an already-built configuration, its state() or its run identity, and item assignment/deletion through the built mappings raises TypeError. Values (params dataclasses, metric callables, loaders) keep their identity; serialized state is unchanged.
• BEST-checkpoint selection, the runs/best pointer and ReduceLROnPlateau stepping now use the first finite metric in the validation-error → validation-loss → training-error → training-loss order: NaN/±inf values are skipped with a per-epoch RuntimeWarning naming the rejected field, a non-finite validation error no longer hides the same epoch's finite loss, an epoch without any finite signal skips the plateau step and compares as an unavailable baseline, and legacy NaN-ranked BEST artifacts are replaced by the next finite candidate. Raw observations in run history and checkpoints are unchanged.

0.2.3 (2026-08-17)

Changed

• Maintenance and publishing tooling now use Setuptools 84.0.0, Markdown 3.10.3, and uv 0.12.3 consistently across requirements, workflows, documentation, and the frozen lockfile.

Fixed

• Documentation now renders a portable API catalog on all three surfaces, validates canonical GitHub links and tracked-file ownership, documents recent public APIs, and corrects stale DPO, comparison, tooling, and release-contract guidance.
• PyPI trusted publishing now runs directly in the top-level release.yml workflow dispatched by Release Please, keeping the OIDC publisher and Sigstore attestation identities aligned instead of publishing from an unsupported reusable-workflow context.
• Draft GitHub release assets are verified through the release database ID because GitHub's tag lookup endpoint does not expose draft releases.

0.2.2 (2026-08-09)

Added

• NNTabularDataset accepts an optional target_dtype (a floating-point dtype, e.g. torch.float32) to skip the integer-label cast and contiguity check and fix output_dim=1, adding first-class regression-target support without the previous "build the DataLoaders yourself" workaround; regression loaders yield targets shaped (batch, 1) to match output_dim=1; integer dtypes are rejected with a clear error; the default (None) preserves the existing classification contract.
• NNRun accepts an optional salt string folded into the run.id hash, letting identical (model, net, train) configurations run as distinct experiments without modifying modeled params; salt=None (the default) preserves existing run.id hashes exactly; empty/whitespace salts are rejected with a clear error (mirrors NNTrainParams.data_id).
• Versioned warm-resume bundles now persist optimizer, scheduler, scaler, completed epoch, and Python/NumPy/PyTorch RNG state in generation-addressed sidecars, preserving the previous resumable generation across interrupted commits.
• Smoke coverage imports every numbered example and executes representative basic, callback, and custom-evaluation examples; release tests also verify normalized source distributions are reproducible.
• A project security policy documents supported versions, private vulnerability reporting, and responsible disclosure expectations.
• Trainer scheduler ownership is explicit through auto_step_schedulers, with loop-owned once-per-epoch stepping as the default.
• A manifest now deterministically projects canonical documentation into self-contained MkDocs and GitHub wiki trees.
• dpo_train_step_factory now surfaces reward_chosen, reward_rejected, and reward_accuracy (the fraction of the batch where the implicit chosen reward exceeds the rejected reward) in each NNEvaluationDataPoint.extra, giving DPO training runs a ranking-quality diagnostic alongside the existing log-prob gap.
• GenerativeNNModel.generate() accepts an optional on_token callback invoked after each generated token, enabling streaming output and progress reporting without re-running decode.

Changed

• Documentation now opens with a centered product identity, synchronized health and technology badges, a wide brand banner, and an executive summary; navigation numbering remains independent from the unnumbered product H1 across the repository, Pages, wiki, and PyPI projection.
• GitHub workflows use current SHA-pinned actions, least-privilege permissions, frozen dependency and tooling resolution, a double-build artifact comparison, and one Release Please-only publication path; managed release PRs refresh the lockfile and dispatch required checks, security audits cover pushes and pull requests on both Gitflow branches, and exact PyPI/GitHub hashes and immutable-release attestations are verified before completion.
• Release builds use the locked virtual-environment interpreter, dependency audits include locked tool groups, and the release toolchain tracks wheel 0.47.0.
• Package metadata uses a deterministic PyPI README projection whose repository links and images resolve outside GitHub.
• Optional notebook display support moved out of the core dependency set, while process-safe persistence locking is now an explicit runtime dependency.
• Examples and public documentation now describe experimental GGUF/Ollama export, whole-loader evaluation hooks, partial gradient accumulation, and complete resume semantics accurately.
• Warm resume restores loader and sampler generators by stable identity and validates optimizer type, parameter topology, and scheduler identity before applying state.

Fixed

• Final partial accumulation windows now perform an optimizer step, and custom evaluation runs once per epoch over the complete validation loader.
• Gradient accumulation now identifies the observed final batch and weights uneven batches by sample count, matching one combined effective-batch update.
• Accumulation and evaluation use the effective CE/NLL denominator for class weights and ignored targets, preserving combined-batch loss semantics.
• Accumulation and evaluation honor configured loss hooks and mean/sum reductions; fully ignored cycles fail before optimizer updates.
• Ignored CE/NLL targets are excluded from classification metrics, custom subclasses keep their own normalization contract, and LR finding rejects live persistent-worker iterators it cannot rewind.
• Dataset, sampler, callback, learning-rate finder, trainer, and Ollama option boundaries reject invalid, non-finite, fractional, unsafe, or unknown inputs before execution.
• Run identity includes data lineage, mutable parameter defaults are isolated, callback-finalized trainer checkpoints preserve transforms, and born-again generations receive distinct lineage-aware runs.
• Checkpoint and run writes use process-safe locks plus unique atomic temporary files; LAST commits an epoch only after history persists, callback checkpoints flush after that boundary, and failed commits roll history back.
• Model snapshots preserve PyTorch state-dict metadata and non-tensor extra state; tensor-only safetensors and Hub exports now reject incompatible extra state with a targeted error.
• Learning-rate finding restores Python and NumPy global RNG state, and warm-resume bundles now preserve Apple MPS RNG state alongside CPU and CUDA state.
• Inference and inspection helpers restore mixed per-module train/eval modes exactly instead of flattening every child to the root mode.
• Mode restoration invokes recursive Module.train() only at actual mode boundaries, avoiding redundant custom-hook side effects in homogeneous subtrees.
• Mixed-mode restoration invokes each custom Module.train() hook exactly once, and checkpoint cleanup enumerates literal directories so metacharacters in root paths cannot affect siblings.
• Child-aware mode hooks retain access to registered modules, probability-target cross entropy reports class-index metrics, and pickle checkpoints/sidecars load portably onto CPU by default.
• Mode restoration suppresses direct calls to any descendant hook, and legacy feed-forward/transformer file loaders also default to CPU with an override.
• Transformer token logits now flatten correctly for standard cross-entropy training, mode hooks restore bottom-up, high-level checkpoint reconstruction accepts a device override, and safetensors honors load placement.
• Transformer probability targets flatten along the class-last layout, predictions select the final class axis, malformed KV caches fail clearly, Hub loads override serialized device metadata, and safetensors rejects unsupported location mappings.
• Transformer-only class-last handling no longer changes class-first segmentation outputs, and Hub loading rejects indexed devices that the serialized device contract cannot preserve.
• Checkpoint reconstruction preserves subclasses, Transformer KV caches validate their full tensor contract, and Hub safetensors loading normalizes torch.device locations.
• Transformer cache validation supports active autocast, checks key/value symmetry and tuple shape, and subclass-preserving reconstruction is reflected in static Self typing.
• Multidimensional probability-target metrics flatten consistently, generative checkpoint reconstruction accepts its tokenizer, and Hub artifacts package tokenizer and topology-transform metadata for loadable generative and converted-QAT round trips.
• Classification error follows subset accuracy for multidimensional labels, cross-entropy subclasses retain metric preprocessing, built-in binary classification uses logits with a zero threshold, Trainer exports retain completed topology transforms, and remote Hub artifacts load from one immutable snapshot.
• Soft binary targets are thresholded only for classification metrics, preventing valid BCE training from failing after an optimizer update.
• Persistent NNModel training rejects low-rank surgery without a reconstruction recipe before mutating state; the surgery example uses manual refinement and state-dict export semantics.
• Explicit custom training steps remain available for reconstructibility-aware custom topologies such as diffusion, and frozen evaluation records store custom metrics in an immutable hashable mapping.
• Inherited native CE/NLL losses retain class-weight normalization, and repeated NNModel or Trainer transformations compose reconstruction recipes instead of replacing them.
• Optimizerless checkpoints explicitly reject stale legacy sidecars after interrupted cleanup, warm resume rejects GradScaler presence changes, and mode restoration honors custom Module.train() hooks.
• Run identifiers reject glob metacharacters before sidecar cleanup, and subclasses of built-in elementwise losses retain the custom-loss normalization contract.
• Documentation projection rejects missing local links, malformed nested manifests, unsafe output slugs and collisions, source paths outside the repository, symlink destinations, repository ancestors, and nonempty unmanaged outputs.
• Documentation projection checks are non-mutating, fenced examples are ignored during link processing, and cleanup rejects symlinks in every output path component.
• Documentation projection preserves inline code and link titles, handles nested fence-like examples correctly, and rejects unmapped repository-file links instead of erasing them.
• Documentation projection rewrites reference-style links, rejects nonportable root-relative URLs, and validates setext headings and explicit anchor IDs.
• Documentation projection uses balanced escape-aware link scanning, decoded local-path lookup, HTML-comment exclusion, and Python-Markdown anchor slugs.
• Documentation projection tokenizes multiline links/references, indented code, contextual comments, and blockquoted headings correctly.
• Documentation projection rejects image traversal and unknown manifest keys, suppresses blockquoted fences and raw <pre> blocks, and derives wiki anchors from rendered GitHub-style heading text.
• Documentation projection respects blockquote and raw-HTML containers, rewrites and validates HTML links/images, rejects unpublished image sources, and allocates globally unique wiki anchors.
• Documentation projection handles quoted and unquoted HTML targets, list-contained fences/headings, site superfence anchors, and nested wiki asset paths without flattening collisions.
• Documentation projection uses offset-preserving HTML attribute scanning, excludes literal blocks from HTML processing and anchor discovery, preserves nested assets on both surfaces, and rejects duplicate manifest keys.
• Documentation projection enforces output containment, depth-aware raw/list containers, MkDocs extension parity, complete HTML resource attributes, contained assets, unique attributes, and HTML-entity lookup.
• Documentation projection rejects asset-directory symlinks, distinguishes list prose from code, covers literal HTML forms and standard resources, preserves data-URL srcset, and derives wiki anchors from rendered structure.
• Documentation projection handles mixed srcset, complete URL-bearing attributes, blockquoted raw HTML, list references/explicit IDs, independent heading suffixes, and copied HTML/CSS validation.
• Documentation projection separates descriptorless data-URI srcset candidates, honors blank-line raw-HTML termination, projects nested-list IDs, parses CSS comments/imports, accepts the standard macOS /tmp alias, and type-checks its scripts.
• Documentation projection preserves GFM raw-HTML blocks through the following blank line, recognizes nested-list HTML, validates embedded CSS and copied SVG resources, and rejects incomplete projected assets.
• Documentation projection closes raw/comment blocks at list boundaries, enforces declaration case, projects nested Setext IDs, tokenizes CSS functions and strings, validates additional SVG references, and rejects empty manifest groups.
• Documentation projection handles list-contained literal HTML, CSS escapes and image sets, literal-code wiki headings, broader SVG references, and balanced nested SVG diagram extraction.
• Documentation projection closes processing-instruction and declaration literals at list boundaries, consumes CSS escapes while scanning nested functions, and validates SVG stylesheet processing instructions.
• Documentation projection measures tab indentation in Markdown columns, handles escaped CSS imports and line continuations plus nested image(), and ignores SVG stylesheet text inside comments and CDATA.
• Documentation projection applies tab columns to fences and reference definitions, covers directional CSS image(), tokenizes stylesheet PI attributes, accepts standard macOS temporary aliases, and wiki publication can initialize an empty remote.
• Documentation projection preserves partial tab-stop indentation, normalizes invalid CSS escapes to U+FFFD, and rejects duplicate or malformed XML stylesheet pseudo-attributes.
• Documentation projection recognizes a dedented closing marker for a list-contained fence instead of reopening suppression across following prose.
• Safetensors checkpoints reject missing or unsupported format versions instead of attempting an incompatible load.
• The architecture diagram now renders completely on desktop while retaining a contained, readable horizontal scroller on small screens.

Security

• Dependency alerts, Dependabot security updates, private vulnerability reporting, secret scanning with push protection, CodeQL default setup, immutable releases, and protected merge-only Gitflow rules are enabled at the repository level.
• Required gitflow checks are bound to the GitHub Actions integration rather than accepting same-named statuses from arbitrary writers.
• The v0.2.1 GitHub release now carries the exact wheel and source distribution published to PyPI; repository immutable releases are enabled for subsequent publications.

0.2.1 (2026-07-22)

Features

• ConvNN (LeNet-style) + NNConvParams + Nets.CONV (#89) (46fdfef)
• dataset: full-batch node-classification loader (no pyg-lib) for NNGraphDataset (2b8b163)
• FeedFwdMoENN net + Nets.FEED_FWD_MOE + NNMoEParams (5d3956f), closes #88
• per-layer activation & dropout on NNParams (net-wide → optional lists) (4a5de0c), closes #85
• pluggable eval_step_fn on NNModel.train (mirror train_step_fn) (822b6da), closes #86

Bug Fixes

• complete specialized network exports (91ff6ef)
• harden training lifecycle and type contracts (c6b322e)
• promote GitPython security refresh (1c03f56)
• re-save LAST checkpoint after on_train_end so callback net mutations persist (65e599c), closes #87
• reconstruct converted QAT checkpoints (24d3262)
• reject missing ConvNN scalar activation (3476d2e)
• reject single-expert MoE params (#114) (a395e2d)
• require hidden layers for MoE params (#116) (70bf913)
• support pandas 2 type contracts (d9c523e)

Documentation

• examples: propagate the new primitives (#85–#89) to the examples (16fa58b)
• link the rendered docs site from README (#68) (3cf1682)
• render architecture diagram in a themed page (fix bare-HTML nav break) (#69) (8940c3e)
• Slate + Cobalt theme for the docs site (#66) (7520455)
• synchronize architecture and publishing surfaces (14a4d2c)

Changed

• Release verification now covers NNx Studio's required API floor. The clean-environment post-PyPI smoke imports NNMoEParams, NNConvParams, FeedFwdMoENN, and ConvNN in addition to checking nnx.__version__, preventing a correctly numbered but API-incomplete wheel from satisfying the release gate. Release documentation now explicitly prohibits distributing artifacts built from untagged commits under the static release-please-managed version.
• Public architecture and evaluation surfaces are complete. ConvNN and FeedFwdMoENN now follow the other first-class networks onto the top-level nnx facade; the API reference includes EvalStepContext / EvalStepFn, NNConvParams / ConvNN, and NNMoEParams / FeedFwdMoENN; and the concepts guide documents the shipped custom-validation hook instead of listing it as future work. Full-batch graph behavior and per-layer activation/dropout configuration remain covered by their existing guides and tests.
• GGUF support is labeled experimental at every user-facing surface. NNx writes a structurally valid nnx_transformer container, but stock llama.cpp, Ollama, and LM Studio do not implement that architecture. Docs and examples no longer recommend relabeling it as LLaMA or claim stock serving compatibility, and the quantization recipe points to the official llama.cpp source build.
• CI, docs, security, and release jobs now install the frozen all-extras graph. Pyright warnings are a required CI gate, security audits the exact exported lock, build tools are pinned, release-please invokes the reusable publisher in the release-creation run, and the source-tree fallback version participates in release-please updates.
• Added dependency/security maintenance automation. Dependabot now watches GitHub Actions and Python dependency manifests weekly, and a new security workflow runs pip-audit against the project dependency graph on relevant PRs, weekly schedule, and manual dispatch.
• Pinned GitHub Actions workflows to immutable SHAs. CI, docs, release, release-please, and security workflows now execute exact action commits, with comments preserving the source tag/branch each SHA was resolved from.
• Added a committed dependency lock and drift gate. uv.lock now captures the resolved dependency graph, requirements-tools.txt pins the resolver tool, README documents frozen uv sync, and CI/release gates run uv lock --check before installing test extras.
• Relicensed from MIT to the Apache License 2.0. The LICENSE file now carries the full Apache 2.0 text and the package metadata declares license = "Apache-2.0" (SPDX). Apache 2.0 adds an explicit patent grant and patent-retaliation termination clause on top of MIT's permissions; it remains a permissive license. Copyright holder is unchanged (Kaveh Razavi).

Fixed

• Converted QAT checkpoints reconstruct through NNModel.from_checkpoint. The post-training save correctly persisted torchao's converted state, but the generic loader still built an FP32 network and could not absorb its quantized keys. NNCheckpoint now carries ordered, versioned topology-transform recipes in both pickle and safetensors formats; QATLifecycleCallback records qat_config and groupsize, and the loader replays that conversion before loading weights. Existing FP32 checkpoints keep an empty transform list, unknown recipes fail explicitly, and legacy converted QAT checkpoints without metadata receive a targeted compatibility error.
• NNConvParams requires a scalar activation at construction. Convolution blocks always use the net-wide scalar even when activations supplies per-layer overrides for the fully connected head, so activation=None previously constructed successfully and failed only on the first ConvNN forward. Direct and serialized construction now reject that invalid configuration at the params boundary; base NNParams still permits activation=None when complete per-layer overrides are present.
• NNMoEParams requires at least one hidden layer. The base params type intentionally permits a no-hidden-layer linear classifier, but carrying that allowance into FeedFwdMoENN produced a model with only its plain classifier head: zero MoELinear modules and a silently inactive MoE auxiliary loss. Both direct and serialized construction now reject hidden_dims=None / [], while plain NNParams retains its linear-classifier behavior.
• NNMoEParams now rejects num_experts < 2 at the parameter boundary. The serialized params type previously accepted a single expert even though MoELinear correctly rejected it during model construction, forcing downstream consumers to duplicate the stronger invariant. Direct construction and from_state() now fail early with the same num_experts >= 2 contract; valid MoE state and checkpoint formats are unchanged.
• Trainer LAST checkpoints include on_train_end mutations. The final checkpoint is refreshed after callback finalization, matching the persisted model to the completed lifecycle.
• Empty aggregation and prediction fail clearly. NNEvaluationDataPoint.mean_of([]) and NNModel.predict() on an empty loader now raise targeted ValueErrors instead of producing warnings, NaNs, or NumPy-internal failures.
• Network enum dispatch rejects incompatible specialized params. Convolutional, MoE, and transformer network variants now require their matching params classes instead of constructing a mismatched network and failing later.
• Frozen dependencies contain no known audited vulnerabilities. Pygments, setuptools, torch, and torchvision were refreshed to patched resolutions, and the security workflow now audits that exact graph.
• NNModel.train / Trainer.train callback cleanup is complete and exception-safe. Only callbacks whose begin hook completed are finalized, finalizers run in reverse order, every cleanup hook gets a chance to run, and cleanup errors no longer mask the original training failure. Begin-hook failures still clean up callbacks that started earlier.
• default_train_step rejects non-finite loss before backward/step. The guard previously raised after gradient propagation and optimizer stepping, so a NaN/Inf loss could poison model weights before the exception reached the caller. The finite-loss check now runs immediately after loss computation on both AMP and non-AMP paths, before any gradients or optimizer state are mutated.
• ONNX export fallback no longer masks real TypeErrors. NNModel.to_onnx and nnx.viz.netron_export now check whether the installed torch.onnx.export signature supports the dynamo keyword before deciding whether to omit it for legacy torch versions. Genuine exporter/model TypeErrors now propagate with their original message instead of being rewritten as an old-torch compatibility error.
• VisUtils.two_dim_tsne_checkpoint_logits now honors n_samples across test-loader batches. The helper previously read only the first test batch and then sliced it, so larger n_samples values silently capped at batch size. It now collects batches until the requested sample count, validates that at least two samples are available, uses a valid small-sample t-SNE perplexity, and defaults random_state=0 for reproducible plots.
• drop_layer rejects an ambiguous importance= scorer for a single target. The scorer only participates when layer_name is a candidate list; passing it with one string target previously looked meaningful but was silently ignored. drop_layer(..., layer_name="layers.1", importance=fn) now raises a clear ValueError instead, while drop_layer(..., layer_name=[...], importance=fn) keeps the minimum-score selection behavior.
• NNOptimParams fails fast on out-of-range accumulate_grad_batches / grad_clip_norm. Both fields constructed fine but misbehaved deep in the train loop: accumulate_grad_batches=0 died mid-training with ZeroDivisionError on batch_idx % accumulate_grad_batches (after printing the whole run-config table), a negative value silently scaled the loss by 1/N < 0 and performed gradient ascent, and grad_clip_norm=0.0 passed the is not None clip-enable check and zeroed every gradient so training ran to completion making no progress. __post_init__ now raises a clear ValueError for accumulate_grad_batches < 1 and for grad_clip_norm <= 0 (use None, not 0, to disable clipping) — same construction-time fail-fast convention already used for param_groups and the dataset classes.
• train_contrastive fails fast on a non-positive grad_clip_norm. The high-level embeddings fine-tune loop validated n_epochs / batch_size / temperature at the boundary but not grad_clip_norm, so grad_clip_norm=0.0 slipped past the is not None clip-enable check and clip_grad_norm_(..., 0.0) zeroed every gradient — the backbone trained to completion making no progress. It now raises a ValueError for any grad_clip_norm <= 0 (use None, not 0, to disable), matching the same fail-fast convention NNOptimParams already enforces for the identical footgun.
• NNModel.train / Trainer.train no longer print "Run saved to X" before the save actually completes. Both wrappers ended with print(...); return run.save() — if .save() raised (disk full, permissions, path inaccessible), the success line had already reached the user, contradicting the exception that followed. Both now save first and print only after, so a failed save surfaces only the real exception. Incremental per-epoch checkpoints already wrote to runs/<id>/ throughout the loop, so the directory existed when the print fired; the message ordering is the only thing that changed.
• NNParams / NNTransformerParams fail fast on non-positive architectural dimensions. Several __post_init__ gaps let a degenerate config construct silently and misbehave far downstream: on NNTransformerParams, d_model=0 passed the 0 % n_heads == 0 divisibility check and then zeroed head_dim (a zero attention-scale divisor) and the FFN width, while n_layers<=0 built an attention-free embed→norm→head model with no error at all — both producing a stable-looking run.id for a silently-wrong model. The base NNParams likewise validated none of input_dim / output_dim / hidden_dims / dropout_prob. NNParams.__post_init__ now requires input_dim > 0, output_dim > 0, every hidden_dims entry > 0, and 0.0 <= dropout_prob <= 1.0; NNTransformerParams.__post_init__ additionally requires vocab_size / n_layers / d_model / max_seq_len / ffn_mult all > 0 (with d_model checked before the n_heads divisibility test so the zero case can't mask itself) and 0.0 <= attn_dropout <= 1.0 / 0.0 <= resid_dropout <= 1.0 so all three dropout knobs validate together (an out-of-range attn_dropout/resid_dropout previously surfaced only at the first training forward). Same construction-time fail-fast convention already used for n_heads, accumulate_grad_batches, and the dataset classes; no field touches state(), so every valid config keeps its run.id.
• NNSchedulerParams validates its numeric fields. The dataclass had no __post_init__ at all, so factor<=0, negative min_lr / threshold / patience / cooldown, and non-positive variant knobs (step_size / T_max / max_lr / total_steps / warmup_steps) constructed fine and only surfaced — if ever — deep inside the relevant torch scheduler constructor. It now raises a clear ValueError at construction time (factor > 0, the non-negative bounds, and > 0 for any present variant knob). Validation only rejects invalid configs — it never mutates state() — so a run.id never shifts for any valid config (the optional variant knobs still omit themselves at their defaults; the required fields are emitted unchanged as before).
• NNTrainParams / NNTrainerParams fail fast on n_epochs < 1. n_epochs drives range(params.n_epochs) in both train loops, so 0/negative made training a silent no-op (empty idps, a degenerate saved run, no BEST checkpoint) after printing the run-config table, rather than erroring. Both now raise a ValueError at construction (NNTrainParams gains its first __post_init__; NNTrainerParams extends its existing one), symmetric across the two train paths.
• ViTNN / ViTBlock / JEPAPredictor / QATLifecycleCallback fail fast on non-positive dimensions. These public constructors take raw numeric kwargs that bypass the params dataclasses, so they carried the same silent-degenerate footgun the params guards close: n_layers=0 built an attention-free model, ffn_mult=0 a zero-width SwiGLU FFN, d_model=0 masked itself through the n_heads divisibility check (and image_size negative through the patch_size one), and a non-positive QAT groupsize silently mis-quantized (or crashed cryptically at groupsize=0) deep inside prepare(). ViTNN, ViTBlock, and JEPAPredictor now validate every architectural dimension > 0 at the top of __init__ (positive-dim checks ordered before the divisibility tests so the zero case can't mask itself), and ViTNN/ViTBlock additionally bound their attn_dropout/resid_dropout to [0, 1]; QAT routes through the shared _build_quantizer chokepoint which now requires groupsize > 0. DiffusionMLP (which already validated input_dim/time_embed_dim) now also requires every hidden_dims entry > 0, matching NNParams. None of these touch any serialized state.
• NNGraphDataset neighbor sampling is now reproducible (the last dataset holdout). It was the only dataset whose loaders both shuffle (train_loader) and spawn worker processes (n_workers=4 by default) yet threaded neither a generator nor a worker_init_fn — so neighbor sampling was non-deterministic across runs even after nnx.set_seed(...), unlike the random_split-based siblings (NNDataset / NNTabularDataset / NNPreferenceDataset) which already exposed seed. It now accepts seed: Optional[int] = None and threads a seeded generator plus the shared dataloader_worker_init_fn into all three NeighborLoaders. seed=None keeps the prior behavior by falling back to torch.default_generator. First behavioral test coverage for the class accompanies the fix.
• MoELinear fails fast on non-positive in_features / out_features. The public MoELinear constructor validated num_experts / top_k but not its feature dims, so in_features=0 or out_features=0 built a silently-degenerate routing layer — PyTorch emits only a UserWarning on a zero-dim nn.Linear. It now raises ValueError for either, completing the positive-dimension fail-fast convention the params dataclasses and the other public net constructors (ViTNN / ViTBlock / JEPAPredictor / DiffusionMLP) already enforce. No serialized state is affected.
• I-JEPA jepa_train_step_factory zeros the predictor's gradients explicitly. The step called model.net.zero_grad() but not predictor.zero_grad(). The recommended path (predictor registered as a model.net submodule) was already covered, but in the docstring's alternate path — the predictor's parameters added to the optimizer directly, outside model.net — its gradients accumulated across steps while the optimizer still stepped them. The step now zeros the predictor explicitly (a no-op in the submodule path).
• PEFT adapter loaders restore GPU-saved adapters on a CPU-only machine. load_lora_weights / load_ia3_weights / load_prefix_weights / load_prompt_weights resolve a checkpoint path through the shared _resolve_source_to_state_dict helper, which used weights_only=True but no map_location — so an adapter trained and torch.save-d on CUDA raised when loaded on a CPU-only host. It now loads with map_location="cpu" (the downstream load_state_dict copies the tensors onto the target module's device), matching nnx.finetune.load_pretrained.
• nnx.viz.gradient_flow accepts an NNModel. It was the only nnx.viz function typed nn.Module-only; the other five (summary / weight_histogram / activation_map / attribute / netron_export) accept an NNModel and unwrap to .net, so passing an NNModel (as callers do for the others) raised AttributeError. It now unwraps .net like its siblings.
• NNParamGroupSpec fails fast on a non-positive lr / lr_multiplier or a negative weight_decay. A per-group lr or lr_multiplier of 0 silently zeroed the group's effective learning rate (its parameters never updated), and a negative value ran gradient ascent on that group — the same lr/scale footgun NNOptimParams and NNSchedulerParams already guard. __post_init__ now requires lr > 0, lr_multiplier > 0, and weight_decay >= 0 (0 still disables weight decay for the group). No serialized state is affected.
• Multi-optimizer Trainer rejects optimizers that don't scope their parameters. With more than one optimizer, an NNOptimParams with param_groups=None routes to all net.parameters(), so two such optimizers silently double-stepped every parameter — violating the disjoint-partition contract the multi-optim Trainer exists to provide. Trainer.train now fails fast with an actionable error naming the unscoped optimizers; the NNTrainerParams object stays constructible (serialization / builder round-trips are unaffected).
• NNOptimParams fails fast on a negative max_lr or weight_decay. Validated at params-construction time (the NNx fail-fast convention) rather than surfacing deep in optimizer construction or — for a negative weight_decay, which some PyTorch optimizers accept silently — as gradual weight growth during training. The per-group NNParamGroupSpec already rejected a negative weight_decay; the top-level fields now match. __post_init__ requires max_lr >= 0 and weight_decay >= 0. max_lr == 0 stays valid as an explicit "freeze updates" value and weight_decay == 0 disables decay; no serialized state is affected.

[0.2.0] — 2026-06-13 — Expansion megamerge + Month-1 cluster + overnight-maintenance + Builder rollout + PyPI rename

Spans the PR #29 megamerge (20 sub-projects) + PRs #30–#41 (Month-1 cluster + first overnight-maintenance pass) + PRs #42–#46, #48 (security fix + 5-PR Builder-pattern rollout + LogitsChain) + PRs #47, #49 (two-step PyPI distribution rename nnx → nnx-pytorch → thekaveh-nnx) + PR #50 (post-#49 overnight-maintenance — Builder correctness backfill) + PR #51 (post-#50 overnight-maintenance — 5 correctness fixes + docs sync) + PR #52 (post-#51 overnight-maintenance — phase-tag refactor + Builder boundary + 4 docs/test fixes) + PR #53 (post-PR-#52 overnight-maintenance — subpackage __all__ consistency + PEFT source-resolution DRY + idiom + Raises:) + PR #54 (post-#53 overnight-maintenance — reproducibility hardening + ergonomics) + the post-PR-#54 overnight-maintenance pass (the Fixed section directly below). Test suite is 881 tests; 879 pass, 2 skip (the CUDA-gated 2:4 semi-structured sparsity path skips on CPU runners; the network-gated test_pypi_lists_the_current_distribution_name skips on PyPI 404 until the first release ships).

Fixed — post-PR-#54 overnight-maintenance pass

• Transformer params no longer silently downgrade on load. All three deserialization paths (NNRun.load, the NNCheckpoint safetensors reader, hub from_pretrained) called NNParams.from_state(...) unconditionally, so a state written by NNTransformerParams.state() came back as a base NNParams with vocab_size / n_layers / d_model / max_seq_len dropped — the reloaded run re-hashed to a different run.id and rebuilding the net crashed. New NNParams.resolve_from_state(state) dispatches on the transformer keys; every loader funnels through it. Regression tests include an end-to-end NNRun save/load asserting the reloaded run keeps its NNTransformerParams and its original run.id.
• activation=None now survives the state() round-trip. NNParams.state() serialized str(None) == "None", which Activations("None") could never parse back (ValueError on reload). state() now stores a real null and both from_state constructors restore None; a fully absent key still falls back to the legacy LEAKY_RELU default on the transformer path.
• Loaded runs no longer get NaN-filled val_edp on every row. NNIterationDataPoint.from_state checked is not None on flattened CSV cells, but pd.read_csv yields NaN (not None) for cells that were None at save time — so every reloaded idp got a NaN-filled NNEvaluationDataPoint, contradicting the documented only-the-last-idp-of-each-epoch contract and breaking idp.val_edp is not None consumers. NaN cells now map back to None, both for val_edp presence and for the optional loss / error fields.
• Dataset seed=None genuinely falls back to the global torch RNG. All three dataset classes (NNDataset, NNTabularDataset, NNPreferenceDataset) passed a fresh torch.Generator() to random_split when seed=None — but a fresh Generator always carries the same fixed default seed, so unseeded splits were bit-identical across runs and deaf to torch.manual_seed / set_seed, contrary to the documented contract. seed=None now passes torch.default_generator; a regression test pins the contract by asserting torch.manual_seed controls the split.
• KV-cache decoding is positionally correct across sliding-window overflow. On overflow the cache path dropped the oldest cached k/v — but cached entries are RoPE-stamped at the absolute position they were written, so every later token's rotary offset was pinned at max_seq_len - 1, drifting the sampler's logits ~0.1–0.5 vs the no-cache path on a tiny model (the existing token-level parity test passed only by argmax luck). The cache is now rebuilt from the current window on overflow — identical math to the no-cache sliding window, so greedy/seeded parity is exact; the O(T) cache win still applies within the window. New logits-level regression test spies on the raw logits entering the processor chain (greedy's temperature-0 scaling collapses post-chain logits to ±inf one-hots, which is what masked the drift).
• generate(max_new_tokens=0) emits zero tokens on both decode paths. The cache path's prefill unconditionally sampled one token before entering the loop, violating the documented hard cap (no-cache emitted 0, cache emitted 1). Plus first test coverage for the stop= contract.
• train() fails fast on a missing train_loader. Both NNModel.train and Trainer.train accepted params.train_loader=None (the dataclass default), printed the run-details table, then crashed mid-loop with a raw TypeError: 'NoneType' object is not iterable. Both now raise an actionable ValueError at the boundary alongside the existing params is None guards.
• from_pretrained honors strict and rejects unknown kwargs. _from_pretrained computed strict=strict if strict else True — always True — so strict=False was impossible; it is now forwarded (default True, matching the real prior behavior). Unexpected **model_kwargs raise TypeError instead of vanishing silently (the mixin-injected net_params / params config dicts are dropped knowingly), and the docstring's claim that map_location was ignored is corrected (it was always forwarded to safetensors).
• CI + release gates install all 13 declared extras. tensorboard, wandb, and onnx were missing from both workflows' install lines, so the TensorBoardCallback events-file test silently skipped on every CI and release-gate run (the PR #29 failure mode), onnx coverage was only transitively satisfied via onnx-dynamo, and WandbCallback had zero happy-path exercise. release.yml also gains the workflow-level permissions: contents: read floor that ci.yml already had.
• Version bumped to 0.2.0. The dated CHANGELOG [0.1.0] section (2026-05-18 extraction baseline, never tagged or published) already claims 0.1.0, so shipping the first real release as 0.1.0 would have produced two changelog sections claiming one version. The first PyPI release will ship this [Unreleased] content as 0.2.0.
• PrefixTuner KV-cache decode corrected. The patched attention forward cached the prefix-injected K/V, so every cached decode step re-prepended the learned prefix on top of the cached copy (n_prefix duplicate slots per step) and computed the RoPE offset from a length inflated by the prefix — generate(use_cache=True), the default, was silently wrong on prefix-tuned models (cached logits drifted ~2.0 from the full forward). The cache now snapshots real-token K/V before prefix injection; a parity regression test pins per-step logits equality and cache length.
• PEFT loaders report tensors actually loaded. All four load_*_weights returned the source-dict size, but load_state_dict(strict=False) silently drops keys the target doesn't have — loading a LoRA checkpoint into an un-adapted model reported 4 loaded when 0 landed. unexpected_keys are now subtracted at all four sites.
• Surgery primitives thread device (and deepen threads dtype from the right layer). deepen's identity layer and low_rank_factorize's two replacement Linears were constructed without device=, splicing CPU layers into CUDA-resident models (widen already did this correctly). deepen's dtype probe also peeked at parent[idx-1] instead of the Linear that sourced the hidden dim, so a Dropout/norm between the Linear and the ReLU got a float32 layer spliced into a float64 model.
• export_to_safetensors handles tied weights. .contiguous() is a no-op on an already-contiguous storage-sharing tensor (tied embedding/head — TransformerNN's default), so safetensors raised "Some tensors share memory"; the cleaner now clones.
• write_gguf no longer leaks its file handle on error (GGUFWriter lifetime wrapped in try/finally); a tokenizer-parse failure during special-token collection now emits a RuntimeWarning instead of silently emitting a GGUF with no CONTROL-marked tokens.
• Edge-case crashes: Utils.print_tree({}) no longer raises ValueError from max() over zero keys; sinusoidal_time_embed(dim=2) no longer divides by zero (half==1 degenerates to a single unit frequency); the __version__ fallback literals missed by the 0.2.0 bump are corrected.
• Documentation truthfulness pass: the processor-chain order is documented as NNx's own canonical order (temperature deliberately last for the temperature-0 greedy markers) instead of falsely claiming HF equivalence; concepts.md's generate() signature, forward_with_cache name, non-destructive-contract count (ten sites), and finalize_step consumer list corrected; comparison.md's auto-resume cell names both resume kwargs; embeddings.md's LangChain hand-off snippet now wraps the raw index instead of calling load_local on a bare file; README/lm.md KV-cache speedup claims qualified to within max_seq_len.
• GNN training no longer leaks labels across splits. _fwd_pass scored every node in a NeighborLoader subgraph, but only the leading batch_size rows are the batch's seed nodes — the appended sampled neighbors can belong to other splits, so val/test labels entered the training loss and train labels inflated val metrics (maximally so under the default full-split batch size). New GraphNNBase.seed_count exposes the seed-row count; the default train step, evaluate, and predict's loader loop slice to it. Plain full-graph Data batches are unaffected.
• DPO excludes pad positions from response log-probs. NNPreferenceDataset right-pads chosen/rejected responses, and the response log-prob summed every position — pad terms don't cancel between policy/reference or chosen/rejected, so the objective was biased and the gradient also trained the policy to emit pads. dpo_train_step_factory now takes pad_token_id (example 22 and docs pass the dataset's pad id); None preserves the legacy behavior for genuinely unpadded responses.
• LINEAR_WARMUP_DECAY no longer trains the entire first epoch at LR=0. The warmup lambda returned 0/warmup_steps at step 0 and the scheduler steps once per epoch; the ramp is now 1-based.
• Mixup/CutMix respect set_seed. Both factories self-seeded np.random.default_rng() from OS entropy, so λ draws and CutMix boxes differed across identically-seeded runs; the numpy seed is now drawn through the torch RNG.
• EarlyStopping resets per run. _best/_wait persisted across train() calls, so a reused instance compared the new run against the previous run's best and could stop it immediately; state now resets in on_train_begin.
• MoELinear accepts (..., in_features) input like the nn.Linear it documents itself as a drop-in for (3-D sequence batches previously raised a cryptic IndexError; tokens now route independently with leading dims restored), and the MoE step factory clears stale last_aux_loss tensors before each forward so a registered-but-unexercised MoE layer can't trigger "backward through the graph a second time".
• PrefixTuner docstring + deepen ModuleList path device (the pass-2 surgery fix covered Sequential parents only); examples honesty pass: example 01 now learns (labels derive from inputs instead of pure noise), example 19 re-prunes after fine-tuning instead of reporting a silently-regrown 4.3%-sparse network as the 50%-sparse result, example 23 documents the per-generation runs/<id>/ overwrite, example 09 acknowledges possible mode collapse.
• train_contrastive rejects degenerate batching (batch_size=1, or a dataset of < 2 pairs): NT-Xent on a single pair is identically 0.0 with zero gradients, so such configs silently trained nothing while printing 0.0 epoch means; the trailing size-1 batch is dropped only when it would actually have size 1. release.yml asserts the pushed tag matches pyproject's version before the PyPI upload (a typo'd tag would have published the wrong version permanently, caught only post-upload).
• The importorskip class is closed completely: test_checkpoint_safetensors.py, test_hub_mixin.py (hub deps), and test_viz_netron.py (onnx — needed by torch.onnx.export's proto save in every test) gained module-level guards; verified by running the three files with the extras import-blocked (3 clean skips, zero hard failures).
• The two viz test files gained the importorskip guard every other optional-extra file already had — running the shipped suite without the [viz] extra hard-failed 10 tests with raw ImportErrors instead of skipping.
• The sdist ships a complete, runnable test suite — setuptools' legacy default glob included tests/test_*.py but silently dropped conftest.py/tests/__init__.py, leaving distro packagers a half-suite that fails with "fixture not found" (the conftest defines load-bearing fixtures and the macOS OMP guards). A MANIFEST.in now includes the suite whole.
• save_pretrained works on default (tied-embedding) transformers — the Hub writer was the third member of the tied-weights class (after export_to_safetensors and NNCheckpoint.to_file) still missing the .clone(); every default TransformerNN crashed save_pretrained/push_to_hub with "Some tensors share memory". Round-trip restores the tie and greedy-generation parity.
• torch.save of a whole prefix-tuned net round-trips — the MethodType binding pickles by name-lookup on the instance, which a class-level alias now resolves (previously the save succeeded but the load died with AttributeError: a silently unloadable artifact).
• deepcopy of a prefix-tuned net is now fully independent. The patched attention forwards were instance closures, which copy.deepcopy treats as atomic — a copy's attention silently kept reading the ORIGINAL weights and prefix params, corrupting quantize-on-copy, born-again frozen teachers, and surgery copies of prefix-tuned nets. The patch is now a MethodType-bound module-level function with its references stored on the MHA module, so deepcopy rebinds to the copy and re-references through the memo.
• Prompt-tuned generate() survives long generations: the sliding window was sized from net_params.max_seq_len, but a PromptTuner consumes n_prompt_tokens of those slots — generations crashed mid-stream once window + soft prompt exceeded the wrapped model's window. PromptTuner now advertises effective_max_seq_len and generate() honors it.
• NNTokenizerParams.of creates the destination's parent directory (the lm/dpo quickstarts' path="artifacts/tok.json" from a fresh cwd previously failed with a cryptic Rust-side "No such file or directory"); surgery.md's freeze-old-rows recipe corrected twice over and now VERIFIED end-to-end: the expanded embedding must be reattached, and the embedding needs a decay-free param group because Adam applies weight decay inside step(), after gradient hooks.
• write_gguf creates the destination's parent directory (the quickstart's write_gguf(net, tok, "out/model.gguf") from a fresh cwd previously raised FileNotFoundError; the ollama exporter already mkdir'd).
• NNTabularDataset rejects target_col inside feature_cols — that config silently trained the model on its own label (near-perfect val accuracy from the classic feature_cols=list(df.columns) mistake).
• Schedulers reject an explicit total_steps < n_epochs — NNx steps schedulers once per EPOCH (not per batch, the HF habit), so a short total_steps made OneCycle raise mid-train at epoch total_steps+1 (losing that epoch's idps) and LINEAR_WARMUP_DECAY silently train every remaining epoch at LR=0. NNRun.load also rejects an empty/truncated run.yaml with the file path (safe_load → None previously died on a bare AttributeError).
• Boundary hardening (pass 14): a fully-frozen model (freeze('*')) is rejected at train() entry on both loops instead of dying mid-loop with torch's raw does-not-require-grad error; NNTabularDataset rejects NaN cells in modeled columns (NaN targets cast to int64 UNDEFINED — silently class 0 on ARM); malformed-artifact errors in NNRun.load name the file that's actually corrupt (a dropped idps.csv column is no longer blamed on run.yaml).
• NNOptimParams rejects plain dicts in param_groups at construction with a wrap-it TypeError — they previously crashed much later inside state() during NNRun hashing with an opaque AttributeError.
• Dataset construction fail-fast pass: NNDataset validates val_proportion, rejects transform-less PIL samples with an actionable message ("pass ToTensor()"), and handles val_proportion=0.0 (previously DataLoader(batch_size=0) crashed); NNTabularDataset rejects non-contiguous labels (e.g. {0, 5}) at construction instead of letting nunique()-sized models fail much later inside cross-entropy; DiffusionMLP validates time_embed_dim at init; the dataset base class types val/test loaders as Optional to match the documented empty-split contract.
• Modelfile injection guard: export_ollama_modelfile rejects triple-quotes in system/template, whitespace in parameter keys, and newlines/quotes in string parameter values — before the expensive GGUF write — instead of emitting a Modelfile where user strings could terminate blocks early or inject directives.
• Plotting correctness: VisUtils.confusion_matrix pins sklearn's labels= so an absent class no longer shifts every later row/column onto the wrong name; generate_colors no longer gives the first and last class identical hues; multi_line_plot's legend mapping now matches its documented (group_labels, line_labels) contract (the two no-trace legend blocks were swapped against the data-trace hover names).
• lr_finder actually reaches end_lr (the exponent used 1/num_iter, stopping one multiplicative step short of the documented sweep ceiling); load_pretrained raises on remap collisions (prefix/key_map rewriting collapsing two source keys onto one target silently loaded whichever came last); freeze/prune/PEFT glob matching uses fnmatch.fnmatchcase at all 8 sites (plain fnmatch is case-insensitive on Windows, so patterns matched differently across platforms); train_contrastive drops the size-1 trailing batch (NT-Xent on a single pair is exactly 0.0 loss with zero grads, but the step still moved weights on stale Adam momentum and deflated the verbose epoch mean); weight_histogram skips empty tensors as its docstring promised; the freezing docstring example gains its missing wildcard.
• Checkpoint format sniffing no longer misroutes safetensors files whose header length ≡ 128 mod 256. A safetensors file's first byte is the low byte of the u64 header length, which can legitimately be 0x80 — the pickle PROTO opcode — so NNCheckpoint.from_file routed such files to torch.load, which died with a confusing UnpicklingError. Byte 8 (the JSON header's {) now positively identifies safetensors before the pickle check.
• NNCheckpoint.to_file(format="safetensors") handles tied weights — every default TransformerNN (tied tok_embed/lm_head) crashed with "Some tensors share memory"; the writer now clones, and the reload keeps the tie intact.
• runs/best tracking fixes: the pointer repoint is now atomic (temp symlink + os.replace — a crash mid-repoint can no longer leave no pointer, which made the next save claim best unconditionally); the symlink is created with target_is_directory=True (Windows-with-developer-mode created an untraversable file symlink, silently breaking best tracking); and _best_err falls back val→train / error→loss via the shared resolver, so runs whose steps leave .error unset (custom trainer/GAN step functions; the shipped diffusion/SimCLR/DPO factories do populate .error) no longer all score +inf and freeze best on whichever run saved first.
• NNRun(...).save() with the dataclass-default idps=None writes an empty idps.csv instead of raising a bare TypeError; NNTokenizerParams.of saves the tokenizer atomically (temp + rename) so an interrupt can't leave a truncated tokenizer.json that from_state can never load.
• generate() restores train mode on early raises and scopes stop= to the continuation. The eval() switch preceded tokenizer/parameter validation while the restoring try/finally only wrapped the decode loop, so an early ValueError stranded the net in eval mode; and stop strings were searched in the full decoded text, so a prompt already containing one halted after a single token.
• Both train loops reject zero-batch epochs (dataset smaller than batch_size with drop_last=True, or a one-shot iterable exhausted in epoch 0) with an actionable ValueError — previously the first epoch crashed on a bare IndexError, and a later zero-batch epoch would have silently attached its val_edp to the previous epoch's last idp.
• TransformerNN gains GPT-2/LLaMA-style embedding init. nn.Embedding defaults to N(0,1); with the tied LM head the input token's own logit then includes e·e ≈ d_model, so an untrained model started at CE ≈ d_model (123 measured at d_model=128 — worse than the uniform-random baseline ln(vocab) ≈ 5.7) and decoding degenerated into repeating the last prompt token regardless of sampling settings. std=0.02 init on the shared tensor covers the tied head; example 11's generations recover.
• GNN seed-slicing gated on input_id so PyG multi-graph Batch.from_data_list collations (which also carry batch_size = num_graphs) aren't truncated to the graph count; plus predict()-path coverage.
• MoELinear also accepts the unbatched (in_features,) form; example 22 pads with the real <pad> id (1) instead of the default 0 = <unk>, which the DPO mask would otherwise drop from genuine responses.
• Surgery primitives no longer consume the global torch RNG. widen's comment claimed RNG-cleanliness (its duplication choices do use a local generator), but the fresh replacement layers in all four module-building primitives (widen ×2, deepen's identity Linear, expand_embedding, low_rank_factorize ×2) ran their default kaiming/normal init off the default generator before being fully overwritten — so any seeded caller pipeline silently diverged after a surgery call (the same reproducibility class as the examples seed-helper bugs). All six sites now build via torch.nn.utils.skip_init (meta-device construction, zero RNG draws on any device); four regression tests pin torch.get_rng_state() equality across each primitive.
• lr_finder and viz.summary(input_size=) no longer perturb seeded pipelines. The follow-up sweep to the surgery skip_init fix: lr_finder's docstring promised the pre-flight sweep wouldn't disturb "any subsequent reproducibility", but its try/finally restored only weights + mode — the sweep's dropout draws (the helper itself calls .train()) and the DataLoader base-seed draw leaked into the caller's RNG stream, so seed → lr_finder → train produced different weights than seed → train. The snapshot/restore now covers CPU + active-device RNG state in the same try/finally — plus any loader-attached generator= at all four attachment points — DataLoader(generator=), an explicit sampler='s, one wrapped inside an explicit batch_sampler=, and a custom batch sampler owning its generator directly (the official PyTorch reproducibility recipe routes shuffle permutations through these streams, which the global restore can't reach; three follow-up refutation passes each caught an escape). Non-torch generator objects riding the same attribute name (e.g. a numpy Generator on a custom sampler) are skipped instead of crashing the snapshot — they stay caller-owned. With persistent_workers=True the sweep also discards the cached iterator it created on the loader — otherwise the caller's first epoch _reset()s that cache instead of drawing a fresh worker base seed, shifting their batch stream even with every RNG snapshot faithfully restored. viz.summary(input_size=...) (the mid-pipeline probe idiom concepts.md recommends) advanced global RNG via torchinfo's torch.rand dummy-input synthesis even though the values never affect the statistics; the wrapper now restores RNG around that path (input_data= was already clean). Both red-verified with state-equality regression tests.
• Reentrancy + filesystem-robustness pass: PrefixTuner rejects an already-prefix-tuned net — a second tuner silently hijacked the first (the patched forwards read the MHA's _nnx_prefix_tuner ref, which the second tuner overwrote, so tuner-1's parameters stopped receiving gradients while its forward injected tuner-2's prefixes; training tuner-1 became a silent no-op). runs/best symlinks now use the sibling run-dir basename as the target — the raw run_path target resolved relative to the symlink's own directory, so a relative root= produced a dangling link from birth and every save took the repoint-unconditionally branch (best tracked the most recent run, not the best); the basename target also survives relocating the runs root. A zero-byte idps.csv (external truncation) now raises the per-file malformed idps.csv at <path> error instead of pandas' context-free EmptyDataError.
• Coverage backfill: NNDataset behavioral tests + train-step-factory convention test. NNDataset (the torchvision facade) previously had only an hasattr smoke assertion — tests/test_nn_dataset.py now exercises split arithmetic, full-split batch resolution, dim/class introspection, and both halves of the seed contract against an in-memory VisionDataset. tests/test_imports.py gains an AST-based scan asserting every *_train_step_factory in the package is re-exported at top-level nnx.* and listed in nnx.__all__ (the convention PR #54 had to repair by hand). Plus shared-helper consolidation: the classification step epilogue (4 copy-pasted sites) and the Hinton softened-KL block (2 sites) now live in one place each, and Trainer delegates its checkpoint/tqdm tails to the NNModel implementations instead of re-implementing them.

Changed — PyPI distribution name

• PyPI publication-verification tests + post-publish CI smoke job (PR #49). Three pytest layers in the new tests/test_pypi_publication.py: (1) pyproject.toml [project] version must equal nnx.__version__; (2) importlib.metadata.version(<dist name>) must succeed and equal nnx.__version__ (regression test catching the _version("nnx") lookup that PR #47 and the first cut of PR #49 both left stale — src/nnx/__init__.py now uses _version("thekaveh-nnx")); (3) a network-gated check that PyPI lists the current dist name (skips on 404 when not-yet-published, skips on network error, skips on NNX_SKIP_PYPI_TESTS=1 for offline contributors). Plus a new verify-published job in .github/workflows/release.yml that runs after pypa/gh-action-pypi-publish succeeds: spins a fresh venv, retries pip install thekaveh-nnx==<tagged-version> up to 5× with 30s backoff (handles PyPI CDN propagation lag), then imports nnx and asserts __version__ matches the tag. This is the contract the user really wants from "is the latest code actually on PyPI" — runs only on release events, asserts the exact uploaded version is reachable. Test count 745 → 748 (3 new tests; the network-gated one skips locally when offline / not-yet-published).
• Renamed package on PyPI: nnx → thekaveh-nnx (final name chosen via PR #49, after an intermediate pass through nnx-pytorch in PR #47). The nnx name on PyPI is owned by an unrelated, abandoned JAX library by cgarciae (nnx==0.0.8, last release 2023-07-12), so pip install nnx would silently install the wrong package — a JAX library with no overlap with this PyTorch toolkit. PR #47 first renamed to nnx-pytorch (generic-suffix convention); PR #49 then renamed again to thekaveh-nnx (username-prefix convention, matching the GitHub-handle namespace @thekaveh/). The import path stays import nnx through both renames — only the install name changes (same pattern as Pillow/PIL or scikit-learn/sklearn). No wheel was ever published under either intermediate name, so this is a one-time correction sequence, not a deprecation chain. All install command references throughout this CHANGELOG, README, docs, examples, and source-code docstrings have been retroactively updated to the final name thekaveh-nnx.

Added — Builder pattern rollout

• NNSchedulerParams.builder() — variant-gated construction (PR #43). Classic-GoF Builder reachable via NNSchedulerParams.builder() returning NNSchedulerParamsBuilder (re-exported at top level as nnx.NNSchedulerParamsBuilder). Five variant methods — reduce_on_plateau, step, cosine_annealing, one_cycle, linear_warmup_decay — each set kind plus the variant-specific fields and leave everything else at the dataclass defaults, so the omit-when-default state() invariant is preserved automatically (Builder only forwards user-touched fields). Existing direct-kwarg NNSchedulerParams(...) ctor is untouched; the Builder is purely additive. 10 new tests cover happy-path per variant, state() round-trip, the "last variant wins" overwrite contract, build-without-variant rejection, the omit-when-default invariant, and a top-level re-export smoke check (9 in tests/test_nn_scheduler_params_builder.py + 1 in tests/test_imports.py).
• NNOptimParams.builder() — variant-gated optimizer config (PR #44). Classic-GoF Builder reachable via NNOptimParams.builder() returning NNOptimParamsBuilder (re-exported at top level as nnx.NNOptimParamsBuilder). Four optimizer variants — adam, adam_amsgrad, sgd, sgd_nesterov — plus three optional chained modifiers — grad_clip(norm), accumulate_grad(batches), param_groups(specs). The Adam variants take a PyTorch-native betas: tuple[float, float] kwarg which the Builder maps onto the dataclass's momentum field — the field name stays momentum on-disk for back-compat, but the Builder API uses the PyTorch spelling. SGD variants keep the float momentum= kwarg. Existing direct-kwarg NNOptimParams(name=Optims.ADAM, ...) ctor is untouched. Omit-when-default state() invariant preserved automatically (Builder only forwards user-touched fields). 9 new tests in tests/test_nn_optim_params_builder.py (8) + tests/test_imports.py (1) cover happy-path per variant, the betas-→-momentum field mapping, grad_clip / accumulate_grad / param_groups chaining + state-round-trip, and the top-level re-export.
• NNTransformerParams.builder() — LM-path config (PR #45). Classic-GoF Builder reachable via NNTransformerParams.builder() returning NNTransformerParamsBuilder (re-exported at top level as nnx.NNTransformerParamsBuilder). Six fluent methods — vocab(size) (sets both input_dim + output_dim), layers(n, heads, d_model) (enforces d_model % heads == 0 at call-time), ffn(mult), context(max_seq_len, rope_base), dropout(attn, resid), tied_embeddings(value). The Builder hides the dead parent-NNParams kwargs (hidden_dims, activation, dropout_prob) that TransformerNN doesn't use, so the LM-path API reads in LM-path terms. Existing direct-kwarg NNTransformerParams(...) ctor is untouched. Omit-when-default state() invariant preserved automatically. 8 new tests in tests/test_nn_transformer_params_builder.py (7) + tests/test_imports.py (1) cover happy-path with hidden-parent-kwargs, omit-when-default, the d_model % heads == 0 call-time validation, all four chained modifiers (ffn / context+rope / dropout / tied_embeddings), and the top-level re-export.
• NNTrainerParams.builder() — composite multi-optim Builder (PR #46). Classic-GoF Builder reachable via NNTrainerParams.builder() returning NNTrainerParamsBuilder (re-exported at top level as nnx.NNTrainerParamsBuilder). Composes the Plan 1+2 Builders: .optimizer(name, NNOptimParams) and .scheduler(name, NNSchedulerParams) register each entry under a user-chosen name. .build() enforces schedulers.keys() ⊆ optims.keys() at the Builder boundary with an actionable error naming the unknown scheduler key + listing the known optim names — today NNTrainerParams.__post_init__ enforces the same invariant only after the dataclass ctor runs, so the Builder surfaces it one stack-frame earlier. Plus chained modifiers .n_epochs, .seed, .save_phase_checkpoints, .train_loader, .val_loader, .extra_metrics. The Example 09 GAN recipe (two parallel optimizers with matching schedulers) collapses from ~30 lines to a single fluent expression. Purely additive; existing direct-kwarg NNTrainerParams(optims={...}, schedulers={...}) ctor untouched; on-disk state() / from_state() round-trip unchanged. Omit-when-default invariant preserved (Builder only forwards the user-touched fields, and the empty _schedulers dict isn't passed to the dataclass kwargs when it's untouched, so the field stays at the default_factory=dict default). 10 new tests in tests/test_nn_trainer_params_builder.py (9) + tests/test_imports.py (1) cover happy-path (minimal + GAN recipe), the omit-when-default invariant, the .build()-time key-subset rejection, the actionable error-message contract (names the typo + lists known optim names), build-without-any-optim, chained .train_loader / .seed / .extra_metrics, the full state() round-trip through the inner Plan-1/2 Builders, and the top-level re-export.
• LogitsChain + LogitsChain.builder() — power-user LM decoding (PR #48). New top-level types nnx.LogitsChain and nnx.LogitsChainBuilder. LogitsChain is a thin frozen-dataclass wrapper around list[LogitsProcessor] with an .apply(logits, token_history) method. The Builder chains .repetition_penalty(p), .top_k(k), .top_p(p), .temperature(t), .custom(processor) in any order; .build() sorts the canonical processors into NNx's canonical order (RepetitionPenalty → TopKFilter → TopPFilter → TemperatureScaling; temperature deliberately last — see the post-#54 truthfulness fix in the Fixed section above), with custom processors appended after. GenerativeNNModel.generate(...) gains a new optional logits_chain: Optional[LogitsChain] = None kwarg — when provided, the inline chain construction from temperature / top_k / top_p / repetition_penalty kwargs is skipped. When None (the default), behavior is unchanged. Purely additive; existing generate(temperature=0.8, top_k=50, ...) callers continue to work. 8 new tests: 6 pure-torch in tests/test_logits_chain.py (empty-chain pass-through, single-processor, canonical-order enforcement, custom-processor append-after-canonical, multiple-calls-to-same-method overwrite contract, chain.apply equivalent to apply_chain direct call) + 1 import smoke + 1 integration test in tests/test_generative_nn_model.py (gated on the lm extra) that verifies generate(logits_chain=..., seed=...) is reproducible across two calls.

Added — Month-1 cluster (PRs #32–#37)

• PEP 561 py.typed marker (PR #32). Adds an empty src/nnx/py.typed and a [tool.setuptools.package-data] entry so the wheel ships it. Declares NNx as type-checked for downstream pyright / mypy consumers — they now see the existing public-surface annotations (NNModel, NNParams, callbacks, enums) instead of treating every symbol as Any. No NNx-side typing change; the gain is entirely downstream.
• docs/comparison.md page (PR #33). "NNx vs Lightning / HF / fastai / Composer" honest, scope-explicit comparison: quick decision matrix, landscape map, capability-axis tables (training loop / distributed / PEFT / generation / diffusion / GNN / surgery / observability / Hub), when-to-use-what rule-of-thumb, and a "what NNx doesn't ship" call-out. Wired into mkdocs.yml nav and linked from README §5.1.
• nnx.viz.gradient_flow(model) (PR #34). Per-layer L2 gradient-norm bar chart for training-loop diagnostics. Call after loss.backward() and before optimizer.zero_grad(); returns a Plotly Figure. Frozen params and params with grad is None are skipped. Raises ValueError with a helpful message when no parameter has a populated gradient. Six new tests; doc paragraphs in concepts.md §12 (now "Six primitives") and api.md.
• nnx.lr_finder(model, train_loader, *, loss_fn, ...) -> LRFinderResult (PR #35). fastai-style exponential LR sweep (1e-7 → 10.0 over 100 iters by default) returning the suggested one-cycle max_lr via the Smith (2017) steepest-descent heuristic on EMA-smoothed loss, plus a Plotly figure of loss vs log(LR). Non-destructive: model state and training-mode are snapshotted before the sweep and restored on exit. Early-exits on divergence (loss > 4× min observed). Nine new tests covering return type, field shape, suggested_lr in range, weight restoration, three invalid-argument paths, and log-axis figure. Docs: concepts.md §13.1 + api.md.
• NNRun._repr_html_() (PR #36). Jupyter rich-display: when an NNRun is the last expression in a cell, it renders a config table (run.id, net, device, loss, dims, dropout, activation, n_epochs, optim + max_lr) plus a Plotly per-epoch metric chart (train_loss / val_loss / train_err / val_err — val traces appear only when validation data was present). Plotly is lazy-imported inside the chart helper so non-Jupyter cost is zero. Falls back to config-table-only when self.idps is None / empty. Epoch boundaries detected by idp.epoch_idx so train-only runs render correctly. Four new tests.
• Six missing megamerge example scripts (PR #37). Closes PR #31's largest deferred item. New files: 19_prune_mnist.py (magnitude pruning at 50% sparsity + brief fine-tune), 20_surgery_resnet.py (low-rank factorize a Linear at rank=8 + refinement), 21_viz_attribute_xai.py (Captum attribution across 4 methods), 22_dpo_tinystories.py (DPO with a deepcopied frozen reference policy on synthetic preference triples; before/after log-prob comparison shows Δ(chosen − rejected) > 0), 23_born_again_distillation.py (iterated self-distillation across G=3 generations), 24_feature_kd.py (FitNets-style feature distillation with one paired teacher → student layer). Each is CPU-runnable in under 2 minutes. examples/README.md catalog updated with four new sub-sections.

Fixed — Month-1 cluster follow-ups

Listed roughly chronologically by ship date; maintenance-pass follow-ups are grouped where their cluster landed rather than strictly interleaved.

• NNRun.load / NNCheckpoint.load reject path-traversal run ids (post-PR-#41 overnight-maintenance pass). Both NNRun.load(id=...) and NNCheckpoint.load(run=...) accept a public string identifier and join it directly into a path under runs/<id>/.... Internal callers always pass the md5 hex of NNRun.state() (32 hex chars; always safe), but a public caller passing "../../etc/passwd" would have escaped the runs root: os.path.join("runs", "../../etc/passwd") resolves to etc/passwd. Sensitive files sitting next to the working directory would be readable on the next .load(...) call. Mirrors the spirit of PR #40's yaml.safe_load fix: the file is normally application-written, but defense-in-depth says never trust the input to escape the API boundary. Added _validate_run_id() in src/nnx/nn/params/nn_run.py that rejects path separators, .., ., embedded nulls, and empty / non-string ids; called from NNRun.load and from _checkpoint_path (the single path-construction site NNCheckpoint.save / .load / .load_optimizer_state all funnel through). New regression tests test_r3_nnrun_load_rejects_path_traversal_run_ids (8 traversal-shaped inputs) and test_r3_nnrun_load_accepts_normal_md5_id (happy-path md5 hex still passes).
• artifacts/ directory ignored (post-PR-#41 overnight-maintenance pass). Three examples write to top-level artifacts/ subdirectories (11_tinystories_lm.py → artifacts/tinystories_lm/, 17_export_transformer_to_gguf.py → artifacts/lm_export/, 18_publish_to_ollama.py → artifacts/ollama_bundle/). Without ignoring the directory, anyone who ran one of these examples ended up with .gguf binaries / tokenizer.json / Modelfile showing as untracked in git status — one careless git add -A away from sweeping them into the repo. Added artifacts/ to .gitignore under the existing "nnx training + example artifacts" block (retitled from "training artifacts" to reflect the expanded scope) alongside runs/ and tb_logs/. Comment block names the three examples that write there so a future maintainer can trace the entry to its callers.
• pyproject.toml Pygments pin (PR #38) — Pygments 2.20.0 broke pymdownx.highlight (a filename=None propagates into HtmlFormatter.__init__ and crashes with AttributeError: 'NoneType' object has no attribute 'replace'), making mkdocs build --strict fail on every docs page containing a fenced Python code block or a mkdocstrings class source block. Pinned Pygments<2.20 in the [docs] extra; CI install line picks it up automatically.
• Examples 20 + 24 seed-state bug (PR #38) — torch.manual_seed(0) inside _make_data() silently overrode the caller's set_seed(42) from main(). PR #37's review caught the same bug in examples 19 / 21 / 23 and fixed those; the fix was missed for 20 / 24. Now removed; examples still complete end-to-end on CPU.
• nnx.lr_finder correctness pass (PR #39). Three correctness issues deferred from PR #38's audit, all fixed: (1) smoothed-min divergence guard — early-exit now compares an EMA-smoothed loss against an EMA-smoothed running minimum (matches fastai's lr_find), so an anomalously low first-batch loss no longer ends the sweep before the descent region is reached; (2) short-sweep fallback — when fewer than 5 points are recorded, suggested LR is the LR at the minimum observed loss, not start_lr (which would have been the worst possible max_lr); (3) monotonically-rising-loss fallback — when no descent region exists at all, the slope-based heuristic falls back to the min-observed-loss LR rather than returning the steepest-positive-slope index. Also fixes a degenerate empty-trace add_vline crash on extremely short sweeps. NNRun._repr_html_ no longer accepts a stale idps arg. Six new regression tests in tests/test_lr_finder.py.
• nnx.interop attribute-access regression (post-#39 maintenance pass). README §1.2 advertises nnx.interop.write_gguf(...) and nnx.interop.export_ollama_modelfile as the post-import nnx access pattern, matching every other subpackage (nnx.diffusion, nnx.peft, nnx.quantize, …). The package-level src/nnx/__init__.py was missing interop from its eager-bind list, so a plain import nnx left hasattr(nnx, "interop") False and the README's documented call shape raised AttributeError until callers added an explicit import nnx.interop. Fixed by adding interop to the from . import … line (the interop subpackage's own docstring already promises that the plain import is safe when gguf isn't installed, since the dep is lazy-imported inside write_gguf / export_ollama_modelfile). New regression test test_subpackages_attribute_accessible_after_plain_import asserts hasattr(nnx, name) for every advertised subpackage so future drift is caught loudly.
• NNRun.save / NNRun.load use safe YAML APIs (post-#39 maintenance pass). NNRun.load previously used yaml.load(f, Loader=yaml.FullLoader), which under PyYAML>=5 no longer permits arbitrary-code tags by default but still allows instantiation of arbitrary Python objects via tags like !!python/object/.... A tampered or attacker-supplied runs/<id>/run.yaml could escalate a filesystem-write primitive into arbitrary-object construction at load time. Switched to yaml.safe_load (matching the long-standing safe_dump / safe_load pair the sibling metadata.yaml already used — see seeding.py's "yaml.safe_load-compatible" comment). Also tightened NNRun.save to yaml.safe_dump, so any future state() change that smuggles in a non-primitive type fails loudly at write-time instead of producing a run.yaml that no longer round-trips. State is a plain dict of primitive types, so the change is a behavior-preserving drop-in; existing NNRun.save / NNRun.load round-trip tests across the suite continue to pass. New regression test test_r3_nnrun_load_rejects_python_object_tags asserts a poisoned !!python/name:os.system tag raises yaml.YAMLError.
• Example 16 seed-state bug (post-#39 maintenance pass) — same anti-pattern as 19–24: _make_synthetic_loader in examples/16_ijepa_cifar10.py opened with torch.manual_seed(0), silently overriding the caller's set_seed(0) from main(). Today both values happen to be 0, so the override is observable only when a future maintainer edits main() to choose a different seed and the dataset stubbornly stays the same. Removed; added the same "no torch.manual_seed here — caller does it" comment that examples 20 / 24 already carry, so the contract is visible at the helper site.
• nnx.embeddings.pair_collate is now a public re-export (post-#39 maintenance pass). docs/embeddings.md documents pair_collate twice as the recommended DataLoader.collate_fn for ContrastiveTextDataset, and the symbol itself has no leading underscore — so it's semantically public. But the nnx.embeddings subpackage's __init__.py omitted it from both from .contrastive_trainer import … and __all__, so users following the docs had to reach into the implementation path (from nnx.embeddings.contrastive_trainer import pair_collate) — a fragile coupling that would break the first time the trainer module is restructured. Added to __all__ alongside the rest of the public surface and updated the test import to use the public path (from nnx.embeddings import pair_collate). The intentional private helper _is_sentence_transformer keeps its submodule import path, with an explanatory comment noting the asymmetry.
• Cross-platform text I/O is now utf-8-explicit (post-#39 maintenance pass). Five open(...) text-mode call sites — three in src/nnx/nn/params/nn_run.py (atomic write, POINTER.txt read, run.yaml safe_load) and two in src/nnx/nn/nn_model.py (_HUB_CONFIG_FILENAME write and read on the Hub round-trip path) — relied on Python's locale-default text encoding. On Linux / macOS that's utf-8 today; on Windows pre-PEP-686 (so anything before Python 3.15 on the default config) it's the locale-specific code page (cp1252 in many cases), which would silently mis-encode a unicode tokenizer path or a non-ASCII run id round-tripped through state(). Pinned all five to encoding="utf-8" with a one-line rationale at the most central site. No behavior change on Linux / macOS; closes the Windows-locale corruption window.
• LICENSE copyright year + holder name (post-#39 maintenance pass). The MIT license header read Copyright (c) 2023 Kaveh, predating the repo (first commit 2026-05-18) and using a first-name-only holder string. Updated to Copyright (c) 2026 Kaveh Razavi to match the actual repo-init year and the pyproject.toml [project] authors entry. No content / license-terms change.
• pyproject.toml PyPI keywords catch up to the megamerge subsystems (post-#39 maintenance pass). The [project] keywords list still reflected the pre-megamerge thekaveh/ml-extraction scope (pytorch, deep-learning, graph-neural-networks, training, checkpointing, …) and missed every major subsystem PRs #29–#37 shipped as a first-class entry point: quantization, language modeling, LoRA / PEFT, diffusion, knowledge distillation, embeddings / RAG, pruning, model surgery, ONNX export, GGUF interop. Added each as a keyword so PyPI search surfaces NNx for users looking for any of those topics. Verified the resulting Keywords: header in the installed wheel metadata.
• Inference helpers (predict, evaluate, generate, sample, embed_texts) now restore model.training instead of stranding it in .eval() (post-PR-#40 overnight-maintenance pass). Five inference-shaped helpers across NNModel.predict, NNModel.evaluate, GenerativeNNModel.generate, nnx.diffusion.sample, and nnx.embeddings.embed_texts all switched the underlying nn.Module to .eval() mode (needed for correct BatchNorm / Dropout semantics during inference) but never restored the caller's prior training-mode state. The codebase already had two non-destructive precedents (nnx.viz.activation_map and the post-PR-#39 nnx.lr_finder); the inference helpers were the odd ones out. Common train → evaluate → train-more / train → predict → train-more loops silently disabled Dropout masking and BatchNorm running-stats updates on the next training step unless the caller remembered to call model.net.train() themselves. Wrapped each helper's body in try: ... finally: model.net.train() if was_training else None. New tests/test_inference_helpers_preserve_training_mode.py ships 9 regression tests covering both train→inference→train AND eval→inference→eval round-trips per helper (paired *_restores_* / *_preserves_eval_mode_caller cases for predict / evaluate / diffusion.sample / embed_texts), plus a predict() exception path that verifies the finally restore fires even when the body raises. GenerativeNNModel.generate's equivalent round-trip is in tests/test_generative_nn_model.py::test_generate_restores_training_mode_after_call because that file already carries the pytest.importorskip("tokenizers") guard needed for the LM path.
• nnx.lr_finder non-destructive contract held only on the happy path (post-PR-#40 overnight-maintenance pass). The docstring promises "the model's initial weights are snapshotted before the sweep starts and restored on exit" — but the restoration ran after the for-loop, not inside a try/finally. Any exception during the sweep (user-supplied loss_fn raising, model(X) crashing on a malformed batch, loss.backward() failing on a NaN gradient) skipped the restore and left the caller's model permanently mutated: weights stuck at whatever the optimizer's mid-sweep updates produced, and model.training left True even if the caller passed an eval()-mode model in. Wrapped the loop body in try: ... finally: model.load_state_dict(initial_state); model.train(was_training). New regression test test_lr_finder_restores_state_after_exception_in_loss_fn passes an eval-mode model + a flaky loss_fn that raises on iteration 4, then asserts both the weights and the training-mode flag are restored despite the mid-sweep crash. Empirically verified against the pre-fix code path (without try/finally, both invariants fail).
• nnx.generation.sample_next_token NaN guard (post-PR-#40 overnight-maintenance pass). The degenerate-probability fallback in the LM sampler used probs.sum().item() == 0.0 to detect rows that couldn't drive torch.multinomial, falling back to argmax(logits) instead. The check is wrong for one residual case: if upstream produces NaN logits (e.g. a divergent KV-cache decoding step), softmax(NaN) propagates to all-NaN probs and probs.sum() is NaN — and NaN.item() == 0.0 is False, so the guard misses it. The user-visible failure was RuntimeError: probability tensor contains either inf, nan or element < 0 from torch.multinomial. (The two prior guards in the same function — +inf early-return and torch.isinf(logits).all() argmax fallback — don't cover this case: isinf(NaN) is False, so a row mixing NaN with finite values reaches softmax intact.) Switched the check to not torch.isfinite(total) or total.item() <= 0.0, which catches NaN, +/-inf, and zero/negative sums. New tests/test_generation_sampling.py ships 7 unit tests (+inf greedy short-circuit, all--inf fallback, NaN-in-logits regression, normal-draw shape, generator reproducibility, bad-shape rejection, underflow edge); the file is shape-test-only so it runs on every CI matrix row regardless of the lm extra.
• NNTransformerParamsBuilder activation default matches the parent NNParams default (PR #50). NNTransformerParamsBuilder.build() hardcoded activation=Activations.RELU, but the parent NNParams default is Activations.LEAKY_RELU. The two construction paths — Builder vs direct-kwarg ctor — therefore produced different state() dicts and different run.id hashes for what users would call "the same config", silently breaking the omit-when-default invariant for everyone using the LM-path Builder. Caught by the post-#49 overnight-maintenance Pass 1 audit; corrected to Activations.LEAKY_RELU so Builder-built transformers hash identically to direct-kwarg ones. New regression test test_builder_state_equals_direct_ctor_with_parent_defaults asserts the two construction paths produce equal dataclasses AND equal state() dicts. Migration note: any prior run.id hash that came from NNTransformerParams.builder()...build() (rather than direct-kwarg NNTransformerParams(...)) will shift to a new id; the on-disk run.yaml continues to load correctly via NNTransformerParams.from_state(...).
• NNOptimParamsBuilder modifier-before-variant chains no longer silently drop the modifier (PR #50). NNOptimParamsBuilder.adam() / .adam_amsgrad() / .sgd() / .sgd_nesterov() did self._fields = {...} (full replace), which silently wiped any prior .grad_clip(N) / .accumulate_grad(N) / .param_groups(...) call. The documented chain order was variant-first-then-modifiers, but the fluent API didn't enforce it — modifier-before-variant chains compiled successfully and silently dropped the modifier. Refactored via a _set_variant() helper that drops only the variant-keyed fields (name / max_lr / momentum / weight_decay) before applying the new variant, so any modifier-set keys (grad_clip_norm / accumulate_grad_batches / param_groups) survive. Last-variant-wins (the existing documented contract) is preserved. New regression tests cover modifier-before-variant survival across all four variant methods.
• Six stale pip install 'nnx[X]' references swept to 'thekaveh-nnx[X]' (PR #50). PR #49's CHANGELOG promised the install-name rename was complete across all docs/examples/source, but the post-#49 audit caught six stragglers — three example module docstrings and three src/nnx/interop/* lazy-import error strings — that still showed the squatted bare nnx name. Updated; closes the gap PR #49 promised.
• NNTransformerParamsBuilder.tied_embeddings(True) / NNTrainerParamsBuilder.save_phase_checkpoints(True) (PR #50). Both setters originally gated with if value is True/False: to avoid storing the dataclass default — but that broke the documented "last call wins" contract: a prior .tied_embeddings(False) followed by .tied_embeddings(True) left the dataclass at False because the True call was a silent no-op. Fixed by storing unconditionally; the dataclass's omit-when-default state() handles run.id stability automatically (the convention every other Builder already followed). Regression tests test_builder_tied_embeddings_true_after_false_overrides_to_true and the analogous save_phase_checkpoints test added.
• env_snapshot()['nnx'] now queries the renamed distribution (post-PR-#50 overnight-maintenance pass). src/nnx/seeding.py:_nnx_version() still called importlib.metadata.version("nnx") after the PR #49 PyPI rename to thekaveh-nnx. On any clean install of the renamed package the lookup raised PackageNotFoundError, was swallowed by the broad except Exception: block, and metadata.yaml silently recorded nnx: null. metadata.yaml's job is reproducibility — a silent None defeats the whole purpose of the snapshot for the most useful field. Mirrors the lookup nnx.__version__ already uses (src/nnx/__init__.py). New regression test test_v3_env_snapshot_nnx_version_is_resolvable asserts snap['nnx'] is not None and snap['nnx'] == nnx.__version__ so future drift between the two version lookups is caught loudly.
• NNModel.to_onnx now restores model.net.training instead of stranding the caller in .eval() (post-PR-#50 overnight-maintenance pass). The post-PR-#40 non-destructive-helper pass fixed five inference helpers (predict / evaluate / generate / diffusion.sample / embed_texts) but left to_onnx as the lone exception — it called self.net.eval() (correct for tracing) with no paired restore. A user doing model.train(...) → model.to_onnx(...) → model.train(...) silently disabled Dropout / BatchNorm-running-stats on the second train() call. Wrapped the export body in try: ... finally: if was_training: self.net.train(), matching the shape nnx.viz.netron_export already used. Two new regression tests in tests/test_to_onnx_inputs.py cover the train→to_onnx→train round-trip AND the eval-mode caller preservation.
• NNTransformerParamsBuilder.dropout() / .context() honor "last call wins" (post-PR-#50 overnight-maintenance pass). Two setters still carried the if value is not <default>: gating anti-pattern that PR #50 cleaned up everywhere else: .dropout() gated on if attn != 0.0: / if resid != 0.0:, so .dropout(attn=0.5).dropout(attn=0.0) left attn_dropout at 0.5; .context() gated on if rope_base is not None:, so .context(rope_base=500000.0).context(max_seq_len=2048) left the prior 500000.0 in place. Fixed by storing unconditionally on .dropout() (Builder defaults match dataclass defaults; omit-when-default in state() handles run.id stability) and switching .context() to a delete-or-store pattern (pass-through None drops any prior override so the dataclass default governs at build time). Two new regression tests cover both paths.
• nnx.interop.export_ollama_modelfile writes Modelfile with explicit encoding="utf-8" (post-PR-#50 overnight-maintenance pass). The post-PR-#39 utf-8 audit (5 open(...) call sites fixed) used a grep "open(" sweep that didn't catch a later-added Path.write_text(...) in src/nnx/interop/ollama.py. A non-ASCII SYSTEM prompt or TEMPLATE (Asian-language fine-tune, emoji prompt) silently mojibake-encoded on Windows pre-PEP-686 where the platform default is cp1252. Added encoding="utf-8" to the write_text call. New regression test test_export_ollama_modelfile_writes_utf8_for_non_ascii_system writes a mixed-script Japanese + accented Latin + emoji SYSTEM prompt and asserts a clean utf-8 round-trip via read_bytes().decode("utf-8").
• Example 13 + 12 docstring "Requires" + examples/README.md install matrix sync (post-PR-#50 overnight-maintenance pass). By the maintainers' convention, every example using a [project.optional-dependencies] extra must declare it in the module docstring's "Requires" paragraph. examples/13_train_domain_embedder.py imports faiss + references sentence_transformers but had no "Requires" paragraph for [embeddings]. examples/12_quantize_int8.py's Phase-5 step calls model_q.to_onnx(...) — the legacy TorchScript exporter needs the onnx PyPI package — but the install line said only pip install thekaveh-nnx[quantize]. examples/README.md's install matrix listed 15_qat_classifier.py only under [quantize] despite its Phase-5 dynamo export needing [onnx-dynamo]. All three sync'd.
• Phase-checkpoint quartile logic consolidated into a single phase_tag() helper (post-PR-#51 overnight-maintenance pass). The FIRST/Q1/Q2/Q3 epoch-boundary branch chain (if idx_epoch == 0: ... elif idx_epoch == int(n_epochs * k/4) - 1: ... for k = 1, 2, 3) was duplicated verbatim across src/nnx/nn/nn_model.py:968-975 (the NNModel.train write path) and src/nnx/trainer/trainer.py:393-401 (the multi-optim Trainer.train write path). A future fix to the phase-boundary semantics (e.g., handling the small-n_epochs collision documented below) would have had to be made in two places. Extracted phase_tag(idx_epoch, n_epochs) -> Optional[Checkpoints] to src/nnx/nn/enum/checkpoints.py; both call sites now delegate. Pure refactor — no behavior change. Helper docstring documents the small-n_epochs caveat explicitly: for n_epochs in [4, 5, 6, 7] the Q1 index collides with FIRST (int(n_epochs * 1/4) - 1 == 0) and Q1 is silently never written; for n_epochs in [1, 2, 3] Q2/Q3 also miss. Changing the math would shift run.id-relevant trajectories for callers already relying on the current schedule, so the silent-skip is preserved as a documented trade-off rather than fixed. New tests/test_phase_tag.py ships 13 parametrized unit tests locking the schedule in: FIRST-at-zero invariant, canonical n_epochs=8 quartiles, "every quartile reachable for n_epochs >= 6" parametrize matrix, "Q1 dropped for n_epochs in [1,5]" parametrize matrix, and non-boundary-returns-None sanity.
• NNTrainerParamsBuilder.build() surfaces n_epochs missing at the Builder boundary (post-PR-#51 overnight-maintenance pass). Calling NNTrainerParams.builder().optimizer("main", ...).build() previously raised TypeError: NNTrainerParams.__init__() missing 1 required keyword-only argument: 'n_epochs' — the bare dataclass error names nothing actionable. The Builder already catches the harder schedulers.keys() ⊆ optims.keys() invariant at the .build() boundary with an actionable error (the Builders' pre-check); the missing-n_epochs case fell through to the dataclass ctor, which was the one error path that did not raise a Builder-level error naming the fix. Added a pre-construct check that raises ValueError("NNTrainerParamsBuilder.n_epochs() must be called before .build() — n_epochs has no meaningful default. Example: .n_epochs(50).optimizer('main', NNOptimParams(...)).build()"). New regression test test_builder_rejects_build_without_n_epochs asserts the error message names .n_epochs() (the method to call), not the dataclass field.
• PyPI keywords drift caught up (post-PR-#51 overnight-maintenance pass). Per the maintainers' PyPI-keyword convention, every major subsystem shipped as a first-class nnx.* entry point should appear in pyproject.toml [project] keywords for PyPI search discoverability. Three subsystems landed without matching keywords: learning-rate-finder (PR #35's nnx.lr_finder), multi-optimizer (PR #46's composite nnx.trainer.Trainer + NNTrainerParams), and builder-pattern (the cross-cutting <Class>.builder() fluent API on every params dataclass, PRs #43–#46 + #48). Added all three. Note: the GitHub repo's About-section Topics list (curated 20-max subset for browse discovery) was independently updated post-merge; pyproject keywords are the broader feature inventory and don't need 1:1 parity.
• tests/test_imports.py regresses against _PackageNotFoundError namespace re-leak (post-PR-#51 overnight-maintenance pass). PR #51 underscore-aliased importlib.metadata.PackageNotFoundError → _PackageNotFoundError in src/nnx/__init__.py to keep the importlib exception out of the nnx.* public surface, but explicitly shipped no regression test ("pure API hygiene"). A future maintainer dropping the alias to "simplify the imports" would silently re-introduce the leak; pre-emption: added test_packagenotfounderror_is_not_a_public_nnx_symbol asserting not hasattr(nnx, "PackageNotFoundError") with the underscore-alias fix in the error message. Same pattern applies to any future implementation-detail import that lives in the top-level nnx/__init__.py.
• NNParams.n_heads omit-when-default invariant gets an explicit on-disk-shape assertion (post-PR-#51 overnight-maintenance pass). Every other params class (NNTrainerParams, NNTransformerParams, NNOptimParams, NNSchedulerParams, NNTrainParams, NNModelParams) has an explicit assert "<field>" not in obj.state() test pinning the omit-when-default state() contract. NNParams.n_heads was the lone gap — the existing test_nn_params_round_trip (passing n_heads=None) and test_nn_params_round_trip_with_n_heads (passing n_heads=4) cover behavior via from_state(obj.state()) == obj but neither checks the on-disk key-presence shape directly. Added paired test_nn_params_state_omits_n_heads_when_none and test_nn_params_state_emits_n_heads_when_set to lock the run.id-stability contract for the FFN path (where n_heads is always None).
• examples/06_finetune_with_layer_freezing.py seeding convention drift fixed (post-PR-#51 overnight-maintenance pass). Per the maintainers' example-seeding convention, the helper-level torch.manual_seed(...) anti-pattern silently overrides the caller's seed; the established convention (examples 16 / 20 / 24) is that helpers carry an explicit # No torch.manual_seed here — the caller does set_seed(...) in main() comment and main() calls nnx.set_seed(...) once before each phase. Example 06 had two violations: _make_loaders(seed) called torch.manual_seed(seed) (torch-only, not numpy/python — partial seeding), and _make_model(seed) called set_seed(seed) (full seeding inside a helper, hiding the seed-management responsibility from main()). Both helpers now consume the current RNG state; main() calls set_seed(0) before Phase-1 model+loader construction, set_seed(1) before Phase-2 model construction (different init), and set_seed(42) before the Phase-2 loader (distribution B). The two-phase distinct-distribution semantics are preserved.
• text_contrastive_train_step_factory promoted to top-level nnx.* + example 02 seeding moved to main() (post-PR-#53 overnight-maintenance pass). Two small-surface ergonomics + convention catches. (1) nnx exposes 11 *_train_step_factory functions as the canonical extension point for custom training paradigms (KD / feature-KD / SimCLR / Mixup / CutMix / MoE / JEPA / DPO / QAT / Diffusion + Trainer-style multi-optim). 10 of them were at the top level (nnx.kd_train_step_factory, etc.); text_contrastive_train_step_factory was the lone outlier reachable only via nnx.embeddings.*. Asymmetry broke the nnx.<TAB> discoverability expectation (search for *_train_step_factory, find all paradigms). Promoted with a placement comment explaining that the high-level train_contrastive and FAISS-export helpers intentionally stay under nnx.embeddings.* (the top-level surface stays focused on train-step entry points). (2) examples/02_resume_training.py:33 had set_seed(7) inside the _make_model_and_loader() helper — the helper-level seed-override anti-pattern (caught repeatedly before across examples 06/16/19/20/21/23/24). Moved to main() where it's called twice (once per round) so the reproducibility contract is visible at the entry point. Helper carries the canonical "No set_seed here — the caller does it in main()" comment. Behavior preserved: both rounds still build the same initial weights (Round 2's get overwritten by load_state_dict on resume), and the DataLoader shuffle order stays pinned for an apples-to-apples training trajectory.
• Reproducibility hardening: NNDataset + NNTabularDataset seeded splits, set_seed sets PYTHONHASHSEED, NNRun.save pins sort_keys=True (post-PR-#53 overnight-maintenance pass). Three reproducibility gaps caught by the new determinism-audit angle. (1) NNDataset.__post_init__ and NNTabularDataset.__post_init__ called torch.utils.data.random_split(...) WITHOUT a generator= arg. The split consumed the global torch RNG, so two runs with identical set_seed(42) could diverge if any intervening code (model layer init, weight registration, anything that calls a torch random op) consumed RNG state between the seed call and the dataset construction. The sibling NNPreferenceDataset already had the seeded-split contract (PR #42 — seed: Optional[int] field + gen = torch.Generator(); gen.manual_seed(self.seed); random_split(..., generator=gen)). Both holdouts now match: added seed: Optional[int] = None field (the default was intended to keep pre-fix global-RNG behavior for back-compat, but shipped passing a fresh fixed-seed torch.Generator() instead — see the post-#54 fix above, which restored the documented global-RNG fallback); when set, builds a deterministic Generator and passes to random_split. New regression test test_f8_tabular_dataset_seeded_split_is_deterministic asserts (a) same seed → identical train/val/test ids, (b) different seed → different ids (200-row sanity check; collision probability astronomically small). (2) set_seed(seed) previously seeded random + numpy.random + torch + torch.cuda + cuDNN deterministic mode but did NOT touch os.environ["PYTHONHASHSEED"]. DataLoader workers using the spawn start method (default on Windows + macOS/Py3.8+) re-randomize their dict/set hash seed on every spawn — any worker code that iterates a dict/set populated in-worker could differ between runs even when every other RNG is seeded. Now writes PYTHONHASHSEED=<seed> so spawned children inherit. Docstring documents the caveat that the current Python interpreter's hash state was fixed at startup and is NOT affected by this assignment (only spawned subprocesses); for full hash determinism in the current process, set the env var in the shell before invoking Python. (3) NNRun.save called yaml.safe_dump(self.state()) and yaml.safe_dump(env_snapshot()) without sort_keys — relied on PyYAML's post-5.1 default of True. Pinned sort_keys=True explicitly so the on-disk YAML stays alphabetically stable across PyYAML major-version bumps that might change the default again. run.id is md5(str(state())) and unaffected (Python 3.7+ dicts preserve insertion order); this is purely the on-disk YAML shape downstream tooling diffs against.
• src/nnx/nn/__init__.py gets a docstring and explicit empty __all__ (post-PR-#53 overnight-maintenance pass). The nn subpackage's __init__.py was zero bytes — every sibling subpackage (embeddings, peft, paradigms, viz, …) carries a top-level docstring describing its surface. Added a docstring naming the seven internal modules (params, net, dataset, enum, callbacks, nn_model + generative_nn_model, moe), with the intentional convention that users reach the public surface via the top-level nnx namespace (from nnx import NNModel, NNParams, …) not via nnx.nn.* deep imports. __all__: list[str] = [] makes from nnx.nn import * a deliberate no-op — discouraging callers from bypassing the curated top-level surface. Pure documentation; no behavior change.
• Idiom + dataclass consistency catch-up (post-PR-#53 overnight-maintenance pass). Three small-surface consistency drifts caught by the Python-idiom audit. (1) LRFinderResult (src/nnx/lr_finder.py:27) was the lone params-shaped dataclass missing @dataclass(frozen=True, kw_only=True, slots=True) — every sibling result/params class (NNIterationDataPoint, NNEvaluationDataPoint, LoadPretrainedResult, etc.) carries the full triple. Single call site at line 185 already uses kwarg-only construction; tests don't mutate the result; safe to upgrade. (2) Eight .py files with real type hints were missing from __future__ import annotations (97/111 files had it; the 14 outliers split into 6 truly-empty __init__.py files that don't need it and 8 type-hint-carrying files that should). Added to vis_utils.py, nn/net/feed_fwd_nn.py, nn/net/graph_conv_nn.py, nn/net/graph_att_nn.py, nn/net/graph_sage_nn.py, nn/dataset/nn_dataset.py, nn/dataset/nn_dataset_base.py, nn/dataset/nn_graph_dataset.py. Cleaned up two "FeedFwdNN" quoted forward refs that ruff UP037 flagged once the future import lifted the need. (3) Raises: blocks added to all four Builder .build() method docstrings (NNOptimParamsBuilder / NNSchedulerParamsBuilder / NNTransformerParamsBuilder / NNTrainerParamsBuilder) — the Builders' pre-check ValueError contract was previously documented only in prose, leaving Sphinx / IDE tooltips without the canonical exception-doc form that 10+ sibling sites already use.
• Public-API surface: nnx.trainer.NNTrainerParamsBuilder subpackage re-export + 8 missing subpackages added to nnx.__all__ (post-PR-#52 overnight-maintenance pass). Two cross-module API drift issues caught by the round-N+1 hygiene scan. First: from nnx.trainer import NNTrainerParamsBuilder raised ImportError. The top-level nnx.NNTrainerParamsBuilder worked because src/nnx/__init__.py imported it via from .trainer.params_builder import ..., bypassing the subpackage's own __all__. Every other Builder supports the from <subpackage> import <Builder> convention (from nnx.peft import LoRALinear etc.); the trainer subpackage's __init__.py was the lone gap. Added from .params_builder import NNTrainerParamsBuilder + the matching __all__ entry. Second: nnx.__all__ previously listed only four subpackages (viz / embeddings / interop / prune); the other eight specialization subpackages (peft / diffusion / finetune / generation / paradigms / quantize / surgery / trainer) were attribute-accessible (via Python's side-effect attribute binding on from .X import Y) but absent from the documented public surface that IDEs, doc generators, Sphinx autosummary, and from nnx import * all read. All 12 subpackages now consistently listed in __all__. New regression tests in tests/test_imports.py: test_nn_trainer_params_builder_importable_from_subpackage locks the subpackage-level import contract; test_subpackages_appear_in_top_level_all locks the __all__ membership for all 12 subpackages so future additions can't silently drift.
• NNParams.dims property gains type-narrowing assert (post-PR-#52 overnight-maintenance pass). _dims: Optional[list[int]] is set unconditionally in __post_init__ via object.__setattr__, so the Optional is a dataclass artifact (slotted frozen subclasses can't take a non-Optional init=False field without default_factory). Pyright can't model __post_init__-set fields and flagged the dims property as returning Optional[list[int]] where the declared type is list[int]. Added a one-line assert self._dims is not None for both pyright narrowing and runtime sanity (the contract that dims is always available after __post_init__). Surgical fix from the post-#52 pyright-triage audit; one of two-to-three truly fixable items the audit picked out from the 89 pre-existing pyright errors that CI tolerates.
• seeding.py git-helper exception rationale comments (post-PR-#52 overnight-maintenance pass). Per maintainer review, every except Exception: in src/nnx/ should document its rationale inline. _nnx_version already carried a detailed comment (PR #51); the two sibling helpers _git_commit and _git_dirty swallowed their broad exceptions silently. Added one-line comments above each except: env_snapshot() is opportunistic — the caller may not be in a git repo (CI tarball install, fresh PyPI install, tempdir runs), git may not be on PATH, the 2-second timeout may fire, or the subprocess may crash. metadata.yaml omits the field rather than surfacing an exception. Pure documentation; no behavior change.
• PEFT load_*_weights source-resolution consolidated into a shared _resolve_source_to_state_dict helper (post-PR-#52 overnight-maintenance pass). Every PEFT adapter's load_<adapter>_weights(module, source) function accepted either a filesystem path or a state-dict and resolved it to a plain dict before applying adapter-specific key filtering. The resolution step was duplicated verbatim across load_lora_weights / load_ia3_weights / load_prompt_weights / load_prefix_weights and carried the security-critical weights_only=True invariant on the torch.load call. Four sites meant a future tightening (or relaxation under documented context) would have to land in four places — exactly the kind of duplication a round-3 merge audit catches. Extracted _resolve_source_to_state_dict(source, fn_name) to a new private module src/nnx/peft/_source.py; the four callers each shed ~7 lines and gained a single helper call. The nn.finetune.loading.load_pretrained companion intentionally does NOT consume this helper — its surface is wider (also accepts nn.Module, uses map_location="cpu"). All 70 PEFT tests pass; the existing match="path or dict" TypeError regressions in test_peft_lora.py / test_peft_ia3.py still lock the error message format, which the helper preserves verbatim.
• Builder-boundary pre-check applied uniformly to NNOptimParamsBuilder / NNSchedulerParamsBuilder / NNTransformerParamsBuilder (post-PR-#52 overnight-maintenance pass). PR #52 established the Builder-boundary pre-check on NNTrainerParamsBuilder — .build() pre-empts a missing-required-no-default field with an actionable Builder-level ValueError naming the Builder method to call (.n_epochs(...)), not the dataclass field. The three sibling builders were the lone holdouts, each surfacing the dataclass's bare TypeError: missing N required keyword-only arguments and forcing the user to translate dataclass field names back to Builder method names. Now: NNOptimParamsBuilder.build() raises ValueError("NNOptimParamsBuilder: call one of .adam(...) / .adam_amsgrad(...) / .sgd(...) / .sgd_nesterov(...) ...") if no variant was selected; NNSchedulerParamsBuilder.build() raises the analogous ValueError naming its five variants; NNTransformerParamsBuilder.build() walks the missing setters (.vocab(size=...) / .layers(n=..., heads=..., d_model=...) / .context(max_seq_len=...)) and emits a precise list of what's missing. Each builder's existing test_builder_build_without_variant_raises / test_builder_build_without_vocab_raises regression updated to assert the new ValueError message format and lock the Builder-method-naming contract.
• NNTransformerParams joins the centralized test_params_round_trip.py contract suite (post-PR-#52 overnight-maintenance pass). Per the omit-when-default state() invariant, every params class with state() / from_state() should have both a round-trip test and an explicit assert "<field>" not in state() test in the centralized tests/test_params_round_trip.py (the official enforcement mechanism). NNTransformerParams was the lone gap — its omit-when-default coverage lived only in test_nn_transformer_params_builder.py (domain-focused), missing from the central suite that catches integration drift. Added three tests: test_nn_transformer_params_round_trip_defaults (round-trip with every optional knob at default), test_nn_transformer_params_round_trip_with_overrides (round-trip with non-defaults so from_state honors them), and test_nn_transformer_params_state_omits_defaults (explicit on-disk-shape assertion for ffn_mult / rope_base / tie_embeddings / attn_dropout / resid_dropout).
• examples/README.md install matrix completes example 18's [lm] coverage (post-PR-#51 overnight-maintenance pass). examples/18_publish_to_ollama.py uses NNTokenizerParams + train_bpe (HuggingFace tokenizers) — its own module docstring correctly says pip install 'thekaveh-nnx[gguf-write,lm]' — but examples/README.md listed it only under the [gguf-write] line. A user following only the README's install matrix and running pip install thekaveh-nnx[gguf-write] would have hit an ImportError at train_bpe. Added 18 to the [lm] line, mirroring the dual-list pattern already used for example 12 (under [onnx] and [quantize]) and example 15 (under [quantize] and [onnx-dynamo]).

---

Expansion megamerge details (PR #29)

This release integrates 20 sub-projects consolidated on 2026-05-28: HuggingFace Hub interop (safetensors + PyTorchModelHubMixin), PEFT additions (DoRA + IA3 + Prefix + Prompt tuning on top of LoRA + Adapters), quantization (PTQ INT8 weight-only + QAT 8da4w via torchao), pruning (magnitude + 2:4 semi-structured), model surgery (Net2Net widen / deepen + drop_layer + low_rank_factorize + expand_embedding), embeddings (contrastive trainer + FAISS export), decoder-only LM (TransformerNN + NNTransformerParams + NNTokenizerParams + GenerativeNNModel.generate() with KV-cache), GGUF write + Ollama Modelfile bundle, model-internals visualization (torchinfo summary + weight histogram + activation map + Captum attribution + Netron export), I-JEPA self-supervised pretraining (+ small ViTNN encoder), Mixture-of-Experts (MoELinear + moe_train_step_factory with Switch-style aux loss), Born-Again Networks (iterated self-distillation), Feature-KD (FitNets-style), DPO (preference fine-tuning for LMs), LogitsProcessor chain (temperature / top-k / top-p / repetition-penalty), ONNX dynamo export opt-in, and assorted ergonomic improvements.

Every change preserves back-compatibility with existing run.id hashes and on-disk checkpoint formats — new params fields all follow the omit-when-default state() invariant.

Fixed — ONNX input coercion (PR #30)

• NNModel.to_onnx(example_input=np.ndarray) — a single 2-D np.ndarray was being unpacked row-by-row into N rank-1 inputs because np.ndarray is iterable; only torch.Tensor was special-cased in the singleton-wrap branch. torch.onnx.export then raised TypeError: forward() takes 2 positional arguments but N+1 were given. Fix extends the singleton check to (torch.Tensor, np.ndarray); the subsequent per-element coercion handles both consistently. New tests/test_to_onnx_inputs.py regresses the four shapes the docstring promises (Tensor singleton, ndarray singleton, tuple, mixed tuple).

Added — model-internals viz attribution + ONNX dynamo opt-in

• nnx.viz.attribute(model, x, *, method, target, **method_kwargs) — Captum-backed input-attribution wrapper. Single string-keyed dispatch over the six most common methods (integrated_gradients, gradient_shap, deep_lift, saliency, input_x_gradient, occlusion) returning (attribution_tensor, plotly.Figure). The figure renders the attribution as a Plotly Heatmap (3-/4-D image-shaped inputs are mean-pooled over channels first). Captum is lazy-imported at the call site so the rest of nnx.viz keeps working without it; the missing-dep path raises a clear ImportError("nnx.viz.attribute requires captum: pip install captum"). Sensible per-method defaults (baselines=zeros for GradientShap, sliding_window_shapes for Occlusion) preserve the one-call ergonomics. Optional dep promoted into the existing viz extra: pip install thekaveh-nnx[viz] now pulls captum>=0.7.0 alongside torchinfo>=1.8.0. 10 new tests in tests/test_viz_attribute.py (unknown-method ValueError, IG return-shape + figure-type, saliency works, missing-captum ImportError via sys.modules stub, every supported-method key end-to-end via @pytest.mark.parametrize).
• NNModel.to_onnx(..., dynamo=True) opt-in. New dynamo: bool = False kwarg on NNModel.to_onnx. When True, dispatches through PyTorch's new torch.export-based ONNX exporter (default in torch>=2.9; supports >2 GB models via external data; generally faster). The default (False) preserves the existing legacy TorchScript path — no behavior change for existing callers. The dynamo path lazy-imports onnxscript and raises a clear ImportError pointing at the new thekaveh-nnx[onnx-dynamo] extra (pip install thekaveh-nnx[onnx-dynamo]) if missing, rather than letting torch surface a less actionable failure.

Added — quantization (PTQ INT8 weight-only via torchao)

• nnx.quantize package — post-training quantization built on top of torchao (the replacement for the deprecated torch.ao.quantization, which is removed in PyTorch 2.10).
  • nnx.quantize_int8(model: NNModel) -> NNModel — one-call PTQ INT8 weight-only quantization. Deep-copies model.net, applies torchao.quantization.quantize_(net, Int8WeightOnlyConfig(version=2)) to the copy, and returns a new NNModel whose net.Linear weights are stored in int8 per-channel (symmetric). Activations stay FP32 — only the weights are quantized, so accuracy loss is typically a fraction of a percentage point. No calibration data, no retraining — pure post-process. The original NNModel is left untouched so callers can keep both around for an accuracy delta comparison.
  • Vision + GNN compatible — any module exposing nn.Linear submodules is a valid target.
  • ONNX export still works on the quantized model (NNModel.to_onnx routes through torch.onnx.export's legacy tracing path; torchao's quantized tensor falls back to dequantized matmul during the trace, so the exported ONNX is FP32 with the quantized weights baked in). Regression test included.
  • State-dict round-trips through NNCheckpoint.to_file (the existing pickle path); the on-disk file shrinks by roughly the same ratio as the pickled state-dict (≈30% on the example below, closer to ~25% at production-scale dims).
• Runnable demo: examples/12_quantize_int8.py — trains a small classifier (FP32), prints the FP32 val accuracy + state-dict size, calls quantize_int8 once, prints the INT8 val accuracy + size, and confirms the quantized model still ONNX-exports. On the toy task the size reduction lands at ~69% with zero measurable accuracy delta.
• New optional dependency: pip install thekaveh-nnx[quantize] (pulls torchao>=0.17).
• 15 new tests in tests/test_quantize_ptq.py covering: returns a fresh NNModel, preserves output shape, doesn't mutate the source, replaces Linear weights with a torchao-quantized tensor, preserves attached attrs (params / net_params / device / loss_fn), output stays within 5% relative L2 of FP32, pickled state-dict shrinks vs FP32, NNCheckpoint.to_file round-trip shrinks on disk, ONNX export round-trip, deep-copy isolation (mutating quantized doesn't leak back), predict() end-to-end on a deeper model, clear ImportError when torchao is missing, state-dict keys unchanged, idempotency-via-deep-copy (calling twice on the same source produces identical outputs), .train() / .eval() toggle still works.

Also shipped in this sub-section: nnx.quantize.qat_train_step_factory + nnx.quantize.QATLifecycleCallback — Int8DynActInt4WeightQATQuantizer fake-quant during training, real-quant on convert via the QATLifecycleCallback (prepare → train → convert lifecycle pinned to epoch boundaries). Re-exported from nnx.*. Tests in tests/test_quantize_qat.py cover prepare/convert idempotency, end-to-end training, and ONNX export of the converted model.

Deferred from this sub-section: INT4 weight-only PTQ lands in a separate follow-up PR.

Added — PEFT++ (IA3)

• IA3Linear(base) — Infused Adapter by Inhibiting and Amplifying Inner Activations (Liu et al., NeurIPS 2022). The smallest adapter in the PEFT family: a single learned per-output-dim scaling vector applied multiplicatively to a frozen nn.Linear's output. Trainable parameter count per wrapped layer is exactly out_features — roughly two orders of magnitude smaller than LoRA at the same effective adaptation budget. scaling is initialized to all-ones so the forward output at step 0 equals base(x) exactly.
• apply_ia3_to(module, *patterns) — fnmatch-glob in-place wrap mirroring apply_lora_to. Same two-phase traversal and idempotency contract (existing IA3 wrappers are not re-wrapped).
• save_ia3_weights(module, path) / load_ia3_weights(module, source) — persist ONLY the scaling parameters, symmetric to LoRA's save/load idiom. The resulting checkpoint is tiny (a single vector per wrapped layer). Same weights_only=True safety guarantee; same empty-dict-is-zero-op contract; same dict-source convenience overload.
• 19 new tests in tests/test_peft_ia3.py: validation (non-Linear base rejection), base-freezing, zero-init invariant (output == base at step 0, with and without bias), forward shape, trainable parameter set is exactly {scaling}, scaling init is all-ones, in/out features pass-through, scaling actually scales the output by a known non-unit value; apply_ia3_to empty-pattern rejection + selective wrap + wildcard wrap + idempotency + forward-preserves-at-init; save/load round-trip + base-keys-excluded-from-checkpoint + dict-source loading + bad-source-type rejection + empty-dict no-op contract.

Added — PEFT++ (DoRA)

• DoRALinear(base, *, r, alpha, dropout) — Weight-Decomposed Low-Rank Adaptation (Liu et al., NVIDIA, ICML 2024 Oral). Subclass of LoRALinear that adds a trainable per-output-row magnitude parameter and recomposes the layer's weight as W = magnitude * V / ||V||_c where V = W_0 + (α/r) · BA is the LoRA-augmented direction. magnitude is initialized from ||W_0||_c so the forward output at step 0 equals base(x) exactly (combined with LoRA's zero-init B). Often outperforms LoRA at the same rank with only out_features extra parameters — negligible vs LoRA's r · (in + out) baseline.
• apply_dora_to(module, *patterns, r, alpha, dropout) — fnmatch-glob in-place wrap mirroring apply_lora_to. Same idempotency contract (existing LoRA/DoRA wrappers are skipped via the parent-is-LoRALinear check, which covers DoRALinear by inheritance).
• DoRA reuses save_lora_weights / load_lora_weights for the lora_A / lora_B matrices unchanged (the inheritance hierarchy ensures the LoRA filter still matches). The magnitude vector is captured by the standard state_dict() round-trip — single vector of length out_features per wrapped layer.
• 16 new tests in tests/test_peft_dora.py: validation (non-Linear base, r/alpha/dropout ranges), base-freezing, zero-init invariant (output == base at step 0, with and without bias), forward shape, trainable parameter set is exactly {lora_A, lora_B, magnitude}, magnitude init matches ||W_0||_c, in/out features pass-through, LoRALinear subclass relationship; apply_dora_to empty-pattern rejection + selective wrap + wildcard wrap + idempotency on re-application + forward-preserves-at-init; save_lora_weights round-trip via DoRA wrappers.
• nnx.paradigms.feature_kd_train_step_factory(teacher, *, auxiliary_layers, alpha, beta, temperature) — FitNets-style intermediate-layer feature distillation. Extends the existing kd_train_step_factory with an additional MSE term between named teacher / student intermediate-layer activations: L = α · KL_soft · T² + β · MSE(student_act, teacher_act) + (1 − α) · L_hard. Forward hooks register on the auxiliary_layers pairs (teacher_layer_name → student_layer_name, resolved via nn.Module.get_submodule); activations are collected per forward and the MSE term is averaged across pairs so beta's scale is invariant to the pair count. The teacher freeze + eval-mode guarantee carries over from kd_train_step_factory. The v1 factory requires shape-matched paired layers — the FeatureRegressor projector for mismatched widths is deferred. Routes through finalize_step for the standard NaN guard + grad-clip path. Re-exported from nnx.paradigms.* and nnx.*.

Added — born-again self-distillation

• nnx.paradigms.born_again_train(model, *, generations, train_params, **kd_kwargs) -> list[NNRun] — iterated self-distillation wrapper. Generation 0 trains plain (no teacher); each subsequent generation uses a deep-copied, frozen, eval-mode snapshot of the model after the prior generation as the teacher in a Hinton-style KD step (composed via kd_train_step_factory). Returns the per-generation NNRun list so callers can inspect the convergence trajectory. Born-again networks (Furlanello et al., ICML 2018) often match or slightly outperform the original — the soft targets act as an implicit regularizer. 9 new tests covering generations-count validation, KD-factory not invoked on generation 0, KD-factory invoked on generations 1+, teacher snapshot is a deepcopy (not the live model), teacher requires_grad=False + eval-mode at handoff, kwargs forwarding, teacher isolation from subsequent training, top-level re-export, and end-to-end model mutation across generations.

Added — Mixture-of-Experts (tutorial-grade)

• nnx.MoELinear(in_features, out_features, *, num_experts, top_k=2) — sparse top-k MoE drop-in for nn.Linear. Router (bias-less nn.Linear) emits per-expert logits; the top_k experts per token are selected, their outputs are weighted by softmax-renormalized gating values, and the per-token result is the weighted sum. Exposes .last_aux_loss after each forward — the Switch-Transformer load-balancing penalty N · Σ_i f_i · P_i where f_i is the dispatch fraction and P_i is the mean router probability for expert i. The penalty is minimized at value 1 (NOT 0) when routing is perfectly uniform across experts. Validates num_experts ≥ 2, top_k ∈ [1, num_experts] at construction.
• nnx.paradigms.moe_train_step_factory(*, aux_loss_weight=0.01) — supervised training step that adds aux_loss_weight · Σ_layer layer.last_aux_loss to the main loss, summed across every MoELinear in the net. Routes through the shared _step_helpers.finalize_step for the standard NaN-guard + grad-clip tail (same shape as the KD / SimCLR / Mixup / CutMix factories). Works on nets with zero MoE layers too — the aux sum just collapses to 0 and the step is exactly supervised.
• Runnable demo: examples/14_moe_classifier.py — a feed-forward classifier whose hidden layer is an MoELinear (4 experts, top-k=2). Prints router / expert / classifier param breakdown, trains with moe_train_step_factory, and verifies the aux loss decreases across the run (routing balances out).
• 22 new tests across tests/test_nn_moe.py (12) and tests/test_paradigms_moe.py (10): MoELinear input validation + forward shape + router / experts module shape + top-k routing invariant + last_aux_loss populated-after-forward + non-negativity + uniform-routing-equals-1 (minimum-value math) + above-minimum-when-skewed + load-balancing converges under SGD on the aux loss; paradigm factory validation + end-to-end aux-loss-decreases + finalize-step routing (NaN guard fires) + no-MoE-layers no-op + zero-weight collapse to supervised + multi-MoE-layer summation + AMP rejection + grad-clip honored + EDP return shape.
• Scope explicitly limited to tutorial-grade. Production-scale MoE (MegaBlocks block-sparse kernels, expert parallelism across GPUs, token-dropping with capacity factor) is OUT — would be hollow wrapping over specialized libraries.

Added — pruning (nnx.prune)

• nnx.prune package — two complementary network-pruning strategies layered on top of plain nn.Linear submodules, mirroring the nnx.peft package shape (public functions, fnmatch glob patterns, in-place mutation).
  • magnitude_prune(net, sparsity, *, layer_pattern="*", bake=True) — wraps torch.nn.utils.prune.l1_unstructured. For each nn.Linear whose dotted name matches layer_pattern, zeros the round(sparsity · numel) smallest-magnitude entries of its weight matrix. Checkpoint-compat invariant: bake=True (default) calls prune.remove immediately after each layer is pruned, so the state_dict keys stay identical to the pre-prune network — pruned checkpoints load into unpruned-network code under strict=True. bake=False keeps the reparameterization in place (state_dict carries weight_orig + weight_mask instead of weight); use this for iterative pruning schedules where successive magnitude_prune calls need to compose with the existing mask. Validates sparsity ∈ [0, 1). Returns the number of layers pruned (0 if layer_pattern matches nothing).
  • semi_structured_24(net, *, layer_pattern="*") — 2:4 semi-structured sparsity via torchao.sparsity.sparsify_ with semi_sparse_weight(). Swaps each matched nn.Linear's weight with a 2:4 structured-sparse tensor subclass. Real wall-clock speedup on Ampere+ GPUs (~1.1× inference, ~1.3× training per torchao's ViT/SAM benchmarks); CPU and pre-Ampere hardware are unsupported by the underlying sparse kernel. The torchao dep is loaded lazily inside the function body so users on the magnitude-only path pay no dep cost; the dep is installed transitively via the existing quantize (torchao>=0.17) tooling.
• 17 new tests across tests/test_prune_{magnitude,semi_structured}.py: zero-fraction correctness; state_dict-keys preservation under bake=True (THE checkpoint-compat invariant); pattern-filter selectivity; idempotency on already-zeroed weights; iterative bake=False path; sparsity bounds rejection + sparsity=0 no-op + no-match returns 0; smallest-magnitude-go-to-zero correctness; full state_dict round-trip into a fresh unpruned net; CUDA-gated swap-actually-happens (skipped on CPU); monkey-patched pattern-filter selectivity for semi_structured_24 (decouples nnx's filter logic from torchao's CUDA-only kernel); torchao-importorskip guard.
• Structured pruning that REMOVES channels / heads (and so breaks state_dict shape) is deferred — needs a per-architecture surgery API the existing checkpoint format doesn't yet support.

Added — HuggingFace interop (safetensors + Hub mixin)

• safetensors as opt-in checkpoint format. NNCheckpoint.to_file(path, format="safetensors") writes a safe, mmap-friendly file readable by ComfyUI / vLLM / AutoGPTQ / HuggingFace tools. NNParams, NNModelParams, and NNIterationDataPoint are JSON-serialized into the safetensors metadata dict (the spec only allows str -> str metadata, so a JSON wrapper is the cleanest fit). Pickle remains the default format for back-compat; NNCheckpoint.from_file(path) auto-detects via magic-byte sniff (modern torch.save starts with the ZIP container PK\x03\x04, legacy / bare pickle starts with \x80, safetensors starts with neither). Requires pip install thekaveh-nnx[hub].
• NNModel is now HuggingFace-Hub-publishable via PyTorchModelHubMixin. Free model.save_pretrained("./dir"), model.push_to_hub("user/repo"), and NNModel.from_pretrained("user/repo" | "./dir"). The on-disk layout is the canonical Hub flat layout: model.safetensors (weights), config.json ({"net_params": <state>, "params": <state>} using the public state() form NNRun hashes), and an auto-generated README.md model card. Without the hub extra installed, all three methods raise a clear ImportError pointing back at pip install thekaveh-nnx[hub].
• thekaveh-nnx[hub] extra — pulls in safetensors>=0.7.0 and huggingface_hub>=1.4.0. Both deps are runtime-import-guarded, so pip install thekaveh-nnx keeps working without them.
• docs/hub.md — when-to-use guide for both tracks (safetensors checkpoints + Hub mixin), a local save/load walkthrough, the Hub publish/download path, and the explicit non-goals.

Added — embeddings (contrastive trainer + FAISS export)

• nnx.embeddings package — the one RAG-adjacent surface NNx ships. Users train a domain-specific text embedder via the existing SimCLR / NT-Xent machinery, then export the trained model to a FAISS index for any retrieval framework (LangChain / LlamaIndex / Haystack / raw FAISS) to consume. NNx does NOT host the RAG stack — chunking, reranking, prompt orchestration, vector-DB clients are inference-time concerns and explicitly out of scope.
  • ContrastiveTextDataset(pairs) — wraps (anchor, positive) string tuples as a torch.utils.data.Dataset. Validates input shape + types up-front (empty list / non-tuple / non-string entries all raise ValueError).
  • train_contrastive(backbone, dataset, *, n_epochs, batch_size, lr, temperature, ...) — high-level training loop. Builds a DataLoader with the string-aware pair_collate, runs NT-Xent (nnx.nt_xent_loss) updates over the trainable parameters of the backbone (anything requires_grad=True; composes with nnx.freeze / nnx.apply_lora_to), returns the in-place-mutated backbone. Accepts either a sentence_transformers.SentenceTransformer or any plain nn.Module(list[str]) -> Tensor[B, D].
  • text_contrastive_train_step_factory(*, temperature) — lower-level TrainStepFn factory for users who want NNx's full callback / checkpoint / runs/<id>/ machinery wrapped around the contrastive step (drive it through NNModel.train(train_step_fn=...) with a DataLoader that yields (anchors: list[str], positives: list[str]) batches).
  • embed_texts(backbone, texts, *, batch_size, device, normalize) — inference-time encoder; runs torch.no_grad() + eval(). Used by export_to_faiss internally and exposed for ad-hoc similarity probes.
  • export_to_faiss(backbone, corpus, out_path, *, index_type, normalize, ...) — embed corpus → build a FAISS index of the requested type (IndexFlatIP for cosine via normalize-then-IP, IndexFlatL2 for L2 distance, IndexHNSWFlat for approximate ANN with M=32) → write to disk via faiss.write_index. Lazy faiss import with a clear "install nnx[embeddings]" message on the failure path.
  • export_to_safetensors(backbone, out_path) — persist backbone weights for HuggingFace Hub / sentence-transformers reload. Uses the safetensors format when the package is importable (transitive via sentence-transformers≥3); falls back to plain torch.save otherwise.
• embeddings optional extra in pyproject.toml — pins faiss-cpu>=1.7 + sentence-transformers>=2.7. Both are optional at import time; the package imports cleanly without them and the ImportError is deferred to the call site that actually needs each one.
• Runnable demo: examples/13_train_domain_embedder.py — synthesizes 40 (sentence, paraphrase) training pairs, trains a tiny bag-of-words hash embedder from scratch for 5 epochs (mean anchor-positive cosine: 0.61 → 0.98), exports to a FAISS IndexFlatIP, reloads from disk, and runs a top-3 query (the paraphrase comes back at #2 with cosine ≈ 0.99). Network-free, CPU-only, ~10s end-to-end.
• New docs page: docs/embeddings.md — when to use, install, quickstart, full API, composition with nnx.freeze / nnx.apply_lora_to, and the explicit "what this is NOT" list (no chunker, no reranker, no vector-DB client, no RAG-framework wrapper).
• 28 new tests across tests/test_embeddings_{contrastive,faiss_export}.py: dataset validation (empty / non-tuple / non-string entries), pair_collate shape, end-to-end "training reduces anchor-positive cosine distance" assertion on a synthetic 32-pair dataset (the headline TDD test), embed_texts batch-invariance + normalize on/off, text_contrastive_train_step_factory bad-batch rejection + weights-move-on-step, FAISS IndexFlat{IP,L2} + IndexHNSWFlat index construction, "embed 100-text corpus → save → reload → top-1 self-similarity" assertion (the FAISS-export TDD test), explicit-normalize override semantics, safetensors-roundtrip via both safetensors and the torch.save fallback. FAISS / safetensors tests skip gracefully when the optional dep isn't installed.
• tests/conftest.py sets KMP_DUPLICATE_LIB_OK=TRUE + OMP_NUM_THREADS=1 at session start — sidesteps a macOS-specific faiss-cpu segfault in its parallel search kernel when torch's libomp.dylib got loaded first. Harmless on Linux CI.

Added — LM path: TransformerNN + tokenizer + generate

• Nets.TRANSFORMER enum variant — decoder-only LM dispatched through the standard NNModelParams(net=Nets.TRANSFORMER, ...) factory path. Back-compat-safe: existing pre-TRANSFORMER run.yaml files load unchanged.
• TransformerNN — decoder-only stack matching LLaMA / Mistral conventions: token embeddings + N TransformerBlocks (pre-norm with RMSNorm + RoPE + SwiGLU FFN + multi-head causal attention) + final RMSNorm + tied LM head. KV-cache seam wired but the low-level default is use_cache=False; GenerativeNNModel.generate flips it on via forward_with_cache (see "Added — generation: LogitsProcessor chain + KV-cache") without changing call sites.
• NNTransformerParams(NNParams) — frozen dataclass holding vocab_size, n_layers, n_heads, d_model, ffn_mult, max_seq_len, rope_base, tie_embeddings, attn_dropout, resid_dropout. Lifts the GraphAttNN n_heads-on-NNParams pattern by subclassing. Every optional field omits itself from state() when at default — the broken-three-times omit-when-default invariant; covered by regression tests.
• NNTokenizerParams — wraps tokenizers.Tokenizer (HF Rust BPE). state() returns {"path": "<tokenizer.json>"}; the tokenizer payload lives on disk, only the pointer goes into run.yaml. Companion train_bpe(...) helper trains a tiny BPE from either file paths or an in-memory text iterator. Available when the thekaveh-nnx[lm] extra is installed.
• GenerativeNNModel(NNModel).generate(prompt, max_new_tokens, temperature, top_k, top_p, repetition_penalty, stop, seed) — autoregressive decode via a LogitsProcessor chain (RepetitionPenalty → TopKFilter → TopPFilter → TemperatureScaling; canonical order documented in the post-#54 fix above). temperature=0 short-circuits to deterministic greedy; same-seed sampling reproducibility is part of the contract.
• New example examples/11_tinystories_lm.py — end-to-end TinyStories-class training run: train a BPE on the corpus, build a small Transformer, train next-token prediction via a custom train_step_fn, then sample. Ships with an inline fallback corpus so it runs offline; --use-hf downloads TinyStories.
• New docs page docs/lm.md — when/how to use the LM path. Linked from README §1.2 + §5.
• pyproject.toml lm extra — tokenizers>=0.20, datasets>=2.20. Opt-in so the Rust tokenizer binary isn't pulled for non-LM users.

Migration notes

These two fixes shift run.id hashes on disk. Older runs/<id>/ directories on disk continue to load by their existing directory name; recomputed ids land in a fresh directory.

• Default-AMP runs now hash to a different run.id than they did between pass-2 and this audit. The mixed_precision=False default is now correctly omitted from state() (back-compat invariant from before pass-2).
• Plateau-scheduler runs now hash to a different run.id than they did between the Schedulers-enum addition and this audit. The kind=None default + its variant-specific knobs (step_size / T_max / max_lr / total_steps / warmup_steps) are now correctly omitted from state() when at their defaults (same back-compat invariant).

Fixed — back-compat invariant audit

• NNModelParams.state() omits mixed_precision when False. The field was added in pass-2 but always emitted into state(), breaking the omit-when-default back-compat invariant. Every default-AMP run had a shifted run.id versus pre-pass-2 runs with otherwise identical config. One-time hash shift: any existing default-AMP runs/<id>/ directory will recompute to a different id after this fix — load by the on-disk directory name still works; recomputed ids will land in a fresh directory.
• NNSchedulerParams.state() omits kind and the variant-specific knobs (step_size / T_max / max_lr / total_steps / warmup_steps) when None. Same omit-when-default invariant: a plain ReduceLROnPlateau NNSchedulerParams now hashes to the same run.id as it did before the Schedulers enum was added. Existing on-disk runs with explicit-None entries still load (the legacy form is tolerated in from_state).
• In-memory best_checkpoint tracking aligned with on-disk BEST. NNModel.train()'s best_checkpoint reassignment used a different comparison than the BEST write inside _save_checkpoints. When val_loader=None (so every val_edp is None), the in-memory tracker effectively held LAST while the on-disk BEST tracked training error. Both now go through the same _best_err helper.
• _best_err deduplicated. Was triplicated — a local closure in NNModel._save_checkpoints, a module-level helper in nn_run.py, and another module-level helper in trainer/trainer.py. Kept the nn_run.py version as canonical; the other two now import it.
• Paradigm step factories honor grad_clip_norm and guard against non-finite loss. The four paradigm train_step_fn factories (diffusion / SimCLR / Mixup / CutMix) plus KD now route through a shared nnx._step_helpers.finalize_step helper. Previously they silently dropped NNOptimParams.grad_clip_norm, and diffusion / SimCLR / Mixup / CutMix had no NaN/Inf guard — only KD checked. New explicit rejection: the helper raises ValueError if NNModelParams.mixed_precision=True (paradigm steps don't handle the scaler) or if accumulate_grad_batches != 1 (no cycle-aware accumulation). Previously these were silently ignored; users with those knobs set now see a clear error.
• ModelCheckpoint callback actually saves now. The body was if ctx.epoch in self.epochs: pass — a no-op stub. Now writes runs/<run.id>/checkpoints/<tag>_e<epoch>.pt via the atomic-write path on matched epochs.
• FeedFwdNN.from_file uses torch.load(weights_only=True) for consistency with NNCheckpoint.load_optimizer_state and load_pretrained. State-dicts are tensor-only; the strict loader works AND removes the arbitrary-code-execution risk on user-supplied paths.
• Documentation and comment cleanups: docs/index.md listed only pass-2 features (added the five new tracks); docs/concepts.md architecture diagram missed the five new subpackages (extended with a Specializations branch); examples/06's _make_loaders docstring claimed class-conditional Gaussians that the code didn't implement (rewrote); freezing.py docstring incorrectly claimed fnmatch * matches segment-boundaries (it matches across dots); loading.py key_map docstring said "substring replacement" but the code does prefix replacement; KD's loss formula in paradigms/distillation.py module docstring + inline comment AND docs/concepts.md all reversed the KL direction — the math is KL(teacher || student) (standard Hinton), but the doc strings read KL(student || teacher); peft/adapters.py activation docstring said nn.GELU() (instance) but the default is nn.GELU (class factory); README's enums-as-factories bullet was missing NoiseSchedulers. Internal phase labels (Track A / Track B / Track C / pass-2 R2 / R3 / R4) that had leaked into published code/docs/tests have been replaced with descriptions of WHAT the referenced thing is.

Added — model-internals viz (nnx.viz subpackage)

• nnx.viz subpackage — sibling of the existing nnx.vis_utils (which handles run-output viz: training curves, confusion matrices, t-SNE of checkpoint logits). nnx.viz covers the model itself rather than what the run produced. Two primitives ship in this PR; activation_map, netron_export, and Captum attribution land in a later PR.
  • nnx.viz.summary(model, *, input_size=..., depth=4, col_names=...) — Keras-style parameter table via a thin torchinfo.summary wrapper. Returns the torchinfo.ModelStatistics object directly so callers can both print the formatted table AND query .total_params / .trainable_params / .total_mult_adds for programmatic regression assertions. Accepts an NNModel (unwrapped to .net) or any nn.Module.
  • nnx.viz.weight_histogram(model, *, bins=64, cols=3, fig_width=1000, row_height=200) — per-parameter Plotly histogram grid. Walks model.named_parameters() and emits one Histogram trace per tensor in a subplot grid, consistent with vis_utils's Plotly-returning idiom. Useful for spotting dead layers, NaN / Inf weights, or saturation patterns. Raises ValueError on a parameter-less module (which would otherwise produce a silently-empty figure).
• New viz optional extra — pip install thekaveh-nnx[viz] pulls in torchinfo>=1.8.0. nnx.viz.summary raises a clear ImportError pointing at the extra if torchinfo is missing; weight_histogram only depends on plotly (already a core dep), so it works out of the box.

Added — PEFT (LoRA + adapters)

• nnx.peft package — two complementary patterns for parameter-efficient fine-tuning of pretrained networks.
  • LoRALinear(base, *, r, alpha, dropout) — wraps an nn.Linear, freezes the base's parameters (requires_grad=False) on construction, and adds two trainable matrices lora_A (r × in, Kaiming-uniform init) and lora_B (out × r, zero-initialized) whose product is added as a residual scaled by α/r. The zero-init on B means output at step 0 equals base(x) exactly — fine-tuning starts from the pretrained behavior and diverges only as B picks up gradient. Validates r > 0, alpha > 0, 0 ≤ dropout < 1 at construction.
  • apply_lora_to(module, *patterns, r, alpha, dropout) — walks module.named_modules() and replaces every nn.Linear whose dotted name matches any fnmatch glob with a LoRALinear wrapper, in place. Returns the count wrapped. Idempotent: re-applying against patterns that already match LoRA-wrapped layers is a no-op (the inner .base is excluded from the walk). Same glob conventions as nnx.finetune.freeze.
  • save_lora_weights(module, path) — writes ONLY the lora_A / lora_B parameters via torch.save of a filtered state-dict subset. A small percentage of the full state_dict size (single-digit % at production scale; closer to ~40% on tiny demo nets where r/dim is large — see docs/concepts.md §11 for the math).
  • load_lora_weights(module, source) — loads LoRA params from a path (weights_only=True for safety) or directly from a dict, via load_state_dict(strict=False) so the frozen base's missing keys don't raise. Returns the number of tensors loaded.
  • AdapterLayer(dim, bottleneck, activation=nn.GELU) — bottleneck residual block y = x + up(act(down(x))). up.weight and up.bias are zero-initialized so the layer is the residual identity at step 0. Composed by the user into a custom nn.Module — NNx doesn't ship a "wrap every block" helper because adapter insertion points are architecture-specific.
• Runnable LoRA demo: examples/07_lora_finetuning.py — pretrains a small classifier, wraps every Linear with LoRA, fine-tunes on a different distribution, explicitly verifies every base parameter is bit-exactly unchanged across the fine-tuning run, and compares the LoRA-only checkpoint size against a full state_dict snapshot.
• 23 new tests across tests/test_peft_{lora,adapters}.py: LoRALinear validation + base-freezing + zero-init invariant (output == base at step 0) + only-LoRA-trainable invariant + in/out features pass-through; apply_lora_to empty-pattern rejection + selective wrap + wildcard wrap + idempotency on re-application + forward-preserves-at-init; save/load round-trip + base-keys-excluded-from-checkpoint + dict-source loading + bad-source-type rejection; end-to-end PEFT contract (every base param bit-exactly unchanged + every lora_B param has moved); AdapterLayer shape + identity-at-init + parameter-count scaling + gradient-flow + dim validation + custom activation.

Added — training paradigms (KD / SimCLR / Mixup / CutMix)

• nnx.paradigms package — four TrainStepFn factories for non-vanilla supervised paradigms, all consumed via the existing NNModel.train(train_step_fn=...) hook. No new params dataclass, no NNModel changes; each is a self-contained closure.
  • kd_train_step_factory(teacher, *, alpha, temperature) — Hinton-style knowledge distillation. Mixes a temperature-softened KL divergence against the teacher's logits (α · KL · T²) with the standard hard-label loss ((1-α) · L_hard). The factory freezes the teacher's parameters and sets its net to eval mode on call, so the teacher provably cannot drift across the student's training. The hard term goes through the student's loss_fn so KD works with any classification loss.
  • simclr_train_step_factory(*, temperature) — SimCLR contrastive training. The training loader must yield (view1, view2) paired-view tensors per source sample. model.net is forwarded once per view (BatchNorm sees one view at a time). Reports the NT-Xent loss in both .loss and .error.
  • nt_xent_loss(z1, z2, *, temperature) — the SimCLR loss exposed as a standalone for users wanting to compose it into custom training loops.
  • mixup_train_step_factory(*, alpha) — Mixup batch augmentation: x' = λx_a + (1-λ)x_b with λ ~ Beta(α, α). Works for any input rank (tabular, sequence, image). Reports λ-weighted accuracy as the accuracy field; accuracy + error == 1.
  • cutmix_train_step_factory(*, alpha) — CutMix batch augmentation for 4D (B, C, H, W) image batches. Copies a random rectangle from x_b into x_a, then re-weights the loss by the actual cut area (which can be smaller than the nominal Beta draw when the box clips at an edge). Raises a clear ValueError on lower-rank input — CutMix's spatial cut isn't well-defined without H and W.
• Runnable distillation demo: examples/10_knowledge_distillation.py — pretrains a wider teacher (hidden_dims=[64, 64]) then distills into a much smaller student (hidden_dims=[16], roughly 4% of the teacher's parameters). The example explicitly verifies teacher weights are unchanged across the student's training run, demonstrating the factory's freeze guarantee. Honest about scope: doesn't claim to beat a non-distilled baseline on toy tabular data.
• 19 new tests across tests/test_paradigms_{distillation,contrastive,augmentation}.py: factory validation (alpha / temperature ranges), teacher freezing guarantee + teacher-eval-mode assertion, end-to-end loss-decreases (KD α=0.5) + α-boundary cases (α=0.0 collapse to supervised, α=1.0 pure distillation), NT-Xent properties (shape mismatch, finite + scalar output, loss smaller for aligned pairs than random), SimCLR step bad-batch-shape error, Mixup self-consistency (accuracy + error == 1), CutMix non-image input rejection + 4D end-to-end.

Added — diffusion (DDPM)

• nnx.diffusion package — DDPM-style diffusion training and sampling, layered entirely on top of the existing train_step_fn hook on NNModel.train() (no Trainer, no NNModel internals touched).
  • NoiseSchedulers — enum-as-factory with two variants: LINEAR(T, beta_min, beta_max) (original DDPM linear betas) and COSINE(T, s) (Improved-DDPM cosine schedule). Each enum value's __call__ returns a precomputed NoiseSchedule.
  • NoiseSchedule — frozen dataclass holding the derived tensors (betas, alphas, alphas_cumprod, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, posterior_variance). All 1D of length T. .to(device) returns a copy with every tensor migrated. Not state()-serialized — recoverable from (kind, T, kind-specific knobs).
  • DiffusionMLP(input_dim, hidden_dims, time_embed_dim) — small conditional MLP: sinusoidal time embed → projection → concat with flat x → MLP → noise prediction. forward(x, t) → ε_pred. Handles arbitrary-rank inputs by flattening + un-flattening. Intentionally minimal; image-space diffusion calls for a U-Net the user supplies, with the same schedule / step / sampler machinery.
  • diffusion_train_step_factory(schedule) -> TrainStepFn — closes over the schedule and returns a TrainStepFn suitable for NNModel.train(train_step_fn=...). Per batch: samples t ~ Uniform[0, T), samples ε ~ N(0, I), computes x_t, predicts noise, backprops MSE. Reports loss as both .loss and .error on the EDP so BEST tracking + ReduceLROnPlateau work.
  • sample(model, schedule, shape, device=, generator=) — reverse-diffusion sampler. Runs T backward steps under torch.no_grad() and model.net.eval(). The optional generator= enables reproducible sampling for notebooks.
  • sinusoidal_time_embed(t, dim) — standalone helper for the standard sinusoidal positional embedding, exposed for users building their own t-conditioned nets.
• NNModel.train() net-params fallback — the run-construction line now reads self.net_params (always set in __init__) instead of self.net.params (FeedFwdNN-specific attribute). Back-compat-safe: the values are identical for the existing supervised path. Lets callers swap model.net for a custom nn.Module post-construction (the same idiom the multi-optimizer Trainer's GAN demo uses) without breaking NNModel.train().
• Runnable diffusion demo: examples/08_diffusion_2d_mixture.py — DDPM on a 2D mixture of 4 Gaussians at (±2, ±2). Verified end-to-end (loss 1.0078 → 0.6048; samples land in all four modes at roughly equal counts).
• 27 new tests across tests/test_diffusion_{schedules,nets,training,sampling}.py covering schedule shape/monotonicity/clamping, net forward shape, full training + loss-decreases, sampling shape / finiteness / reproducibility / mode coverage.

Added — multi-optimizer Trainer (GAN / actor-critic)

• nnx.trainer package — Trainer class that parallels NNModel.train() for scenarios where the per-batch update isn't a single supervised forward/backward/step. Built around the GAN G/D pattern, but applicable to actor-critic, EBM, contrastive multi-head, or any other multi-optimizer paradigm.
  • Trainer(model: NNModel).train(params, trainer_step_fn, callbacks=) — builds one torch.optim.Optimizer per entry in NNTrainerParams.optims, dispatches to a user-supplied trainer_step_fn(ctx) -> NNEvaluationDataPoint per batch, writes the same NNRun + per-tag NNCheckpoint artifacts as NNModel.train(). No default_trainer_step — multi-optim updates are scenario-specific and silently running the wrong update is worse than requiring an explicit fn.
  • NNTrainerParams — frozen dataclass with optims: Mapping[str, NNOptimParams] (name-keyed multi-optim config), schedulers: Mapping[str, NNSchedulerParams] (one per optim, defaults to ReduceLROnPlateau when missing), plus the standard n_epochs / train_loader / val_loader / seed / save_phase_checkpoints / extra_metrics. Validates non-empty optims and that every scheduler key matches an optim key. state() keys sorted for deterministic run.id.
  • TrainerStepContext — frozen bundle passed into a trainer_step_fn: model, batch, optimizers (dict), schedulers (dict), extra_metrics, batch_idx, epoch_idx. The companion TrainerStepFn type alias is exported.
• Strict param_groups semantics for multi-optim — build_param_groups(..., strict=True) (new keyword) drops parameters that match no spec instead of bucketing them into a default group. Threaded through Optims.__call__(..., strict_param_groups=True). The Trainer passes True so disjoint optimizers don't co-own parameters via implicit default buckets. Default strict=False preserves the fine-tuning semantics introduced by nnx.finetune.param_groups exactly.
• NNRun.trainer: Optional[NNTrainerParams] — populated by the Trainer; None for NNModel.train() runs. Strict back-compat: OMITTED from state() when None so existing NNModel run.id hashes are unchanged. NNRun.load(id) round-trips trainer-mode runs by lazy-importing NNTrainerParams.from_state when the YAML carries a trainer block.
• Runnable GAN demo: examples/09_gan_with_trainer.py — generator + discriminator packed into one nn.Module, two disjoint optimizers scoped via NNParamGroupSpec(name_pattern="G.*" | "D.*"), alternating updates on a 1D mixture-of-Gaussians. Verified end-to-end on CPU.

Deferred from this PR: trainer-mode warm-resume. The Trainer writes only the model net's state_dict to its NNCheckpoints — there is no per-optimizer .opt.<name>.pt sidecar yet. NNTrainerParams does not carry resume_from_run_id / resume_from_checkpoint. Resuming a GAN's Adam state for both G and D will land as its own follow-up PR once the use case is exercised.

Added — fine-tuning infrastructure (freeze / unfreeze / param_groups)

• nnx.finetune package with three submodules:
  • freezing — freeze(module, *patterns) / unfreeze(module, *patterns) / frozen(module). Glob-pattern (fnmatch) toggling of requires_grad on submodule parameters; the standard transfer-learning idiom. NNModel.freeze / NNModel.unfreeze are convenience methods delegating to the free functions.
  • loading — load_pretrained(module, source, *, key_map, strict, prefix) returns a LoadPretrainedResult with loaded_keys / missing_keys / unexpected_keys. Sources: file paths (loaded with weights_only=True for safety), state-dicts, or other nn.Modules. Key remapping handles foreign naming conventions (torchvision / HuggingFace / etc.).
  • param_groups — NNParamGroupSpec (frozen, kw_only, slots dataclass) for declarative per-layer LR / weight_decay overrides. The fine-tuning idiom of "small LR on the backbone, large LR on the head" expressed as a list of specs on NNOptimParams.param_groups. build_param_groups(module, specs, default_lr, default_weight_decay) is the helper the Optims enum factory dispatches through.
• NNOptimParams.param_groups: Optional[list[NNParamGroupSpec]] field. When set, the optimizer factory builds per-group dicts with the spec's lr / lr_multiplier / weight_decay overrides; frozen parameters are dropped. Strict back-compat: param_groups=None (default) is OMITTED from state(), so existing run.id hashes are unchanged.
• NNModel.export_state_dict(path) — saves self.net.state_dict() to disk as a plain torch file (no NNCheckpoint wrapper). Companion to load_pretrained for the round-trip.

Added — train_step_fn hook on NNModel.train() (foundational)

• train_step_fn hook on NNModel.train(). One optional kwarg that swaps out the supervised forward/backward/step for any user-supplied function. Unblocks non-supervised training paradigms (autoencoder, VAE, link prediction, recommendation, diffusion) without modifying NNx core. Default-None path is byte-identical to the prior loop. New public surface: TrainStepContext (frozen dataclass carrying model/batch/optimizer/scaler/grad_clip_norm/extra_metrics/accumulate_grad_batches/batch_idx/epoch_idx), default_train_step(ctx) (the standard supervised step, exported for users who want to layer behavior on top), TrainStepFn (type alias). Seven tests in tests/test_train_step_hook.py; runnable autoencoder example at examples/05_custom_train_step_autoencoder.py.
• Public alias for nnx.PredictResult (was reachable only via nnx.nn.nn_model).

Changed — internal

• NNModel.__fwd_pass → NNModel._fwd_pass. Required so the free default_train_step can reach it without Python name-mangling. Single underscore is still "weak private"; no external consumer touched the mangled _NNModel__fwd_pass name.
• NNModel._train_step becomes a one-line wrapper around default_train_step for back-compat with any hypothetical subclass that overrode it. The train() loop itself no longer dispatches through _train_step.

Fixed

• _save_checkpoints / _step_scheduler / _update_tqdm_postfix now tolerate an NNEvaluationDataPoint with error=None. Custom train_step_fn hooks for non-supervised paradigms (VAE/autoencoder/diffusion) don't always have a classification error to report; the loop falls back through val_edp.error → val_edp.loss → train_edp.error → train_edp.loss and skips the scheduler step entirely if nothing is set. Previously these three sites crashed with TypeError on None < float / float(None) / f"{None:.4f}".

Deferred

• eval_step_fn / predict_fn — same pattern, but evaluate() and predict() still assume supervised classification. First ml-lab task that needs custom eval (autoencoder, VAE, DDPM) will drive that.
• Network registry (Nets.register(...)) — each new architecture lands a Nets enum variant via its task's PR.
• Loss registry — custom losses live inside train_step_fn today (the user computes the loss tensor manually). Lift to a registry when multiple tasks duplicate the same custom loss.

[Pass-2 unreleased] — comprehensive improvements pass 2

Second improvement pass on branch chore/comprehensive-improvements-pass-2, building on pass-1. Strict back-compat preserved throughout — every new field on a params dataclass defaults to its old value and omits itself from state() when the default holds, so existing run.id hashes are unchanged.

Added — features (warm-resume, gradient accumulation, custom epoch checkpointing, etc.)

• Warm-resume training. NNTrainParams.resume_from_run_id and resume_from_checkpoint load weights AND optimizer state from a prior run's checkpoint at the start of train(). Optimizer state is written as a .opt.pt sidecar so the existing pickled NNCheckpoint format is untouched.
• Gradient accumulation via NNOptimParams.accumulate_grad_batches (default 1). Loss is scaled by 1/N; zero_grad/optimizer.step fire on cycle boundaries; AMP unscale + grad-clip both honor the cycle.
• TensorBoardCallback and WandbCallback — stream per-epoch train/val metrics + LR. Lazy import so users not on the path don't pay the dep cost.
• NNModel.to_onnx(path, example_input) — export the network via the legacy torch.onnx.export tracing path (no onnxscript needed). Marks dim-0 dynamic by default.
• NNTabularDataset — wraps a pandas DataFrame into train/val/test loaders matching the NNDatasetBase contract.
• Custom metrics via NNTrainParams.extra_metrics={name: fn}. Each fn(Y, Y_hat) -> float populates the new NNEvaluationDataPoint.extra dict; survives the NNRun.save/NNRun.load round-trip via extra.<name> CSV columns.

Added — reproducibility (seeded RNGs + env snapshot in metadata.yaml)

• nnx.set_seed(seed, strict=False) pins Python random, NumPy, torch CPU+CUDA, and cuDNN. strict=True also calls torch.use_deterministic_algorithms(True).
• nnx.dataloader_worker_init_fn — pass to DataLoader(worker_init_fn=...) for per-worker deterministic seeds.
• NNTrainParams.seed runs set_seed at train() entry; included in state() only when set.
• nnx.env_snapshot() captures library / torch / numpy / python / platform / CUDA / git-commit info. Written by NNRun.save() to runs/<id>/metadata.yaml — separate from run.yaml so it does NOT contribute to run.id.

Added — API ergonomics (predict tuple unpack, file= kwargs, NNCheckpoint helpers)

• NNModel.predict(X) accepts numpy.ndarray, torch.Tensor, tuples thereof, or a DataLoader (labels in batches are discarded). Returns a PredictResult NamedTuple that unpacks positionally as (logits, classes) for back-compat.
• NNTrainParams.save_phase_checkpoints: bool = True. Set False to skip the FIRST + Q1/Q2/Q3 cycle (LAST + BEST still always saved) — useful for tiny experiments or huge models.
• Devices.torch_device() / Devices.get_torch_device() return torch.device directly without the .() dance.
• Utils.print_tree / print_table accept file= for output redirection.
• nnx.__version__ resolves from importlib.metadata; falls back to "0.1.0+local" when editable-installed.
• pyproject keywords expanded (training, checkpointing, callbacks, experiments, reproducibility, neural-networks, research).

Added — reliability (NaN-loss guard, gradient clipping, atomic NNRun.save)

• NaN/Inf guard in NNModel._train_step — raises FloatingPointError rather than letting divergence corrupt checkpoints silently.
• Gradient clipping via NNOptimParams.grad_clip_norm: Optional[float]. AMP-aware (unscales before clipping).
• Incremental persistence — NNRun.save() runs after every epoch, not just at the end. KeyboardInterrupt / OOM mid-training now leaves a loadable partial run.
• SECURITY note on NNCheckpoint.from_file calling out the arbitrary-code-execution risk of weights_only=False on untrusted files.
• Re-pin loss_fn to self.device on every evaluate() call (guards against late device reassignment).

Fixed — correctness (evaluate aggregation, NNOptimParams.is_valid, callback isolation)

• NNOptimParams.is_valid() now returns False (not implicit None) for unknown enum variants — invalid configs no longer slip past the not params.optim.is_valid() pre-flight check.
• NNModel.train() tolerates DataLoaders without __len__ (IterableDataset-backed). Falls back to a tqdm bar with no total.
• NNRun.save() falls back to writing best/POINTER.txt when os.symlink raises (Windows without developer mode).
• NNModel.evaluate() aggregates Y / Y_hat across batches and computes metrics once on the aggregate, fixing unequal-final-batch weighting. Raises ValueError on an empty loader instead of returning NaN.
• NNIterationDataPoint gets a docstring spelling out that val_edp is populated only on the LAST idp of each epoch — readers shouldn't expect it on every row.

Changed — tooling (pyproject extras, conftest hygiene, type-checker config)

• CI runs pytest under coverage (pytest-cov), uploads coverage.xml artifact on Python 3.11.
• CI runs pyright in basic mode (continue-on-error: true today; will tighten to --strict over time).
• NNX_TQDM_DISABLE=1 silences the training progress bar — autouse'd in tests/conftest.py so pytest output stays clean.
• tests/conftest.py exposes shared fixtures (tiny_model, tiny_classification_loaders, tmp_runs_root, ...).
• mkdocs.yml + docs/ skeleton (index, quickstart, concepts, api). New .github/workflows/docs.yml builds with --strict on every push and deploys via mkdocs gh-deploy on main. New thekaveh-nnx[docs] optional extra.
• .pre-commit-config.yaml with ruff + standard pre-commit-hooks.
• CONTRIBUTING.md covering setup, workflow, back-compat invariants, testing.
• .github/ISSUE_TEMPLATE/{bug_report,feature_request}.md + pull_request_template.md.
• .github/workflows/release.yml — tag-triggered build + PyPI publish via OIDC trusted publishing.

Added — docs (concepts.md / quickstart.md / api.md)

• examples/ folder with four runnable scripts: 01_synthetic_classification.py, 02_resume_training.py, 03_custom_metrics.py, 04_onnx_export.py. All verified end-to-end on CPU.

Internal (Utils back-compat shim, vis_utils module aliases)

• Utils.print_tree / print_table / flatten_dict are now module-level functions in nnx.utils. The Utils class is a thin shim binding the same functions as staticmethods, so existing Utils.method(...) callers continue to work with no semantic change.
• VisUtils plotting helpers get module-level aliases (from nnx.vis_utils import confusion_matrix works).

Additional fixes (post-initial-pass)

• runs/best POINTER.txt fallback wasn't read during BEST comparison; env_snapshot subprocessed git on every save; .gitignore missed runs/, tb_logs/, *.onnx, coverage.xml, site/.
• Critical: NNOptimParams.state() unconditionally emitted grad_clip_norm=None, changing every existing run.id hash. Plus callbacks.py top-level IPython import (pulling IPython into every import nnx); NNRun.all() crashed on missing runs/ and tried to load stray files.
• mkdocs build --strict had 4 warnings (specs in docs but not nav; griffe couldn't parse a docstring; missing type annotation on to_onnx.example_input).
• Critical: NNEvaluationDataPoint.extra didn't actually round-trip through idps.csv. json_normalize flattened the dict on save but NNIterationDataPoint.from_state never reassembled the train_edp.extra.* columns. The pass-2 claim that "extra survives idps.csv" was false until this fix.
• pytest-cov listed in dev extras but not installed locally. CI handles via pip install -e ".[dev]"; surfaced via cov-run on a fresh venv.
• NNEvaluationDataPoint.mean_of silently dropped the extra dict from inputs. NNCheckpoint.load_optimizer_state now uses weights_only=True (the state dict is structured tensors + dicts; the strict loader works AND removes the ACE risk).
• Six conftest fixtures (tiny_model, tiny_classification_loaders, etc.) defined but unused — premature abstractions deleted; CONTRIBUTING.md updated to match.
• NNTabularDataset now validates feature_cols / target_col against df.columns up-front with a clear KeyError; new test for env_snapshot cache (introduced in R1 but never explicitly tested).
• stray leading blank line in nn_graph_dataset.py.
• Real recovery gap: NNRun.save()'s three writes (run.yaml, metadata.yaml, idps.csv) were non-atomic. A Ctrl-C mid-write left half-written files. New _atomic_write_text helper does tmp + fsync + os.replace.
• NNCheckpoint.to_file had the same non-atomic gap (torch.save direct to destination). New _atomic_torch_save helper applies the same tmp + rename pattern to both the main checkpoint and the .opt.pt sidecar.
• Atomicity also applied to the Windows POINTER.txt fallback; helper reordered (defined before its caller); pyproject filterwarnings for the upstream torch_geometric.distributed / torch.jit.script DeprecationWarnings; fix the scheduler test's optimizer-before-scheduler step order so the runtime UserWarning doesn't fire.
• README "Other models" was a non-functional snippet (imported classes without showing how to wire them through NNModel). Replaced with concrete NNModelParams(net=Nets.GRAPH_*) examples + a pointer at the examples/ folder. Added README subsections for Reproducibility, Warm-resume, and Custom metrics so the pass-2 features are visible from the top-level doc.
• test_imports.py was missing smoke imports for nnx.seeding, nnx.nn.callbacks, nnx.nn.net.graph_nn_base, nnx.nn.dataset.nn_tabular_dataset, and nnx.nn.enum.schedulers. The test predated pass-1 and never grew with the codebase. Closed the gap so the cheapest-possible refactor signal is exhaustive again.
• release.yml skipped twine check between python -m build and the PyPI upload step. A malformed README or invalid classifier would only surface when PyPI rejected the upload — by then the tag is burned. Added a twine check dist/* verification step; also added cache: pip to the setup-python step for parity with the other workflows.
• Final sweeps: ran the literal README quickstart end-to-end, manually exercised the four predict() input forms (ndarray, tensor, tuple-of-each), verified all internal markdown links resolve, and confirmed mkdocs build --strict is silent. No additional actionable findings.

Deferred (with rationale)

• D3 (split NNModel.train() into a TrainingLoop runner): the existing helpers (_train_step, _save_checkpoints, _step_scheduler, _build_scheduler, ...) already break the loop body into testable units. A full extraction would be churn without proportional value.
• D7 (versioned state-dict checkpoint format with a versioned reader): too risky for this back-compat pass. The pickled NNCheckpoint continues to work; the weights_only=False security note in the docstring guards against the supply-chain risk.
• D8 (Storage protocol for cloud backends): broad I/O abstraction touching every save/load site. Better as its own focused PR.
• N5 (md5 of str(state) → json.dumps(sort_keys=True)): would change every existing run.id. Can't ship under strict back-compat.
• O5 / O6 / O7 (NNTrainParams config-vs-runtime split, callbacks-as-params, NNModel __init__ param rename): API breaks. Deferred.
• O4 (frozen _CallbackContext view): would change the surface callbacks can mutate — defer to a callback API revision.
• P1 / P2 (per-batch device sync, loss.item() sync): would sacrifice per-batch metric granularity (idp.train_edp). Deferred.
• E7 (move ipython/kaleido to optional extras): would break pip install thekaveh-nnx for users relying on the default extras. Deferred.

[Pass-1 unreleased] — comprehensive improvements pass 1

The pass-1 series landed on branch chore/comprehensive-improvements-pass-1. Strict back-compat preserved: no public API renames, no on-disk format breaks, deep imports still resolve.

Fixed — correctness

• NNDataset now carves the validation slice out of the source train=True split, keeping the source train=False split intact for final evaluation. Previously val was a slice of test, leaking the test pool. Reported val metrics will differ between pre/post versions.
• NNDataset random_split sizes are computed as (total - val, val) instead of two truncated halves. Fixes the crash on odd-length source train sets.
• NNRun.save no longer crashes when comparing against a prior BEST run that has no val_edp (e.g., a no-validation experiment). A new _best_err helper falls back to train error, then +inf.
• NNEvaluationDataPoint.of now defaults average="macro" for f1 / recall / precision. The prior "micro" hardcoding made all three numerically identical to accuracy for single-label multi-class tasks. Pass average="micro" to opt back in.
• VisUtils.multi_line_plot: removed dead cs = px.colors.qualitative.Plotly[...] assignment that was immediately overwritten; replaced the ls[:len(ys)] legend loop (which depended on a leaked inner-loop variable) with n_lines_per_series = len(yss[0]); raises ValueError on empty yss.
• Activations.SOFTMAX returns a closure that supplies dim=-1, avoiding the implicit-dim warning and ambiguity from F.softmax.
• NNModel._train_step: detach train_loss before float() to avoid UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.

Fixed — deprecations / future breakage

• Migrated torch.cuda.amp.{autocast,GradScaler} to torch.amp.* with explicit device_type="cuda". The torch.cuda.amp module has been deprecated since torch 2.4.
• NNCheckpoint.from_file calls torch.load(weights_only=False). Without it, torch ≥ 2.6 (where weights_only defaults to True) raises UnpicklingError on any saved NNCheckpoint (checkpoints pickle the full Python object, not a bare state dict).
• NNGraphDataset reads the underlying Data via dataset[0] instead of dataset._data. The private accessor was renamed/removed across PyG versions.
• Removed the top-level from IPython.display import clear_output import in nn_model.py. The actual use is in callbacks._LegacyCallback; leaving the top-level import made every consumer of nnx.nn.nn_model pull in IPython.

Added — initial release scaffolding (top-level re-exports, persistence root, viz figures)

• nnx/__init__.py re-exports the curated public surface (NNModel, params, callbacks, enums, nets, datasets, utils) with an explicit __all__. Deep imports (from nnx.nn.net.feed_fwd_nn import FeedFwdNN) still work for existing code.
• NNRun.save / load / all / checkpoints and NNCheckpoint.save / load accept an optional root: Optional[str] = None kwarg. Default is unchanged (cwd-relative); callers wanting to redirect persistence can now pass one.
• NNEvaluationDataPoint.of accepts average: str = "macro".
• VisUtils.{multi_line_plot, scatter_plot, two_dim_tsne_checkpoint_logits, confusion_matrix} now return the plotly.graph_objects.Figure they build. The .show() call is gated on a non-None renderer so headless test envs no longer crash.
• tests/test_params_round_trip.py — contract test asserting obj == from_state(state()) for every params dataclass. Fails loudly when fields drift.
• tests/test_train_integration.py — end-to-end NNModel.train() coverage on a tiny in-memory TensorDataset, plus NNRun.load round-trip and NNModel.from_checkpoint reconstruction.
• NNOptimParams.momentum docstring explaining the SGD-vs-Adam dual meaning.
• NNDataset docstring documenting that val is carved from train.
• This CHANGELOG.md.

Changed — tooling

• Ruff lint now selects E, F, W, B (bugbear), I (isort), UP (pyupgrade). Style-preserving ignores: E701 (case style), B024 (structural base class), UP007 / UP045 (keep Optional over X | None). 213 auto-fixes applied (mostly import ordering).
• CI matrix adds Python 3.12.
• CI ruff step no longer has continue-on-error: true — lint gates merges.

Internal

• nn_dataset_base.py: trimmed 9 unused imports.
• nn_model.py: removed empty class NNModel(): parens.
• nn_dataset.py: switched to a local resolved_batch_sizes so downstream loaders don't read self.batch_sizes while it still holds the default tuple.

[0.1.0] — 2026-05-18

Initial extraction from thekaveh/ml.