torch-dimensions / PLAN.md
Celsia's picture
Upload folder using huggingface_hub
ecc81b3 verified
|
Raw
History Blame Contribute Delete
54.4 kB
# torch-dimensions β€” Build Plan
Companion to [DESIGN.md](DESIGN.md). That document says *what* the library is; this one says *in what order it gets built, how we know each step worked, and how much of the surface each step actually covers*.
**Sequencing principle:** each phase ends with something runnable and tested, and the riskiest unknowns get resolved before anything is built on top of them. Concretely that means the tensor bookkeeping (`Lattice`) is proven before any model exists, the test suite exists before the hard math, and the config surface is written **last** β€” a config schema authored before the block signatures settle just freezes a guess.
**Legend.** `[x]` done and verified (a test or a shipped artifact backs the tick β€” a tick with no evidence is not a tick). `[ ]` open. `[~]` partially done, with the gap stated inline. Sizes are relative (S β‰ˆ half a day, M β‰ˆ 1–2 days, L β‰ˆ 3–5 days, XL β‰ˆ 1–2 weeks) for one person already fluent in this material.
**Coverage** lines state what fraction of the surface a phase's ticks actually exercise β€” ranks, dtypes, devices, densities, model families β€” because "phase done" and "surface covered" are different claims, and conflating them is how silent gaps ship.
---
## Current position (2026-08-02, third pass)
**Track A:** Phases 0–11 complete except where noted. **9 and 10 closed this pass** β€” reproductions with numbers (sMNIST 99.53%, psMNIST 97.79%, a sparse-lattice baseline that had no published prior) and benchmarks that contradicted two of the design's own predictions. v0.1.0 is live on PyPI. Phase 12 (autoregressive stepping) is designed and is now the largest single piece of unbuilt core.
**Track B (viewer):** V1, V1.5 and **V2 done** β€” `td.viz.show(model)` ships inside the wheel. V3 (shape flow) is next and unstarted.
**Track C (performance):** the checklist exists; **no CUDA kernel has ever run**. This is the biggest gap between what the library claims and what it has demonstrated.
**Track D (ecosystem):** both extension-point guides written and executed by tests; no docs site, no notebooks, no tech report.
**Track E (research):** parked by design.
Next up, in order of value: **run the CUDA checklist** (fifteen minutes, and it converts the project's largest unverified claim), **three seeds per RESULTS.md row** (every row is currently n=1), **a ViT/CNN reproduction row** so the new families meet the same standard as the old ones, **Phase 12 AR stepping**, **Viewer V3**.
**Third pass added** the non-sequence families (Phase 7b), the first verification of our *composition* against a published N-D method (Phase 7c: S4ND at 1.8e-17): `td.CNN`, `td.TCN`, `td.ViT`, and `td.flatten` as a fourth method of multidimensionality β€” plus LTI.md, which measured a claim this plan had been making without evidence and found it wrong. Two more bugs documented (#26 the spec describing a model that never runs, #27 absent cells reaching the output through the residual stream), 27 total.
What the second pass changed, in one line each: 23 bugs now documented (four new, two found by *looking at what shipped* rather than at the build); ranks 5–6 tested so the README caveat could be deleted; the Kronecker conformance check runs for the first time; `td.Transformer` completes the family the README table always claimed; coverage floored at 95% (measured 97%).
---
# Track A β€” Core library
## Phase 0 β€” Skeleton and packaging Β· S Β· no deps
Make the repo installable and enforceable before any logic exists.
- [x] `pyproject.toml` (hatchling), `src/` layout so tests import the installed package rather than the working tree.
- [x] Dependencies: `torch>=2.4` required and **nothing else**. Optional extras: `[mamba]`, `[fla]`, `[dev]`, `[all]`. (An `[s4]` extra existed until the portable S4D landed in-tree and made it pointless.)
- [x] `ruff` (lint + format), `pytest`, `pytest-cov`. Type hints throughout; `mypy` advisory β€” *and executed in CI*, because "advisory" once quietly meant "never run" (DEBUG.md #7).
- [x] CI: CPU matrix on Python 3.10–3.12, full-history checkout (DEBUG.md's citation test needs it).
- [x] `LICENSE` (Apache-2.0), `README.md`, `CONTRIBUTING.md`, `CHANGELOG.md`.
- [x] Publishing workflow: tag-triggered, PyPI trusted-publisher OIDC, sdist-size guard (the 35 MB node_modules incident, now a CI check).
- [x] sdist scoped to the library (`src`, `tests`, `examples`, docs) β€” the sdist is the library, not the repo.
- [ ] Python 3.13 in the CI matrix once torch ships stable wheels for it.
- [ ] A `nightly` CI job against torch nightly, so upstream breakage is our alarm and not our users'.
- [ ] Windows runner (even one smoke job) β€” path handling and MPS-vs-CUDA device pick are the risks.
**Coverage:** packaging is fully exercised (clean-venv wheel install is tested; PyPI install is live). CI covers Linux CPU only β€” macOS/MPS runs locally, Windows never.
**Acceptance:** `pip install -e ".[dev]"` then `pytest` passes on a machine with no CUDA and no optional deps installed. βœ… Also now: `pip install torch-dimensions` from PyPI works cold.
---
## Phase 1 β€” `Lattice` Β· M Β· needs Phase 0
The foundation. Pure tensor bookkeeping, zero models, zero learnable parameters.
- [x] `shape`, `names`, `valid`, `time`; axis name β†’ position resolution; degenerate `shape=()` + `time=True` sequence lattice.
- [x] Permutation generation and inverse for "move axis `k` to sequence position"; inverse checked against `torch.argsort`, not hand-derived twice.
- [x] Flat indices, scatter-to-dense, gather-from-dense; broadcast validity mask; per-axis valid counts for masked pooling.
- [x] Immutable after construction (`__setattr__` refused) β€” a value object that is hashed from and cached from must be a value (DEBUG.md #2).
- [x] Defensive copies: `valid` cloned on entry, `mask()` returns a fresh tensor, never a view (DEBUG.md #13).
- [x] Device correctness: `flat_idx` indexes on the data tensor's device, both mismatch directions (DEBUG.md #18, proven on MPS).
- [x] Property/fuzz tests: random shapes and masks, ranks 1–4, fold/scatter/permutation round-trips against independent references.
- [x] Extend the fuzz envelope to ranks 5–6 β€” done, and both families now pass **full conformance at ranks 5 and 6** with a rank-5 model trained end to end. The README caveat is deleted because it stopped being true.
- [x] Stress shapes: every axis length 1, a single valid cell, a 10,000-long axis β€” explicit cases rather than fuzz luck (`test_stress_shapes_that_fuzz_would_have_to_be_lucky_to_draw`).
- [x] `Lattice.sliced(**axes)` / `Lattice.merge` β€” sub-lattice views for splits over *space*. `sliced` returns the lattice **and** the tensor selection as one object, so the two cannot drift apart; integer indices are refused because a rank that changes with a train/test split is not a rank.
- [ ] Serialization guarantee documented: a lattice pickled/JSON-round-tripped on one torch version loads on the next (add a stored-fixture test).
**Coverage:** ranks 1–6 fuzz-tested against independent references; dense and sparse; CPU + MPS; float32/float64; degenerate and extreme aspect ratios as explicit cases. Not covered: pickling across torch versions, ranks β‰₯ 7 (untested, and now merely untried rather than claimed).
---
## Phase 2 β€” `ScanPlan` Β· S Β· needs Phase 1
Pure data. No tensors, no modules.
- [x] `.cyclic()`, `.paired()`, `.from_list()`; `__repr__`; to/from dict; name-or-index axes resolved against a lattice; warning when a plan leaves an axis unswept.
- [x] Per-axis `bidirectional` (time stays causal while space does not); the `set("time")` string-iteration trap guarded.
- [x] Direction flips per *cycle*, not per layer β€” the even-axis-count aliasing that silently pins every axis one way (DEBUG.md #4; exists in published research code).
- [x] `_warn_if_pinned`: a layer budget that cannot deliver the requested bidirectionality warns instead of silently downgrading.
- [x] Immutable and hashable together (DEBUG.md #1).
- [x] Plan algebra: `+`, `* k`, `.reversed()` (layer order) and `.flipped()` (sweep direction), named apart because calling either one "reversed" alone is how they get confused. The identity that earns it: `p + p.flipped()` is bidirectional everywhere `p` swept at all.
- [x] `plan.coverage(lattice)` β€” per-axis sweep counts, directions, layer indices, `unswept` and `pinned`. It never warns, because a report that warns cannot be used to decide whether to warn; `spec()` now derives its sweeps section from it and gained `pinned_axes` for free.
- [ ] Schedule catalog: named constructors for the published schemes beyond Mamba-ND's paired (e.g. S4ND's simultaneous-separable as a degenerate plan, zigzag variants) with citations. **Read LTI.md first:** "simultaneous" is only equivalent to "sequential in any order" for scalar-valued filters, so a degenerate-plan constructor has to say which order it collapses to.
- [ ] Property test: for every constructor, every axis mentioned is swept β‰₯ 1 time or the warning fires β€” fuzz over axis counts 1–8 and layer counts 1–64.
**Coverage:** constructors, resolution, serialization, algebra, coverage reporting, and the aliasing signature are fully tested (even *and* odd axis counts β€” one parity finds nothing). Not covered: exotic published schedules.
---
## Phase 3 β€” `AxialScan` + RNN family Β· M Β· needs Phases 1–2
- [x] `axial_apply`: permute β†’ fold β†’ mixer β†’ unfold β†’ inverse permute; `reverse`; `chunk` (validated β‰₯ 1); shape-contract error naming `(M, A, H)`.
- [x] `AxialScan`: pre-norm residual per layer, per-layer or shared mixers, absent cells zeroed on entry **and** after every layer.
- [x] `LSTM`, `GRU` β€” one class each, 1-D without a lattice, N-D with one; no `LSTMND`.
- [x] Single layer on a rank-1 lattice matches `nn.LSTM` **bit-for-bit**; multi-layer matches to float tolerance (fold normalizes memory layout; torch RNN kernels are layout-sensitive β€” expected, not a defect, and documented).
- [x] Verified against independent references: `cumsum`, position-weighting β€” mixers that torch itself can check.
- [ ] Batch-fold chunking auto-tuned from device limits rather than user-supplied. Phase 10 now says what the tuner should do: **chunking never won on time** here β€” `chunk=64` cost 3Γ— by launching 3Γ— the kernels β€” so it is a memory-pressure valve, and the default should be off until an allocation actually fails.
- [x] `td.testing.Recorder` β€” a mixer that computes nothing and records every call, droppable into a real model via the new `mixer=` argument. That argument needed a guard: a substituted mixer cannot go into the recipe, so `save()` refuses rather than silently rebuilding a *different* model that loads without complaint.
- [ ] Gradient-checkpointing option per layer (rank-4 lattices at real `d_model` will need it; measure in Phase 10 first).
**Coverage:** ranks 1–4, dense + sparse, time on/off, forward + reverse + chunked, float64 gradcheck, MPS parity. Not covered: activation checkpointing, mixed precision (Track C), very long folded batches near memory limits.
---
## Phase 4 β€” Conformance suite Β· M Β· needs Phase 3
The single highest-leverage phase: it is what keeps a (models Γ— ranks Γ— sparsity Γ— devices) matrix from rotting.
- [x] `td.testing.check_block` β€” the seven checks: shape, gradients + gradcheck, rank-1 equivalence vs the bare mixer, Kronecker identity (kernel family), absent-cell inertia, axis-storage covariance, `torch.compile` (opt-in).
- [x] Skips are recorded, never silently passed (`_Skip` is the mechanism; "we never ran that one" can never read as "that one passed").
- [x] Checks run only at ranks the caller requested (DEBUG.md #16).
- [x] `check_trainable`: fresh data per step, held-out scoring, and a negative control that must fail β€” a learning test without one measures capacity, not learning (DEBUG.md #5).
- [x] Fuzz suite (`test_fuzz.py`): seeded, checked against slow independent references, mutation-verified.
- [x] Device suite (`test_device.py`): runs against whatever accelerator exists (MPS locally, CUDA elsewhere), skips visibly otherwise.
- [x] Check #4 β€” the Kronecker identity **runs now**. It was an unconditional skip since it was written; `check_block(kernels=...)` takes an adapter returning the operators a block actually used, and CaFA passes at ranks 2–3 with a negative control that fails.
- [ ] Check #8 β€” parallel/sequential equivalence (lands with Phase 12; specified there).
- [x] Mutation testing as a scheduled CI job (`scripts/mutate.py`, weekly). Seven mutations, each one a bug this project actually had, each naming the tests that must fail; **7/7 caught**. A survivor is reported as a hole in the *tests*, not a bug in the code. It refuses to run on a dirty `src/` β€” a crash mid-mutation must never be confusable with someone's uncommitted work.
- [x] Coverage floor in CI at 95%, measured at **97%** across 447 tests. `pytest-cov` had been installed and unused since Phase 0 β€” DEBUG.md #7's shape, one tool over.
- [x] `td.testing.check_data_source(source)` β€” five checks including slice-seam consistency and picklability, both with negative controls. The pickle check matters most: DEBUG.md #9's version of that failure *hung* rather than raised.
- [x] Golden-file specs, one per family, compared verbatim with a field-path diff and a `TD_UPDATE_GOLDEN=1` regeneration path. Every previous spec test asserted particular keys, so a silently added or renamed field passed them all and reached the viewer as a blank panel.
**Coverage:** every shipped block passes all applicable checks at ranks 1–3 (RNN + SSM + kernel + attention families) and ranks 5–6 for the scan and kernel families, dense and sparse, CPU float64 + accelerator float32; data sources and spec documents have their own checkers. Not covered: automated mutation runs, fp16/bf16 (Track C1).
---
## Phase 5 β€” Data Β· M Β· needs Phase 4
Building a lattice from data is lattice construction and belongs here; running a training loop is not and never will be.
- [x] `from_coords` β€” vocabularies, shape, valid mask inferred from observed tuples; unknown values raise; encode/decode round-trip fuzz-tested.
- [x] `from_table` β€” long-format rows to `(T, *shape, F)`; duplicate `(time, cell)` rows refused (a join bug must not become a plausible dataset).
- [x] `LatticeWindow` β€” pure index arithmetic; split drops straddling windows on both sides; unsorted timestamps refused (a silent nonsense split is the quietest possible leakage bug β€” DEBUG.md #15).
- [x] `LatticeSource` protocol + `TensorSource`; `LatticeDataset`; `collate_lattice` (ragged refused; mixed target presence refused β€” DEBUG.md #14; lattice kept out of the batch).
- [x] Multiprocessing-safe: samples and batches pickle; `DataLoader(num_workers>0)` proven end-to-end (DEBUG.md #9 β€” it used to *hang*).
- [x] `td.data.MemmapSource` β€” the on-disk reference, written to fail the way on-disk sources fail: the handle opens lazily and is dropped on pickling, because a live mmap in a `DataLoader` worker hangs rather than raises. Passes `check_data_source`.
- [ ] A zarr source behind an extra β€” one on-disk implementation is a demonstration; two would be a pattern.
- [x] `td.data.masked_stats` / `Normalizer` β€” statistics over present cells only. A plain `series.mean()` on a 30%-sparse lattice is dragged 30% toward zero by structural zeros, the scale with it, and nothing about the resulting model looks wrong. NaN is treated as absent too.
- [ ] Ragged-time policy documented and tested: what a source with per-cell history lengths should do (today: build the union lattice and mask; write the recipe down with a test).
- [ ] Streaming/windowed iteration for series too large for memory (windows over a source that only supports sequential reads).
- [ ] An end-to-end "CSV to trained model" example with a real public dataset (feeds Phase 9 and Track D).
**Coverage:** in-memory and memory-mapped paths tested including worker processes; property-tested windowing; masked normalization with its failure mode as an explicit test. Not covered: zarr, ragged series, >memory scale.
---
## Phase 6 β€” Kernel family Β· M/L Β· needs Phase 4
- [x] `axial_contract` + `kron_operator`; the factorization checked against the explicit Kronecker product, not against itself.
- [x] Sparse renormalization with the **relative** cancellation guard β€” degeneracy is cancellation, and cancellation is relative to the absolute mass (DEBUG.md #3 β†’ #12: the first fix repeated the bug one level up; float32 blew up 7,000Γ— under an absolute epsilon).
- [x] No NaN laundering: an input NaN must leave (DEBUG.md #11).
- [x] `AxialKernel`: per-line scores (axial attention) or pooled per-axis kernels (CaFA); learned relative-position bias per axis; softmax or leaky_relu gate; per-line renormalization *is* masked softmax for the softmax gate.
- [x] The hybrid form: kernels own space, the model's mixer owns time; CaFA pools over *other spatial axes only* β€” a kernel at time t built from the future would leak through a causal model, and the causality test holds the hybrid to bitwise on the past.
- [x] A mixer on a time-less lattice refused as dead weight.
- [x] `td.axial_attention`, `td.cafa` as strategies; registered by name; `method=` as the short spelling of `nd_method=`.
- [x] Module-level Kronecker equivalence for CaFA β€” closed, via the `kernels=` adapter in `check_block`. Worth recording *what* the adapter returns: CaFA pools the current activation, so the second axis's kernel is built from the first axis's output. The factors are the ones actually applied, not the ones a static reading predicts.
- [ ] Multi-head kernels (currently single-head per axis) β€” heads are the difference between this and what a Transformer person expects to configure, and the CaFA reference carries them, so this is now a confirmed gap rather than a suspected one. The contraction is the work: a per-head kernel means contracting feature slices independently, which `axial_contract` takes one call per head to express.
- [x] `qk_norm` and `kernel_residual` (`K + gamma*I` before the gate), both from CaFA's `LowRankKernel` and both defaulted off so existing models are bitwise unchanged. The residual is the interesting one: it starts a contraction near "keep your own value" instead of fully mixed.
- [x] `td.Transformer` / `AttentionMixer` β€” attention as the swept mixer, which makes "N-D Transformer" literal and gives both constructions in the literature a name. Non-causal by default (a mixer is not told which axis it sweeps, so masking "the future" of a spatial axis is meaningless); the causal mask is tested bitwise in both directions.
- [x] Cost model, measured rather than derived: BENCHMARKS.md's "Where factorization starts winning". The answer was **not** what this line assumed β€” per-line attention is faster everywhere below ~50 cells per axis, and factorization only leads at 64Β³.
- [ ] `gate="softmax"` temperature / learned-per-axis option (CaFA paper ablates this).
**Coverage:** both strategies pass full conformance at ranks 1–3 and 5–6, dense + sparse, both gates, learnability, hybrid causality, the module-level Kronecker identity, MPS. Not covered: multi-head kernels, gate temperature.
---
## Phase 7 β€” SSM family Β· L Β· portable core done; fast paths open
**Portable core (done):**
- [x] `S4DMixer` β€” diagonal kernel, **bitwise identical** to upstream's with copied weights.
- [x] `S4Mixer` β€” the full DPLR kernel: HiPPO-LegS NPLR init, Cauchy resolvent, rank-1 Woodbury, bilinear transform with the **Nyquist-pole guard built in** (upstream survives that pole only by rounding luck; on MPS it NaNs β€” recorded below). Matches upstream at 3e-8; matches a dense state-space matrix-power reference at 4.6e-16 **in CI**.
- [x] `MambaMixer` β€” gated conv + selective scan, mirroring the reference non-fused path; scan matches `selective_scan_ref` at 1e-6.
- [x] `td.S4`, `td.S4D`, `td.Mamba` via the shared `LatticeModel` base; explicit N-D names `S4ND`/`S4DND`/`MambaND` with mandatory, checked `dim` (`dim=1` refused by name β€” one spatial axis is the 1-D model, and code reading "S4ND" must not be running S4).
- [x] Causality tests (S4 causal to 2.6e-15 through the FFT β€” mathematically causal, not bitwise, and the test says exactly that; Mamba held to bitwise).
- [x] Full conformance, learnability, device parity for all three.
**S4/Mamba portability dossier (established empirically, 2026-07-30, MPS):** the upstream repos run on Mac once vendored copies drop `pytorch_lightning`/hub/`transformers` imports and fix the `_omega` Nyquist pole β€” full S4Block at 1.13e-6 CPU↔MPS parity, Mamba v1 at 2.98e-8, Mamba2 via an `ssd_minimal` adapter at 4.77e-7 (adapter verified against a sequential recurrence at 8.9e-15 over 40 ragged-length cases). Import shims: stub `selective_scan_cuda`, stub triton *after* torch imports with `.configs`-bearing autotune, stub the hub mixin. **Mamba3 ships no reference implementation** β€” every path is a 712-line triton kernel; a portable version means writing the rotary trapezoidal recurrence from the paper.
**Fast paths and breadth (open):**
- [ ] `mamba-ssm` fast-path adapter: when installed and on CUDA, `MambaMixer` delegates to the fused kernel; the conformance suite runs *both* paths and asserts they agree β€” the shim harness from the portability work is the seed.
- [ ] `causal-conv1d` fast path for the conv (same pattern, same agreement test).
- [ ] `Mamba2Mixer` (SSD): portable chunked implementation β€” the verified `chunk_scan_ref` adapter *is* this, one packaging step away β€” plus the triton fast path when available.
- [ ] `Mamba3Mixer`: write the pure-torch reference from the paper (own task, XL); verify on a CUDA box against upstream triton; then it joins the family.
- [ ] S4 options parity where they earn their keep: rank-2 low-rank correction, `legt`/`fourier` measures, learnable-vs-fixed `dt` per the paper's ablations β€” each lands only with a test against the dense reference.
- [ ] `fla` (flash-linear-attention) adapter behind the existing `[fla]` extra β€” currently the extra installs a dependency nothing uses; that is a small dishonesty with a deadline.
- [ ] Sequential-scan speedup for the portable Mamba path: block-parallel (associative) scan in pure torch β€” the O(A) python loop is correct and honest, and a log-depth scan would make the portable path usable for long 1-D sequences too.
**Coverage:** the three shipped mixers are verified against upstream references *and* independent in-repo references, on CPU/MPS, ranks 1–3, all conformance checks. Not covered: any fused kernel, Mamba-2/3 as mixers, CUDA execution anywhere.
---
## Phase 7b β€” Convolutional and vision families Β· M Β· **done this pass**
Added because the abstraction's claim is about *lattices*, not sequences, and until now every mixer in the library was a sequence model. These two are the falsification test: a convolution has no state and no notion of "so far", and a Vision Transformer's whole trick is *not* factorizing the grid. Both fit without either side bending.
- [x] `ConvMixer` β€” 1-D convolution along the swept axis; causal or centred; optional depthwise-separable; `activation=None` gives a strictly LTI operator. Even kernels refused when centred (a window with no centre would shift the output half a step along one axis).
- [x] `TCNMixer` / `td.TCN` β€” the published TCN defaults, causal and bitwise so. **Dilation grows per axis, not per layer**: layer 6 of a rank-3 cyclic plan is the third sweep of *its own* axis and dilates by 4, not 64. `AxialScan` passes `sweep=` only to factories whose signature names it, which is what kept the feature invisible to every other mixer β€” and what let `TCNMixer` swallow it into `**kw` and run at dilation 1 everywhere until a test looked at the schedule.
- [x] `td.CNN` β€” the separable N-D convolutional network, proven equal to a full N-D convolution with a rank-1 kernel (`F.conv2d`/`F.conv3d`, 1e-12, negative control included).
- [x] `td.receptive_field(model)` β€” per-axis span against axis size. Refuses the kernel family explicitly rather than duck-typing to a confident wrong number.
- [x] `td.flatten` β€” the fourth nd_method: no factorization, every axis in one sequence. The baseline the axial methods exist to beat, and the only composition where a sparse lattice is *cheaper* rather than merely masked.
- [x] `td.ViT` + `td.PatchEmbed` β€” rank-generic patching (a 3-D volume uses the same code), factorized positional embedding by default. Tested: ViT and axial-ViT have identical parameter counts and different outputs.
- [x] **LTI.md** β€” `check_lti` measures linearity and time-invariance for every shipped mixer, and the N-D consequences are measured rather than asserted. This corrected the project's own folklore; see the coverage note.
- [ ] Multi-head is inherited from `AttentionMixer`; `ConvMixer` has no grouped-channel option beyond depthwise.
- [ ] A ViT reproduction row in RESULTS.md (CIFAR-10 at toy scale) β€” the family is tested but has no published-comparison number, which is the same standard Phase 9 held the others to.
- [ ] `join_time=False` is untested against a real video task; it is exercised only synthetically.
**Coverage:** both families pass full conformance at ranks 1–3, dense + sparse, plus separability, causality, dilation schedule, receptive field, config round-trip and learnability. Not covered: any reproduction number, grouped convolutions, video.
**What LTI.md corrected.** The received claim β€” repeated in this project's own working notes, and the implicit basis for Phase 2's "S4ND's simultaneous-separable as a degenerate plan" item β€” is that LTI mixers commute across axes, so for an LTI mixer the sweep order carries no information. False for every multichannel network: a `groups=1` convolution is a *matrix-valued* filter, offset `j` carrying its own channel matrix, and those commute only if the matrices do. Measured order gap for a perfectly linear, perfectly time-invariant conv: **3e-01**. Build the filter as a channel matrix times a spatial filter and it drops to **3e-16**. So `ScanPlan` matters for a linear CNN too, and "separable" and "order-free" are different claims β€” only the first one follows from LTI.
---
## Phase 7c β€” Verification against the published N-D methods Β· S Β· **started**
Every mixer in this library is checked against its source; the **composition**
never was. That is the part the library claims as its own, so it was the part
resting on internal consistency alone.
The policy, arrived at by trying the alternative: **transcribe the paper's
composition rule as a test oracle; do not vendor the paper's code as a
dependency.** An upstream 1-D block satisfies the mixer contract and adapts
cleanly (Track C2); an upstream N-D model carries a config framework, a
trainer, a rank-specific einsum table and a license question, all to run
twenty lines of arithmetic. `kron_operator` and the dense state-space
reference already set this precedent.
- [x] **S4ND** β€” its composition is not a sweep: one kernel per axis, outer
product in Fourier, one N-D FFT. Our sequential one-axis-per-layer sweep
reproduces it at **1.8e-17**, rank 2 and rank 3, with a channel-mixing
negative control. `tests/test_published_composition.py`. This is the
strongest evidence for the premise: the plan reproduces exactly a model that
was never written as a sweep.
- [x] **Mamba-ND** β€” *not* blocked on CUDA after all. `Block.__init__` takes
`mixer_cls` as an argument, so the selective scan is injected: give it our
portable `MambaMixer` and the whole N-D scan runs on MPS. Their block with
`n_dim_pos=1` is **bitwise identical** to `td.axial_apply` on CPU and MPS
across all six order/direction combinations, and `ScanPlan.paired`
reproduces their 12-layer schedule exactly, layer for layer. The repo still
**ships no LICENSE file**, so it stays un-vendored and the harness imports
it from a local clone (`dossier/verify_mamba_nd.py`).
- [x] Upstream S4ND runs on MPS too, once `pytorch_lightning` is stubbed and
the Cauchy kernel falls back to pure torch. Verified against *their*
kernels, not a transcription: bitwise at rank 3 on CPU, 4.7e-10 on MPS
(`dossier/verify_s4nd.py`, numbers in `dossier/README.md`).
- [ ] **Axial Transformer** β€” `lucidrains/axial-attention` is MIT and pure
torch, so this one can be checked numerically here.
- [x] **CaFA** β€” [BaratiLab/CaFA](https://github.com/BaratiLab/CaFA), MIT.
Their contraction is two sequential per-axis einsums, which `axial_contract`
reproduces at **2.0e-16 relative** (float64). `FABlockS2` is hardcoded to
two named axes, so the rank-generality is ours. Adopted their `qk_norm` and
their `K + gamma*I` kernel residual (both off by default); declined RoPE and
the spherical quadrature weights. Their `normalize_to_one=True` path raises
`AttributeError` β€” a guard on an attribute never assigned.
- [ ] A dossier section recording each comparison's number, as Phase 7 does
for the 1-D mixers.
**Coverage:** two of four published N-D methods verified numerically against
their own running code, on CPU and MPS. The CI test needs no vendored code and
no upstream clone; the dossier scripts need a clone and are not packaged.
**A method neither we nor the papers name.** Mamba-ND scans the last
`n_dim_pos` axes of a per-layer ordering *as one sequence*, batching the rest.
`n_dim_pos=1` is `td.axial_scan`; `n_dim_pos=rank` is `td.flatten`; their
published configs use the values in between (`factorization='hw_t'` β†’
`(2, 2, 4)`). Our two methods are the endpoints of a spectrum whose interior
we cannot currently express. See Track E.
---
## Phase 8 β€” Registry, config, save/load Β· M Β· done
- [x] `model.config` β€” the construction recipe recorded at build, plain JSON-able types, `n_layers` recorded as the plan's true depth (a recipe that could disagree with itself is not a recipe).
- [x] `td.build(dict | yaml path)` β€” registry by kind (`lstm`, `gru`, `s4`, `s4d`, `mamba`, + ND names); unknown keys are a hard error naming the offender.
- [x] `model.save(path)` / `td.load(path)` β€” one checkpoint file: format version, config, state dict; rebuilt model is the same model (outputs bitwise equal; validity mask included); incompatible format versions refused.
- [x] The round-trip proven for every model family and both composition families.
- [x] `safetensors` as an optional container, chosen by file extension, with the config carried in its metadata β€” still one file, and one that cannot execute code when opened. A test asserts the written file contains no pickle.
- [ ] Checkpoint migration policy written down: what version N promises to load from version Nβˆ’1, and a stored-fixture test per released version (the fixtures directory *is* the compatibility contract).
- [x] `td.build(checkpoint)` builds the architecture a checkpoint records and leaves the weights alone; `weights=True` is exactly `load`. Plus `td.read_config(path)` β€” what a checkpoint claims to be, in plain data, readable without the architecture still existing.
- [x] Entry-point registration (`torch_dimensions.models` group): a plugin's module is imported only when the registry is first consulted, and a plugin that fails to import warns rather than making this library unimportable.
**Coverage:** every shipped model round-trips config and weights bitwise, in both containers (torch pickle and safetensors); third-party registration and its failure path are tested. Not covered: cross-version loading β€” there is still no released fixture set, which is the one thing here that only time can produce.
---
## Phase 9 β€” Reproduce a published result Β· L Β· needs Phases 5–8 Β· **DONE for the sequence and sparse rows**
The make-or-break test of the abstraction: a published N-D result, reproduced from a config file alone, with no code written outside this library. Until this passes, "the unification holds" is a design claim, not an empirical one.
Candidates, in order of feasibility on available hardware:
- [x] **sMNIST / psMNIST with `td.S4D`** β€” **99.53%** and **97.79%** test accuracy, 20 epochs each, 201k parameters, ~21 min per run on one Apple Silicon laptop. The S4D paper reports ~99.6% and ~98.5% with far longer schedules, so both land inside the one-point target. This isolates the mixer from the N-D machinery: the lattice is the degenerate one.
- [x] **2-D lattice image classification** (`examples/repro/image_nd.py`) β€” MambaND with the paired schedule at **97.29%** and S4DND at **98.33%**, 2 epochs each over a real 2-D lattice with no pixel flattened into a sequence. The wall clock is the finding: 19.1 min against 1.2 for twice the parameters, which is BENCHMARKS.md's portable-Mamba-scan cost appearing in a real run. CIFAR-10 is wired and checksummed and waits only on a download.
- [x] **Sparse-lattice forecasting** on UCI Beijing air quality β€” 12 stations Γ— 6 pollutants, hourly, a genuinely 2-D lattice. `masked` vs `zeros` arms, same data and budget, scored on present cells only so the comparison is about representation and not about which cells count. Sparsity is **induced and said to be induced**: a naturally sparse dataset confounds "is missing" with "is different", and removing cells at random isolates the variable.
- [x] **Method comparison as one flag** β€” `cafa` and `axial_attention` arms of the same forecasting experiment, differing in exactly one argument.
- [x] Each ships as a `python -m` runnable, an append-only ledger, a RESULTS.md row with hardware and wall-clock, and a CI smoke variant on synthetic data (`tests/test_repro.py`) β€” including the check that makes the sparse comparison mean anything: a 1e3 perturbation in an absent cell must not move the output by one bit.
- [ ] Seed variance: every row is one seed. Three seeds per row is the next honest increment, and until then no row should be read as a mean. **This is now the blocking issue for the sparse claim**, whose arms differ by less than one percent relative.
- [ ] **The experiment the null result implies.** Masking did not improve accuracy on a *fixed* mask β€” a dense model learns the mask. So the claim has to be restated and retested as the thing masking actually guarantees: train on one sparsity pattern, evaluate on a **different** one. The masked model is provably invariant (bitwise, already tested); the dense model cannot be. That is a property experiment, not a leaderboard, and it is the one that would justify the feature in a paper.
- [ ] CIFAR-10 row (the download stalled repeatedly on this connection; the code path is tested, the number is not taken).
**Acceptance:** met β€” RESULTS.md carries reproducible rows, each one command.
**What the numbers said, including the part that stings.** The sequence reproductions worked: 99.53% and 97.79%, both inside a point of published. The sparse-lattice experiment produced a **null result** β€” masking absent cells did not beat zero-filling at 30% or 60% sparsity, because a mask fixed across training and evaluation is something a dense model simply learns. The library's masking still provides what it always claimed (absent cells provably cannot influence an output, tested bitwise); what it does not provide, on this task, is a better score. Writing that down is worth more than the row that would have been quietly dropped.
**Coverage:** the sequence and sparse claims have numbers at one seed each; the image rows are 2-epoch and say so; nothing here has run on CUDA, and no row is a mean.
**Risk:** dataset licensing and download flakiness β€” vendor nothing; download scripts with checksums, and CI smoke uses synthetic stand-ins.
---
## Phase 10 β€” Benchmarks Β· M Β· **DONE on MPS; CUDA open**
Measure the known risks before making any performance claim.
- [x] Fold overhead: **10–30% of a step** at rank β‰₯ 2 and small `d_model`, falling as width grows. Real, and not the first thing to optimize.
- [x] Dense vs factorized, and then the sweep that matters: the crossover is at **64Β³**, and below ~50 cells per axis per-line attention is 2–3Γ— *faster*. The family earns its complexity at the top end, not everywhere.
- [x] Portable mixer costs, MPS: the Mamba scan is 5–8Γ— the kernel-based SSMs β€” the number a fused path must beat.
- [ ] The same table against a fused path and against CUDA (Track C).
- [x] `torch.compile` is a **0.8Γ— slowdown** on MPS at these sizes, so it is not recommended by default. Re-measure before repeating that about any other machine.
- [x] Chunked fold: never won here; `chunk=64` cost 3Γ—. It is a memory valve, not a speed knob.
- [x] Rank 1β†’5 at fixed cells β€” and the design claim was **wrong**. Cost tracks the *length of the swept axis*, not the cell count and not the rank: rank 1 costs ~95Γ— rank 4 at the same 4,096 cells, because a sequential mixer's cost is its launch depth.
- [x] Published as BENCHMARKS.md with the machine named and the scripts in-repo. Interpretations live in the script as a `finding` field, because regenerating overwrites the markdown and an interpretation a re-run deletes is one nobody rewrites.
- [x] A CI perf-smoke (`tests/test_perf.py`) β€” but on **ratios between two configurations in the same process**, not on absolute times: a loaded runner is 5Γ— slower than a quiet laptop and an absolute threshold flakes exactly where it was meant not to. Machine speed cancels; complexity does not.
**Acceptance:** met on MPS. **Coverage:** one machine, float32, forward and forward+backward. Not covered: CUDA, half precision, multi-GPU, memory (only CUDA tracks an allocation peak β€” see DEBUG.md #20 for the plausible substitute that was removed).
---
## Phase 11 β€” Docs and release Β· M Β· **v0.1.0 SHIPPED; docs open**
- [x] README: the unification table, verified quickstart (every snippet executed before written), honest scope limits, "Correct on purpose" section.
- [x] CHANGELOG discipline (Keep-a-Changelog, real 0.1.0 entry).
- [x] PyPI: trusted publishing, tag-triggered, live β€” `pip install torch-dimensions` works.
- [x] DEBUG.md as a living practice document, with its citations enforced by a test.
- [ ] Docs site (mkdocs-material): API reference from docstrings, the design docs rendered, versioned with releases.
- [x] **[Adding a mixer](docs/adding-a-mixer.md)** β€” a real mixer end to end, with its conformance report, the table of what each check protects against, and the mistakes to expect. The example is executed by `tests/test_examples.py`.
- [x] **[Adding an nd_method](docs/adding-a-method.md)** β€” same treatment, and its best material was an accident: the conformance suite failed the guide's own example twice (DEBUG.md #21, #22), so the page documents what it caught and a test reproduces the broken version.
- [ ] Tutorial notebooks: (1) forecasting a real CSV end-to-end through `td.data`, (2) the method-of-multidimensionality comparison on one task, (3) sparse lattices β€” why masking matters, with the Phase 9 numbers.
- [ ] Docstring pass with a doctest runner in CI (examples in docstrings that execute are examples that stay true).
- [ ] Release automation niceties: CHANGELOG section extracted into the GitHub release notes by the publish workflow.
---
## Phase 12 β€” Autoregressive stepping Β· L Β· needs Track C fast paths Β· v0.2
Generation over a lattice: process one timestep at a time with carried state, so the same models that trained in parallel can roll forward step by step.
**The shape of it.** Autoregression is a property of the *time* axis only β€” the spatial axes are fully materialized at every step. So the primitive is
```python
state = model.init_state(batch) # per-layer, opaque
y_t, state = model.step(x_t, state) # x_t: (B, *shape, H) β€” one timestep
```
and a layer's behavior at step time follows from what it sweeps: a layer sweeping **time** consumes and updates its slice of the state; a layer sweeping a **spatial** axis runs exactly its normal forward on the single timestep (nothing to carry). The kernel family needs no new mechanism at all β€” its spatial kernels are already timestep-local, which the causality test proves today.
- [ ] `Mixer.step(x, state) -> (y, state)` as an optional protocol method; a mixer without it cannot be stepped and the model refuses AR mode loudly at `init_state`, never at step 500.
- [ ] `LSTM`/`GRU` state: `(h, c)` β€” free, `nn.LSTM` already does this.
- [ ] `Mamba` state: conv ring buffer (`d_convβˆ’1` inputs) + selective-scan state β€” mirrors upstream's inference cache, and must be *the same design* as the fused path's cache (why this phase follows Track C).
- [ ] `S4`/`S4D` state: the recurrent view of the same SSM β€” materialize `(dA, dB)` from the learned parameters with the *same* discretization as the kernel; DPLR steps in the diagonalized basis, O(N) per channel per step. The dense-reference test already proves the recurrence equals the kernel.
- [ ] Refusals: a plan sweeping time backward (bidirectional time says the future is visible β€” generation cannot honor it); any step-less mixer. Constructor-time, pointed messages.
- [ ] Conformance check #8: parallel `forward` equals sequential `step`ping β€” bitwise for RNNs, float tolerance for SSMs. The single check that catches the classic AR bugs (off-by-one state, conv buffer misalignment, leaked normalization statistics).
- [ ] `save`/`load` round-trips mid-generation state (a checkpoint that cannot resume generation is half a checkpoint).
- [ ] Viewer: step mode in the run panel (Track B V5) β€” generate cell-by-cell along time and watch the lattice fill.
**Deliberately out of scope, permanently:** sampling policy, beam search, KV caches for attention over time, readout heads β€” the caller's, same as training loops. `step` is the primitive; generation loops are five lines of user code around it.
---
# Track B β€” Viewer / GUI mode
The library emits a versioned JSON spec (`td.spec`, SPEC_VERSION 1); the viewer consumes it. The boundary is a document, never a socket into the library.
## V1 β€” Static architecture viewer Β· **done**
- [x] React + Vite + react-three-fiber; lattice rendered as cubes, absent cells genuinely absent; sweep wavefront animated per layer; per-layer timeline; sample specs for 2-D sparse LSTM, 3-D paired Mamba-ND, 4-D S4D.
- [x] Rank β‰₯ 4 via **dimensional stacking**: a cube of cubes, spaced groups along the fourth axis (the "16-pack" idea).
- [x] Open-a-spec-JSON from disk; sidebar with model card, directions, unswept-axis warning.
## V1.5 β€” Live runs and controls Β· **done**
- [x] `examples/viewer_live.py`: state to `run.json` after every step; viewer polls, auto-switches, loss chart (train + held-out, log scale), progress bar, status lifecycle.
- [x] Run controls: the script **waits** β€” nothing trains until Start is pressed; pause froze the counter (proven 5255 β†’ 5255), resume, stop; control server with CORS; state machine.
- [x] Presets: `lstm2d`, `mamba4d` (18-layer Mamba-ND over a sparse 4-D lattice β€” trained live to held-out 9e-5).
## V2 β€” Ship it in the wheel Β· M Β· **done**
- [x] `td.viz.show(model)` β€” the bundle rides inside the wheel (387 KB), so there is no node, no network and no build step at install time. `serve()` is the non-blocking half.
- [x] The publish workflow builds the bundle and greps the built **wheel** for `viz/static/index.html` rather than trusting the build log; CI has a viewer job.
- [x] `show(spec_dict)`, `show(path.json)`, `show(checkpoint)`. Anything else raises β€” an empty viewer is a far worse error message than an exception.
- [x] A stale local training run got baked into the first bundle and the viewer loaded it *in preference to* the model passed in (DEBUG.md #19). Stripped by the installer, guarded by a test.
- [ ] Playwright smoke in CI: the three samples render, the sidebar populates. The accessibility-tree assertions are proven by hand; making them a script is what remains.
## V3 β€” Shape flow Β· M
- [ ] Per-layer tensor-shape inspector: hover a layer, see `(B, [T,] *shape, H)` β†’ fold β†’ `(M, A, H)` β†’ restore, with the actual numbers for a chosen batch size.
- [ ] The fold visualized: which axes collapse into `M` for the hovered layer (this is the single hardest thing to explain in prose; it is one animation).
- [ ] Parameter breakdown per layer (bar, hover for tensor names) from the spec's `n_params`.
## V4 β€” Metrics dashboard Β· L
- [ ] Multi-run compare: `run.json` grows a run id; the panel lists past runs (ring buffer on disk), overlays loss curves.
- [ ] More channels: learning rate, gradient norm, per-axis gradient norms (does the `w` sweep learn faster than `h`? β€” a genuinely novel N-D diagnostic), step time.
- [ ] Live lattice heatmap: per-cell loss contribution at eval time, painted on the 3-D lattice β€” *where* on the grid the model is wrong, over training.
- [ ] Export: PNG of the scene, CSV of metrics.
## V5 β€” Full GUI mode (the original vision) Β· XL
- [ ] Config editor in the panel: edit d_model/layers/method/plan with validation via `td.build`'s schema errors, see the architecture update live *before* any training.
- [ ] Launch from the panel: the control server grows a `POST /launch` that spawns the training script with the edited config (opt-in flag, localhost only, explicitly not a deployment tool).
- [ ] Expected-output preview: run one forward pass on synthetic data at build time; show output shapes and ranges ("press continue to train").
- [ ] Plan editor: drag layers, flip directions, see the coverage report (Phase 2's `plan.coverage`) recompute β€” the aliasing bug rendered as a picture would have caught DEBUG.md #4 in a viewer.
- [ ] AR step mode (with Phase 12): generate along time, watch the lattice fill cell by cell.
- [ ] Sparse-mask painter: click cells on/off, export the `valid` mask β€” the fastest way to build a demo lattice.
---
# Track C β€” Performance & hardware
## C1 β€” CUDA verification Β· M Β· **procedure written, not yet run**
- [x] The procedure exists as [docs/cuda-checklist.md](docs/cuda-checklist.md): what to run, what each item protects, and what it deliberately does not cover. The README now says CUDA is a design claim rather than a test result, which it should have said from the start.
- [x] The procedure is now one runnable file, `scripts/cuda_check.py`: every CUDA claim with the number behind it, `skip` with a reason where there is no device, non-zero exit on any failure. Verified on CPU (1 pass, 14 skips) so the GPU run cannot fail on a harness bug.
- [ ] Run it. Two commands on a free Colab T4. **Nothing in this project has executed a single CUDA kernel** β€” and since 0.3.1 that includes the whole fused-vs-portable dispatch and Mamba-3's transcription, which has never met the Triton kernel it was transcribed from.
- [ ] The bitwise rank-1 LSTM claim re-checked under cuDNN (expected: holds per single layer; document whatever is true).
- [ ] fp16/bf16/AMP conformance additions: the relative-cancellation guard was *designed* dtype-aware; prove it at fp16 where the accidental-safety finding (rounding to exact zero) came from.
- [ ] A `gpu` CI lane (GitHub GPU runners are paid β€” decide when Phase 7 fast paths land; until then, a documented manual Colab checklist per release).
## C2 β€” Fast paths Β· L (see Phase 7 open items)
- [ ] The adapter pattern: portable path is the reference; fused path must agree with it in the conformance suite on the same machine. Never ship a fast path whose only test is "it ran".
## C3 β€” Throughput engineering Β· L Β· gated on Phase 10 numbers
- [ ] Fused/permute-avoiding fold if Phase 10 says permute dominates (the only justified kernel).
- [ ] `chunk` auto-tuning from device properties.
- [ ] Associative-scan (log-depth) portable selective scan.
- [ ] Activation checkpointing option; memory profile at rank 4.
- [ ] `torch.compile` graph-break audit per family (the conformance hook exists; make the graphs clean).
## C4 β€” Distributed sanity Β· M Β· v0.3+
- [ ] DDP: one multi-process CPU test that gradients sync correctly (buffers like `cell_mask` must not desync β€” the `persistent=False` choice needs a test under DDP broadcast).
- [ ] FSDP smoke with the SSM families (complex-as-real parameters interact with flat-param sharding; find out now, not in a user issue).
---
# Track D β€” Ecosystem & community
- [ ] Docs site + guides (Phase 11 items; listed here because they gate community, not release).
- [ ] Example gallery in-repo: one runnable file per (family Γ— method) pair β€” the matrix is the product; show it as a matrix.
- [ ] GitHub issue templates that ask for `td.spec(model)` output β€” the spec doubles as a bug report's reproduction seed.
- [ ] `CITATION.cff` + a short tech report (arXiv) once Phase 9 rows exist β€” the sparse-lattice baseline is the publishable nugget.
- [ ] Hugging Face Hub integration via Phase 8 checkpoints (config + safetensors maps cleanly; do it after safetensors lands, not before).
- [ ] "State of the lattice" doc: honest comparison table vs mamba-ssm, s4, FLA β€” what to use when, including "not us".
- [ ] Conformance badges: a downstream mixer repo can run `check_block` in its CI and claim conformance β€” write the two paragraphs and the badge that make that a thing.
---
# Track E β€” Research extensions (parked deliberately; each is a design, not a stub)
- [ ] **Space-filling traversals** (Hilbert/Morton) as a *second step kind* carrying a full cell permutation β€” rejected from `ScanPlan` v1 because a step is `(axis, reverse)` and a curve is not axis-aligned; the design doc for step-kind-2 must handle sparse lattices (curve over present cells only?) before any code.
- [ ] **Dynamic-shape lattices**: `dim` without `shape` β€” rank known, sizes bound at first forward. The scan family needs only rank; masks and kernel-family bias tables need sizes. Doable for dense scan-only models; the design question is what `spec()` and `save()` mean before binding.
- [ ] **Learned schedules**: differentiable relaxation over sweep order (soft mixture of axis sweeps annealed to a hard plan). Research-grade; the `ScanPlan`-as-data design is what makes it even expressible.
- [ ] **Hierarchical / nested lattices** (multi-resolution grids, lattice-of-lattices): the S4ND paper's multi-scale appendix and weather models both want this; `Lattice.merge/slice` (Phase 1) is the substrate.
- [ ] **Masked-pretraining utilities**: random cell masking as a training-time lattice transform β€” the `valid` machinery already guarantees inertia; the utility is three functions and a tutorial, but only after Phase 9 proves the supervised story.
- [ ] **Scan factorization** β€” the Mamba-ND paper's own name for it (Β§4.1,
Fig. 5b): a policy choosing how many axes are flattened into each scanned
sequence. "No factorization" is one sequence over every cell β€” `td.flatten`.
Maximal factorization is one sequence per line β€” `td.axial_scan`. The
interior (D 2D sequences out of a 3D lattice) is what their published
configs actually use, and this library cannot express it. Probably a
`groups=` argument giving per-layer ordered axis subsets rather than a new
strategy; the design question is whether a `ScanPlan` step should be allowed
to name a *set* of axes rather than one, which would make all three policies
a single schedule type.
- [ ] **Continuous-time lattices**: irregular timestamps per cell (the `dt` in SSMs is *built* for this β€” S4's continuous-time parameterization applied per-observation). The biggest research swing in the list; would make td the only library doing irregular N-D series natively.
---
## Critical path
```
β”Œβ”€β”€β–Ί 5 βœ… ┐
0 ──► 1 ──► 2 ──► 3 ──► 4 ──┼──► 6 βœ… ┼──► 8 βœ…β”€β”€β–Ί 9 βœ…β”€β”€β–Ί 10 βœ…β”€β”€β–Ί 11 (docs: guides βœ…, site open)
(all done) └──► 7 βœ… β”˜ β”‚
β”‚ └──► Track D (has 9's numbers now)
└──► [C1 CUDA] ──► C2 fast paths ──► 12 AR (v0.2)
Track B (viewer): V1 βœ… ── V1.5 βœ… ── V2 βœ… ── [V3] ── V4 ── V5 (V5 wants 12)
```
The critical path now runs through **C1**, and it is fifteen minutes of
someone else's GPU. Everything downstream of it β€” the fused fast paths, the AR
state design that must match the fused cache, the half-precision claims β€” is
blocked behind a machine this project has never had. That is the single
highest-leverage hour available to anyone reading this.
Phases 5, 6, 7 were independent once the conformance suite existed, and that bet paid: all three landed without touching each other. The same principle now says **9 and C1 are independent** β€” reproduction runs on MPS while CUDA verification happens on Colab β€” and both feed 10. The library remains useful and shippable without a single CUDA kernel; that stays deliberate.
---
## Repo layout
```
torch-dimensions/
β”œβ”€β”€ pyproject.toml hatchling; torch>=2.4 only; extras [mamba] [fla] [dev] [all]
β”œβ”€β”€ README.md Β· DESIGN.md Β· PLAN.md Β· VIEWER.md Β· DEBUG.md Β· CHANGELOG.md Β· CONTRIBUTING.md
β”œβ”€β”€ .github/workflows/ ci.yml (full-history checkout) Β· publish.yml (OIDC, size guard)
β”œβ”€β”€ src/torch_dimensions/
β”‚ β”œβ”€β”€ lattice.py Β· plan.py Β· spec.py Β· config.py Β· testing.py
β”‚ β”œβ”€β”€ compose/ scan.py Β· kernel.py Β· attention.py Β· __init__ (strategies, registry)
β”‚ β”œβ”€β”€ mixers/ base.py Β· rnn.py Β· ssm.py
β”‚ β”œβ”€β”€ models/ base.py Β· rnn.py Β· ssm.py
β”‚ └── data/ coords.py Β· table.py Β· window.py Β· source.py Β· collate.py
β”œβ”€β”€ tests/ one file per module + conformance, device, fuzz, trainable, debug_md
β”œβ”€β”€ examples/ train_nd.py Β· viewer_live.py
└── viewer/ Vite + React + react-three-fiber (excluded from sdist)
```