"""Every mixer against every method — the claim, checked as a matrix. The library's thesis is one sentence: *an N-D model is a 1-D mixer plus a plan for sweeping it over an N-D lattice*. If that is true then any mixer composes with any method, and the product of the two lists is the product the library ships. Until now that was spot-checked — each mixer was tested under its usual composition, each method under a convenient mixer — and the diagonal of a matrix is not the matrix. 13 mixers x 4 methods is 52 cells, and every one of them is a configuration a user can write in a single line. The invariants here are the ones that must hold for *any* combination, and they are deliberately few, because this file's job is breadth: shape, finite values, gradients reaching the mixer, absent cells staying inert, and batch independence. Depth belongs in the per-family files, which know what each mechanism additionally claims. All 52 run in about 13 seconds, so the matrix is exhaustive rather than sampled. A cell that cannot work should be a refusal with a reason, not a silent absence from the test suite. """ from __future__ import annotations import itertools import pytest import torch import torch_dimensions as td from torch_dimensions import mixers as M pytest.importorskip("einops", reason="the vendored mixers need the [upstream] extra") pytest.importorskip("hydra", reason="the s4 pipeline needs hydra-core") D_MODEL = 32 MIXERS = [ "LSTMMixer", "GRUMixer", "S4DMixer", "S4Mixer", "MambaMixer", "UpstreamS4DMixer", "UpstreamS4Mixer", "UpstreamMambaMixer", "UpstreamMamba2Mixer", "Mamba3Mixer", "AttentionMixer", "ConvMixer", "TCNMixer", ] METHODS = ["axial_scan", "axial_attention", "cafa", "flatten"] CELLS = list(itertools.product(MIXERS, METHODS)) # The head layout of the Mamba-2/3 blocks constrains d_model; nothing else # needs an argument, which is itself part of the claim. EXTRA = {"UpstreamMamba2Mixer": {"headdim": 16}, "Mamba3Mixer": {"headdim": 16}} def lattice(sparse: bool = True) -> td.Lattice: if not sparse: return td.Lattice(shape=(4, 5), names=("h", "w"), time=True) gen = torch.Generator().manual_seed(5) valid = torch.rand(4, 5, generator=gen) > 0.3 valid[0, 0] = True return td.Lattice(shape=(4, 5), names=("h", "w"), valid=valid, time=True) def build(mixer: str, method: str, lat: td.Lattice): cls = getattr(M, mixer) kw = EXTRA.get(mixer, {}) torch.manual_seed(0) return td.LSTM( # the host class only supplies the composition; `mixer=` is the mechanism D_MODEL, 2, lat, d_input=1, mixer=lambda d, _c=cls, _k=kw: _c(d, **_k), method=getattr(td, method), ) def inputs(lat: td.Lattice, batch: int = 2, t: int = 3) -> torch.Tensor: torch.manual_seed(1) return torch.randn(batch, t, *lat.shape, 1) @pytest.mark.parametrize("mixer,method", CELLS, ids=[f"{m}-{s}" for m, s in CELLS]) def test_the_cell_runs_and_keeps_the_lattice_shape(mixer, method): lat = lattice() model = build(mixer, method, lat).eval() x = inputs(lat) with torch.no_grad(): y = model(x) assert y.shape[:-1] == x.shape[:-1] assert y.shape[-1] == D_MODEL assert torch.isfinite(y).all() @pytest.mark.parametrize("mixer,method", CELLS, ids=[f"{m}-{s}" for m, s in CELLS]) def test_gradients_reach_the_mixer(mixer, method): """Not merely that a backward pass runs: that it reaches the *mixer*. A composition that dropped the mixer out of the graph — folding it away, or detaching somewhere in the axis bookkeeping — would still train the projections and still show a falling loss.""" lat = lattice() model = build(mixer, method, lat) model(inputs(lat)).pow(2).mean().backward() inner = [p for name, p in model.named_parameters() if name.startswith("nd.")] assert inner, "the composition exposed no parameters of its own" assert any(p.grad is not None and p.grad.abs().sum() > 0 for p in inner), ( "no gradient reached any parameter inside the composition" ) @pytest.mark.parametrize("mixer,method", CELLS, ids=[f"{m}-{s}" for m, s in CELLS]) def test_absent_cells_stay_inert_in_every_cell_of_the_matrix(mixer, method): """The sparse guarantee is a property of the *composition*, so it has to survive every mixer put through it — including the ones whose own code knows nothing about masks.""" lat = lattice() model = build(mixer, method, lat).eval() mask = lat.mask(torch.float32) x = inputs(lat) * mask noise = torch.randn_like(x) * 50 * (1 - mask) with torch.no_grad(): assert torch.equal(model(x), model(x + noise)) @pytest.mark.parametrize("mixer,method", CELLS, ids=[f"{m}-{s}" for m, s in CELLS]) def test_batch_independence_in_every_cell_of_the_matrix(mixer, method): lat = lattice() model = build(mixer, method, lat).eval() x = inputs(lat) poisoned = x.clone() poisoned[1:] += 100.0 with torch.no_grad(): assert torch.equal(model(x)[0], model(poisoned)[0]) @pytest.mark.parametrize("mixer,method", CELLS, ids=[f"{m}-{s}" for m, s in CELLS]) def test_a_dense_lattice_works_in_every_cell(mixer, method): lat = lattice(sparse=False) model = build(mixer, method, lat).eval() with torch.no_grad(): assert torch.isfinite(model(inputs(lat))).all() def test_the_matrix_is_complete(): """The count itself is the claim. If a mixer or a method is added and this file is not, the product silently stops being the product — so the size is asserted against the exported lists rather than against a literal.""" exported_mixers = [n for n in M.__all__ if n.endswith("Mixer") and not n.startswith("Mixer")] assert sorted(MIXERS) == sorted(exported_mixers), ( "the matrix does not cover every exported mixer; " f"missing {sorted(set(exported_mixers) - set(MIXERS))}" ) # `ND_METHODS` is a registry users extend at runtime — `register_nd_method` # is public API, and the guide in docs/adding-a-method.md adds "pyramid" to # it for real. So the claim cannot be equality with the registry; it is # equality with the methods the *library* ships, which are exactly those # the registry maps to a `td.`-exported callable of the same name. shipped = sorted(n for n in td.ND_METHODS if getattr(td, n, None) is td.ND_METHODS[n]) assert sorted(METHODS) == shipped, ( f"the matrix does not cover every method the library ships; missing " f"{sorted(set(shipped) - set(METHODS))}" ) assert len(CELLS) == len(MIXERS) * len(METHODS)