"""ViT and the joint (flatten) method of multidimensionality. The claim under test is that ViT is not a special model but a composition choice: the same patches and the same parameter count, with a different answer to "how are the extra axes handled". """ import pytest import torch import torch_dimensions as td from torch_dimensions.models.vit import PatchEmbed, ViT # -- the flatten method ------------------------------------------------------- def _factory(**kw): def build(lat, d_model, plan=None): return td.Transformer( d_model, lat.n_axes, lat, plan=plan, method=td.flatten, mixer_kwargs={"n_heads": 2}, **kw, ) return build def test_flatten_passes_every_applicable_check(): report = td.testing.check_block(_factory(), d_model=4, ranks=(1, 2, 3)) assert report, str(report) def test_flatten_sees_every_cell_in_one_layer(): """The defining property: one layer, and every cell can reach every other. The axial methods cannot do this — a single axial layer mixes along one axis only, so two cells differing in two coordinates are unreachable from each other until the second layer. """ lat = td.Lattice(shape=(3, 4), names=("h", "w")) torch.manual_seed(0) joint = td.Transformer(8, 1, lat, method=td.flatten, mixer_kwargs={"n_heads": 2}).double() axial = td.Transformer(8, 1, lat, mixer_kwargs={"n_heads": 2}).double() x = torch.randn(1, 3, 4, 8, dtype=torch.float64) bumped = x.clone() # A cell differing from (0, 0) along *both* axes. The perturbation must # vary across features: adding a constant to every feature of a token is # removed exactly by the pre-norm LayerNorm, so a constant bump is # invisible to any pre-norm block and would make this test vacuous. torch.manual_seed(1) bumped[0, 2, 3] += torch.randn(8, dtype=torch.float64) * 10.0 joint_reach = (joint(x) - joint(bumped))[0, 0, 0].abs().max().item() axial_reach = (axial(x) - axial(bumped))[0, 0, 0].abs().max().item() assert joint_reach > 1e-6, "joint attention did not reach across both axes in one layer" assert axial_reach < 1e-12, "a single axial layer must not reach across two axes" def test_absent_cells_are_dropped_from_the_sequence_not_masked_into_it(): """The flatten family's one structural advantage on sparse lattices.""" valid = torch.ones(3, 4, dtype=torch.bool) valid[1, 1] = valid[2, 3] = False lat = td.Lattice(shape=(3, 4), names=("h", "w"), valid=valid) model = td.Transformer(8, 2, lat, method=td.flatten, mixer_kwargs={"n_heads": 2}) assert model.nd.seq_len == 10 == lat.n_valid def test_time_can_be_joined_or_kept_out_of_the_sequence(): lat = td.Lattice(shape=(2, 3), names=("h", "w"), time=True) joined = td.Transformer(8, 2, lat, method=td.flatten, mixer_kwargs={"n_heads": 2}).double() apart = td.Transformer( 8, 2, lat, method=td.flatten, join_time=False, mixer_kwargs={"n_heads": 2} ).double() x = torch.randn(2, 5, 2, 3, 8, dtype=torch.float64) assert joined(x).shape == apart(x).shape == x.shape # With time held out, each timestep is mixed alone: perturbing t=4 cannot # reach t=0. With time joined, it can. Non-constant, per the note above. torch.manual_seed(1) bumped = x.clone() bumped[:, 4] += torch.randn(2, 3, 8, dtype=torch.float64) * 10.0 assert (apart(x) - apart(bumped))[:, 0].abs().max().item() < 1e-12 assert (joined(x) - joined(bumped))[:, 0].abs().max().item() > 1e-6 def test_flatten_refuses_to_be_the_model_without_a_mixer(): lat = td.Lattice(shape=(3, 4), names=("h", "w")) with pytest.raises(ValueError, match="no operator at all"): td.flatten(None, td.ScanPlan.cyclic(lat.axis_names, 2), lat, 8) def test_the_spec_says_no_axis_is_swept(): lat = td.Lattice(shape=(3, 4), names=("h", "w")) s = td.spec(td.Transformer(8, 2, lat, method=td.flatten, mixer_kwargs={"n_heads": 2})) assert s["nd_method"]["family"] == "flatten" assert [layer["kind"] for layer in s["layers"]] == ["flatten", "flatten"] assert all(layer["axis"] is None for layer in s["layers"]) assert s["sweeps"]["joint_axes"] == ["h", "w"] assert s["sweeps"]["directions"] == {} # -- patch embedding ---------------------------------------------------------- @pytest.mark.parametrize( "image,patch,grid", [((32, 32), 4, (8, 8)), ((32, 32), (4, 8), (8, 4)), ((8, 8, 8), 2, (4, 4, 4))], ) def test_patch_embed_is_rank_generic(image, patch, grid): embed = PatchEmbed(image, patch, in_channels=2, d_model=16) assert embed.grid == grid x = torch.randn(3, *image, 2) assert embed(x).shape == (3, *grid, 16) assert embed.lattice().shape == grid def test_patch_embed_preserves_patch_contents(): """Each output cell must be a function of exactly its own patch. A permute bug here shuffles pixels between patches and produces a model that trains, badly, forever — the archetypal N-D failure this library exists to make loud. """ embed = PatchEmbed((4, 4), 2, in_channels=1, d_model=8).double() x = torch.zeros(1, 4, 4, 1, dtype=torch.float64) base = embed(x) x[0, 0, 0, 0] = 1.0 # a pixel in patch (0, 0) only out = embed(x) changed = (out - base).abs().sum(dim=-1)[0] > 1e-12 assert changed[0, 0] and changed.sum() == 1, changed def test_a_partial_patch_is_refused_rather_than_cropped(): with pytest.raises(ValueError, match="divide the image exactly"): PatchEmbed((32, 30), 4, in_channels=3, d_model=8) # -- the model ---------------------------------------------------------------- def test_vit_returns_per_patch_features(): model = ViT(32, 2, image=(16, 16), patch=4, in_channels=3, n_heads=2) assert model.grid == (4, 4) assert model(torch.randn(2, 16, 16, 3)).shape == (2, 4, 4, 32) def test_the_method_is_the_only_difference_between_vit_and_axial_vit(): """The library's central claim, at its most literal: same patches, same parameter count, one argument apart.""" kw = dict(image=(16, 16), patch=4, in_channels=3, n_heads=2) joint = ViT(32, 4, **kw) axial = ViT(32, 4, method=td.axial_scan, **kw) assert sum(p.numel() for p in joint.parameters()) == sum(p.numel() for p in axial.parameters()) assert type(joint.nd).__name__ == "Flatten" assert type(axial.nd).__name__ == "AxialScan" x = torch.randn(2, 16, 16, 3) assert joint(x).shape == axial(x).shape # Same shape, different computation — otherwise the comparison is empty. assert not torch.equal(joint(x), axial(x)) def test_factorized_position_embedding_is_smaller_and_still_positional(): kw = dict(image=(16, 16), patch=2, in_channels=1, n_heads=2) fac = ViT(32, 1, pos_embed="factorized", **kw) full = ViT(32, 1, pos_embed="full", **kw) n_fac = sum(p.numel() for p in fac.pos.parameters()) n_full = sum(p.numel() for p in full.pos.parameters()) assert n_fac == 2 * 8 * 32 and n_full == 8 * 8 * 32 assert n_fac < n_full # "Positional" means two identical patches at different positions get # different embeddings. Without that the table is decoration. torch.manual_seed(0) for m in (fac, full): torch.nn.init.trunc_normal_(m.pos.tables[0], std=0.5) x = torch.zeros(1, 16, 16, 1) out = m.patch_embed(x) assert not torch.allclose(m.pos(out)[0, 0, 0], m.pos(out)[0, 1, 0]) def test_vit_refuses_a_lattice_it_would_have_to_reconcile(): with pytest.raises(ValueError, match="builds its lattice"): ViT(32, 1, image=(8, 8), patch=2, lattice=td.Lattice(shape=(4, 4))) def test_vit_builds_from_config_and_round_trips(tmp_path): model = td.build( { "kind": "vit", "d_model": 32, "n_layers": 2, "image": [16, 16], "patch": 4, "in_channels": 3, "n_heads": 2, } ).eval() x = torch.randn(2, 16, 16, 3) path = tmp_path / "vit.td" model.save(path) same = td.load(path).eval() assert torch.equal(model(x), same(x)) assert same.grid == (4, 4) def test_vit_learns(): """A task that needs the patch grid: predict each patch's mean intensity from an image where the informative pixel sits in a different patch.""" torch.manual_seed(0) model = ViT(32, 2, image=(8, 8), patch=2, in_channels=1, n_heads=2) head = torch.nn.Linear(32, 1) opt = torch.optim.Adam([*model.parameters(), *head.parameters()], lr=3e-3) def draw(g): x = torch.randn(16, 8, 8, 1, generator=g) y = x.reshape(16, 4, 2, 4, 2, 1).mean(dim=(2, 4)) # per-patch mean return x, y g = torch.Generator().manual_seed(0) first = last = 0.0 for i in range(120): x, y = draw(g) loss = (head(model(x)) - y).pow(2).mean() first = loss.item() if i == 0 else first last = loss.item() opt.zero_grad() loss.backward() opt.step() assert last < first / 3, (first, last) # -- refusals and reporting --------------------------------------------------- def test_patch_rank_must_match_the_image(): with pytest.raises(ValueError, match="axes, image has"): PatchEmbed((16, 16), (4, 4, 4), in_channels=1, d_model=8) @pytest.mark.parametrize( "bad,match", [ (torch.randn(2, 16, 16), "expected a 4-D tensor"), (torch.randn(2, 8, 8, 3), "expected image dims"), (torch.randn(2, 16, 16, 5), "expected 3 channels"), ], ) def test_malformed_input_is_refused_by_name(bad, match): embed = PatchEmbed((16, 16), 4, in_channels=3, d_model=8) with pytest.raises(ValueError, match=match): embed(bad) def test_patch_embed_repr_states_the_grid_it_produces(): text = repr(PatchEmbed((32, 32), 4, in_channels=3, d_model=8)) assert "grid=(8, 8)" in text and "patch=(4, 4)" in text def test_unknown_pos_embed_is_refused_and_none_is_honoured(): with pytest.raises(ValueError, match="factorized\\|full\\|none"): ViT(16, 1, image=(8, 8), patch=2, in_channels=1, pos_embed="sinusoidal", n_heads=2) plain = ViT(16, 1, image=(8, 8), patch=2, in_channels=1, pos_embed="none", n_heads=2) assert sum(p.numel() for p in plain.pos.parameters()) == 0 # With no positional embedding the stage is a pass-through, and the model # still runs — "none" is a supported choice, not a broken one. x = torch.randn(1, 8, 8, 1) assert torch.equal(plain.pos(plain.patch_embed(x)), plain.patch_embed(x)) assert plain(x).shape == (1, 4, 4, 16) assert "none" in repr(plain.pos) def test_pos_embed_repr_reports_its_parameter_cost(): fac = ViT(16, 1, image=(8, 8), patch=2, in_channels=1, n_heads=2) assert "factorized" in repr(fac.pos) and "grid=(4, 4)" in repr(fac.pos) # -- the flatten composition's own edges -------------------------------------- def test_flatten_reports_what_it_spans(): lat = td.Lattice(shape=(3, 4), names=("h", "w"), time=True) nd = td.Transformer(8, 2, lat, method=td.flatten, mixer_kwargs={"n_heads": 2}).nd assert "space+time" in repr(nd) and "tokens=12" in repr(nd) apart = td.Transformer( 8, 2, lat, method=td.flatten, join_time=False, mixer_kwargs={"n_heads": 2} ).nd assert "space only" in repr(apart) def test_flatten_refuses_a_wrong_width_and_a_shape_changing_mixer(): lat = td.Lattice(shape=(3, 4), names=("h", "w")) model = td.Transformer(8, 1, lat, method=td.flatten, mixer_kwargs={"n_heads": 2}) with pytest.raises(ValueError, match="expected 8 features"): model(torch.randn(1, 3, 4, 5)) class Truncating(torch.nn.Module): def __init__(self, d_model, **_): super().__init__() self.scale = torch.nn.Parameter(torch.ones(())) def forward(self, x): return x[:, :-1] * self.scale bad = td.Transformer(8, 1, lat, method=td.flatten, mixer=Truncating) with pytest.raises(ValueError, match="mixer changed shape"): bad(torch.randn(1, 3, 4, 8)) def test_flatten_chunking_matches_the_unchunked_path(): """`chunk` exists to keep a fused kernel inside its grid limits; it must not change the answer.""" lat = td.Lattice(shape=(3, 4), names=("h", "w"), time=True) torch.manual_seed(0) whole = ( td.Transformer(8, 2, lat, method=td.flatten, mixer_kwargs={"n_heads": 2}).double().eval() ) torch.manual_seed(0) piece = ( td.Transformer(8, 2, lat, method=td.flatten, chunk=1, mixer_kwargs={"n_heads": 2}) .double() .eval() ) x = torch.randn(3, 2, 3, 4, 8, dtype=torch.float64) assert torch.allclose(whole(x), piece(x), rtol=0, atol=1e-12) def test_a_shared_mixer_instance_is_used_by_every_layer(): lat = td.Lattice(shape=(3, 4), names=("h", "w")) shared = td.mixers.AttentionMixer(8, 2) nd = td.Flatten( mixer=shared, plan=td.ScanPlan.cyclic(lat.axis_names, 3), lattice=lat, d_model=8 ) assert all(m is shared for m in nd.mixers) assert nd(torch.randn(1, 3, 4, 8)).shape == (1, 3, 4, 8)