tests/test_spec.py

"""Acceptance for td.spec — the architecture document the viewer reads.

See VIEWER.md. The spec is a contract with a separate codebase, so the tests
care about exact keys and lossless encoding, not just plausibility.
"""

import json

import pytest
import torch

import torch_dimensions as td
from torch_dimensions.spec import SPEC_VERSION, _rle


def rle_decode(runs, total):
    """Independent decoder, so the encoder is checked against something that
    does not share its code."""
    out, value = [], False
    for count in runs:
        out.extend([value] * count)
        value = not value
    assert len(out) == total
    return torch.tensor(out, dtype=torch.bool)


def sparse_lattice():
    valid = torch.tensor([[True, True], [True, False], [False, True]])
    return td.Lattice(shape=(3, 2), names=("state", "sku"), valid=valid, time=True)


# -- envelope ----------------------------------------------------------------


def test_spec_is_json_serializable():
    model = td.LSTM(4, 3, td.Lattice(shape=(2, 3), time=True))
    text = json.dumps(td.spec(model))
    assert json.loads(text)["format"] == "torch-dimensions/architecture"


def test_spec_carries_a_version():
    model = td.LSTM(4, 2)
    assert td.spec(model)["version"] == SPEC_VERSION


def test_spec_refuses_a_module_that_does_not_describe_itself():
    with pytest.raises(TypeError, match="does not describe itself"):
        td.spec(torch.nn.Linear(2, 2))


# -- lattice -----------------------------------------------------------------


def test_lattice_axes_mark_time_as_dynamic():
    spec = td.spec(td.LSTM(4, 3, sparse_lattice()))["lattice"]
    assert spec["names"] == ["time", "state", "sku"]
    assert spec["axes"][0] == {"name": "time", "size": None, "dynamic": True}
    assert spec["axes"][1] == {"name": "state", "size": 3, "dynamic": False}
    assert spec["rank"] == 2 and spec["n_axes"] == 3


def test_presence_encoding_is_lossless():
    lat = sparse_lattice()
    cells = td.spec(td.LSTM(4, 3, lat))["lattice"]["cells"]
    assert cells == {
        "total": 6,
        "present": 4,
        "dense": False,
        "present_rle": cells["present_rle"],
    }
    decoded = rle_decode(cells["present_rle"], 6)
    assert torch.equal(decoded, lat.valid.reshape(-1))


def test_a_dense_lattice_encodes_as_one_run():
    lat = td.Lattice(shape=(2, 3), time=True)
    cells = td.spec(td.LSTM(4, 3, lat))["lattice"]["cells"]
    assert cells["dense"] is True
    assert cells["present_rle"] == [0, 6]


@pytest.mark.parametrize(
    "flags",
    [
        [True] * 5,
        [False] * 4 + [True],
        [True, False, True, False, True],
        [False, False, True, True, False],
    ],
)
def test_rle_round_trips_for_arbitrary_masks(flags):
    t = torch.tensor(flags, dtype=torch.bool)
    assert torch.equal(rle_decode(_rle(t), len(flags)), t)


# -- layers ------------------------------------------------------------------


def test_layers_name_the_axis_each_one_sweeps():
    lat = td.Lattice(shape=(2, 3), names=("h", "w"), time=True)
    layers = td.spec(td.LSTM(4, 4, lat))["layers"]
    assert [layer["axis"] for layer in layers] == ["time", "h", "w", "time"]
    assert all(layer["mixer"] == "LSTMMixer" for layer in layers)
    assert all(layer["n_params"] > 0 for layer in layers)


def test_layers_record_direction():
    lat = td.Lattice(shape=(2, 3), names=("h", "w"))
    model = td.LSTM(4, 4, lat, bidirectional=True)
    assert [layer["reverse"] for layer in td.spec(model)["layers"]] == [
        False,
        False,
        True,
        True,
    ]


def test_layer_count_matches_the_model():
    assert td.spec(td.LSTM(4, 7, td.Lattice(shape=(2, 3))))["model"]["n_layers"] == 7


# -- the thing the viewer exists to show -------------------------------------


def test_directions_report_axes_pinned_to_one_way():
    """An axis swept only forward is invisible in code and obvious in a
    picture. The spec makes it data so the picture can show it."""
    lat = td.Lattice(shape=(2, 3), names=("h", "w"), time=True)
    model = td.LSTM(4, 6, lat, bidirectional=("h", "w"))
    directions = td.spec(model)["sweeps"]["directions"]
    assert directions["time"] == "forward"
    assert directions["h"] == "both" and directions["w"] == "both"


def test_directions_distinguish_backward_only():
    lat = td.Lattice(shape=(2, 3), names=("h", "w"))
    plan = td.ScanPlan.from_list([("h", True), ("w", False)])
    directions = td.spec(td.LSTM(4, 2, lat, plan=plan))["sweeps"]["directions"]
    assert directions["h"] == "backward" and directions["w"] == "forward"


