NNx Architecture

Lightweight PyTorch training / eval / visualization toolkit — public surface, specialization extensions, callback bus, on-disk persistence.

on_*_begin/end writes writes nnx package — pip install thekaveh-nnx Specialization subpackages Persistence (runs/<id>/ on disk) User Code import nnx PyTorch nn.Module / Tensor NNModel net + params + .train() Trainer multi-optim — GAN, A2C train_step_fn hook (optional) trainer_step_fn hook (required) nnx.finetune freeze, load_pretrained nnx.peft LoRA / DoRA / IA3 / Prefix / Prompt nnx.diffusion DDPM, sample() nnx.paradigms KD / SimCLR / MoE / JEPA / DPO nnx.trainer Trainer class _step_helpers finalize_step() Plus: nnx.quantize (PTQ / QAT) • nnx.prune (magnitude / 2:4) • nnx.surgery (Net2Net / drop / low-rank) • nnx.embeddings (contrastive + FAISS) • nnx.interop (GGUF / Ollama / safetensors) • nnx.generation (LogitsProcessor) • nnx.viz (model-internals) Training Loop epoch × batch dispatch NaN + grad-clip inline / _step_helpers _step_scheduler Schedulers enum _save_checkpoints FIRST/Q1/Q2/Q3/LAST/BEST Callback Bus • on_train_begin / on_epoch_begin / on_epoch_end / on_train_end EarlyStopping patience monitor LRMonitor LR history ModelCheckpoint custom epochs TensorBoard opt-in extra Wandb opt-in extra NNRun run.yaml + idps.csv + metadata.yaml runs/<id>/ md5(state)-keyed directory NNCheckpoint *.pt + *.opt.pt

Public entry surface

  • • NNModel — supervised orchestrator
  • • Trainer — multi-optimizer (GAN, actor-critic)
  • • Frozen params dataclasses round-trip via state()
  • • Curated re-exports at top-level nnx

Specialization subpackages

  • • finetune — freeze, load_pretrained, NNParamGroupSpec
  • • peft — LoRA / DoRA / IA3 / Prefix / Prompt / AdapterLayer
  • • quantize — quantize_int8 (PTQ) + QAT 8da4w via torchao
  • • prune — magnitude (unstructured) + 2:4 semi-structured
  • • surgery — widen / deepen / drop_layer / low_rank_factorize / expand_embedding
  • • diffusion — NoiseSchedulers, DiffusionMLP, sample()
  • • paradigms — KD / feature-KD / SimCLR / Mixup / CutMix / MoE / I-JEPA / DPO / Born-Again
  • • trainer — Trainer class (also surfaced as the cyan entry-point)
  • • embeddings — contrastive trainer + FAISS export
  • • interop — GGUF write + Ollama Modelfile + safetensors
  • • generation — LogitsProcessor chain (temp / top-k / top-p / repetition)
  • • viz — torchinfo summary + weight histogram + activation map + Captum + gradient flow + Netron
  • • _step_helpers — shared finalize_step (NaN + grad-clip)

Persistence (runs/<id>/)

  • • run.yaml — config (state() md5 = id)
  • • idps.csv — per-iteration metrics (incremental)
  • • metadata.yaml — env snapshot (NOT in id hash)
  • • checkpoints/*.pt — FIRST/Q1/Q2/Q3/LAST/BEST

NNx • Generated via architecture-diagram skill • SVG mirror at docs/assets/architecture.svg