Skip to content

20. 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]

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). Per feedback_examples_optional_extras_docstring 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 (per [[builder-pattern-shape]] §11b); the missing-n_epochs case fell through to the dataclass ctor, which is the one error path inconsistent with the Builder rubric's "catch errors at the Builder boundary with a fix-naming message." 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 [[pypi-keywords-with-features]] 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 contract from [[omit-when-default-state-invariant]]. 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 [[examples-seed-helper-override]] 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 [[examples-seed-helper-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 §11b 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 [[user feedback]] 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 the [[megamerge-pattern]] round-3 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.
  • §11b Builder-boundary pre-check pattern applied uniformly to NNOptimParamsBuilder / NNSchedulerParamsBuilder / NNTransformerParamsBuilder (post-PR-#52 overnight-maintenance pass). PR #52 established the [[builder-pattern-shape]] §11b convention 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 [[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.