def test_unswept_axes_are_listed():
    lat = td.Lattice(shape=(2, 3), names=("h", "w"))
    with pytest.warns(UserWarning):
        model = td.LSTM(4, 1, lat, plan=td.ScanPlan.from_list(["h"]))
    assert td.spec(model)["sweeps"]["unswept_axes"] == ["w"]


def test_nothing_unswept_when_the_plan_covers_the_lattice():
    lat = td.Lattice(shape=(2, 3), names=("h", "w"))
    assert td.spec(td.LSTM(4, 2, lat))["sweeps"]["unswept_axes"] == []


# -- io ----------------------------------------------------------------------


def test_io_shapes_are_symbolic_in_batch_and_time():
    lat = td.Lattice(shape=(3, 2), time=True)
    io = td.spec(td.LSTM(8, 3, lat, d_input=2))["io"]
    assert io["input"] == ["B", "T", 3, 2, 2]
    assert io["output"] == ["B", "T", 3, 2, 8]
    assert io["cells_per_step"] == 6


def test_io_omits_time_when_the_lattice_has_none():
    lat = td.Lattice(shape=(3, 2))
    assert td.spec(td.LSTM(8, 2, lat))["io"]["input"] == ["B", 3, 2, 8]


def test_d_input_defaults_to_d_model_without_a_projection():
    model = td.LSTM(8, 2, td.Lattice(shape=(3, 2)))
    assert td.spec(model)["model"]["d_input"] == 8


def test_the_one_dimensional_model_still_describes_itself():
    io = td.spec(td.LSTM(4, 2))["io"]
    assert io["input"] == ["B", "T", 4] and io["cells_per_step"] == 1


def test_reported_parameter_count_matches_the_module():
    model = td.LSTM(8, 3, td.Lattice(shape=(2, 3)), d_input=2)
    assert td.spec(model)["model"]["n_params"] == sum(p.numel() for p in model.parameters())


def test_gru_is_described_too():
    spec = td.spec(td.GRU(4, 2, td.Lattice(shape=(2, 3))))
    assert spec["model"]["kind"] == "GRU"
    assert spec["layers"][0]["mixer"] == "GRUMixer"


# -- the spec must describe the family that actually runs ---------------------


def test_the_kernel_family_spec_does_not_claim_spatial_sweeps():
    """A kernel-family layer contracts every spatial axis at once and sweeps
    only time. The spec used to describe it with the scan family's schema, so
    a 3-layer CaFA model claimed to sweep time, then h, then w — three sweeps
    that never happen — and the viewer drew them (DEBUG.md #26).
    """
    lat = td.Lattice(shape=(4, 5), names=("h", "w"), time=True)
    s = td.spec(td.LSTM(8, 3, lat, method=td.cafa))

    assert s["nd_method"]["family"] == "kernel"
    assert [layer["kind"] for layer in s["layers"]] == ["kernel"] * 3
    assert {layer["axis"] for layer in s["layers"]} == {"time"}
    assert all(layer["contracted"] == ["h", "w"] for layer in s["layers"])
    assert s["sweeps"]["contracted_axes"] == ["h", "w"]
    # A contraction has no direction; only the swept axis gets one.
    assert s["sweeps"]["directions"] == {"time": "forward"}
    # Contracted axes are mixed, so they are not "unswept" in the sense the
    # viewer warns about.
    assert s["sweeps"]["unswept_axes"] == []


def test_a_kernel_only_block_reports_no_swept_axis_at_all():
    lat = td.Lattice(shape=(4, 5), names=("h", "w"))
    block = td.AxialKernel(
        mixer=None, plan=td.ScanPlan.cyclic(lat.axis_names, 2), lattice=lat, d_model=8
    )

    class Wrapper(torch.nn.Module):
        lattice = lat

        def __init__(self):
            super().__init__()
            self.nd = block

        def to_spec(self):
            from torch_dimensions.spec import scan_model_spec

            return scan_model_spec(self)

    s = td.spec(Wrapper())
    assert [layer["axis"] for layer in s["layers"]] == [None, None]
    assert [layer["mixer"] for layer in s["layers"]] == [None, None]
    assert s["sweeps"]["directions"] == {}
    assert s["sweeps"]["contracted_axes"] == ["h", "w"]


def test_the_scan_family_spec_is_unchanged_in_shape():
    lat = td.Lattice(shape=(4, 5), names=("h", "w"), time=True)
    s = td.spec(td.LSTM(8, 3, lat))
    assert s["nd_method"]["family"] == "scan"
    assert [layer["axis"] for layer in s["layers"]] == ["time", "h", "w"]
    assert [layer["kind"] for layer in s["layers"]] == ["scan"] * 3
    assert s["sweeps"]["contracted_axes"] == []