thekaveh-nnx 0.2.0__tar.gz

thekaveh_nnx-0.2.0/CHANGELOG.md

@@ -0,0 +1,479 @@
+# Changelog
+
+All notable changes to NNx are documented here. Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning is roughly [SemVer](https://semver.org/) — pre-1.0, we allow behavior changes (typically bug fixes) without renaming public APIs.
+
+## [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`](https://github.com/pytorch/ao) (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 `TransformerBlock`s (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 `NNCheckpoint`s — 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.Module`s. 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 `DataLoader`s 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`.