"""LTI and non-LTI mixers, and what the difference does to N-D composition. The library's premise is that any 1-D mixer can be swept over a lattice. True, but not uniform: *which* mixer decides whether the sweep order is a modelling choice or a no-op. This file measures the property and then measures the consequence. See LTI.md for the table these tests produce and what it means. """ from functools import partial import pytest import torch import torch_dimensions as td from torch_dimensions.mixers.attention import AttentionMixer from torch_dimensions.mixers.conv import ConvMixer, TCNMixer from torch_dimensions.mixers.rnn import GRUMixer, LSTMMixer from torch_dimensions.mixers.ssm import MambaMixer, S4DMixer, S4Mixer D = 4 # (name, factory, expected verdict). The expectations are the claim; the # measurement is in check_lti. CLASSIFICATION = [ ("ConvMixer(activation=None)", partial(ConvMixer, D, activation=None), "LTI"), ("ConvMixer(gelu)", partial(ConvMixer, D), "time-invariant, nonlinear"), ("TCNMixer", partial(TCNMixer, D), "time-invariant, nonlinear"), ("S4DMixer", partial(S4DMixer, D), "time-invariant, nonlinear"), ("S4Mixer", partial(S4Mixer, D), "time-invariant, nonlinear"), ("MambaMixer", partial(MambaMixer, D), "neither"), ("LSTMMixer", partial(LSTMMixer, D), "neither"), ("GRUMixer", partial(GRUMixer, D), "neither"), ("AttentionMixer", partial(AttentionMixer, D, 2), "neither"), ] @pytest.mark.parametrize("name,factory,expected", CLASSIFICATION) def test_every_shipped_mixer_classifies_as_documented(name, factory, expected): report = td.testing.check_lti(factory) # "LTI (affine)" is LTI with a bias; the distinction is reported, not # asserted here. assert report.verdict.startswith(expected), f"{name}: {report}" def test_linearity_is_measured_not_assumed(): """A linear convolution is linear to floating point; adding one GELU is enough to destroy it. If both came out the same, the measurement would be measuring nothing.""" linear = td.testing.check_lti(partial(ConvMixer, D, activation=None)) nonlinear = td.testing.check_lti(partial(ConvMixer, D)) assert linear.additivity < 1e-14 assert nonlinear.additivity > 1e-3 # Both remain time-invariant: a pointwise nonlinearity does not care where # in the sequence it is applied. assert linear.time_invariant and nonlinear.time_invariant def test_a_recurrence_with_biased_gates_is_not_at_rest(): """Why the RNN and Mamba families are not time-invariant despite being perfectly causal. Time invariance is a claim about a system *at rest*: feed it nothing, then feed the signal later, and the same thing comes out later. A gated recurrence fed zeros does not sit still — its biases drive the state — so a delayed signal arrives to a different machine than the original did. """ for factory in (partial(LSTMMixer, D), partial(MambaMixer, D)): report = td.testing.check_lti(factory) assert not report.time_invariant, report # But the deviation is small: the state does settle, so this is a # transient rather than a wholesale failure. Worth stating, because # "not time-invariant" and "wildly different" are different claims. assert report.shift_equivariance < 1e-1, report def test_attention_is_not_shift_equivariant_because_padding_adds_tokens(): """Self-attention is permutation-equivariant, which is *stronger* than shift-equivariance on a fixed token set — but delaying a signal means padding, and padding a sequence adds tokens to attend over. The set changes, so the output does.""" report = td.testing.check_lti(partial(AttentionMixer, D, 2)) assert not report.time_invariant and report.shift_equivariance > 1e-3, report # -- the consequence for N-D -------------------------------------------------- def _two_orders(mixer, lat): """The same shared mixer, swept in both axis orders. A shared instance rather than two builds: the claim is about the *order*, so the operators have to be identical objects or the comparison also contains a weight difference. `norm=False` because LayerNorm is nonlinear and would break commutation on its own; the residual is linear and stays. """ a = td.AxialScan( mixer=mixer, plan=td.ScanPlan.from_list(["h", "w"]), lattice=lat, d_model=D, norm=False ) b = td.AxialScan( mixer=mixer, plan=td.ScanPlan.from_list(["w", "h"]), lattice=lat, d_model=D, norm=False ) return a.double().eval(), b.double().eval() def _order_gap(mixer, lat, x): a, b = _two_orders(mixer, lat) with torch.no_grad(): return (a(x) - b(x)).abs().max().item() / x.abs().max().item() def test_a_scalar_valued_lti_filter_commutes_across_axes(): """**The result this file exists for, and it is narrower than folklore.** "Separable convolutions commute, so the order does not matter" is true for *scalar* filters and false for the multichannel filters every real network uses. The operator commutes across axes exactly when it factorizes as a channel matrix times a spatial filter, ``W[out, in, j] = A[out, in]·b[j]``: then each axis contributes ``C_axis ⊗ A``, the spatial parts commute because they act on different axes, and the channel parts are the same matrix. Constructed rather than trained, because the point is which *structure* commutes — see the next test for what a generic filter does. """ lat = td.Lattice(shape=(4, 5), names=("h", "w")) x = torch.randn(2, 4, 5, D, dtype=torch.float64) torch.manual_seed(0) mixer = ConvMixer(D, activation=None, bias=False).double() with torch.no_grad(): channel = torch.randn(D, D, dtype=torch.float64) spatial = torch.randn(3, dtype=torch.float64) mixer.convs[0].weight.copy_(channel[:, :, None] * spatial[None, None, :]) assert _order_gap(mixer, lat, x) < 1e-12 @pytest.mark.parametrize( "name,factory", [ ("generic linear conv", partial(ConvMixer, D, activation=None, bias=False)), ("depthwise + pointwise", partial(ConvMixer, D, activation=None, depthwise=True)), ("gelu conv", partial(ConvMixer, D)), ("S4D", partial(S4DMixer, D)), ("LSTM", partial(LSTMMixer, D)), ("Mamba", partial(MambaMixer, D)), ], ) def test_everything_else_depends_on_the_sweep_order(name, factory): """Including a perfectly linear, perfectly time-invariant convolution. A generic multichannel convolution is a *matrix-valued* filter: offset ``j`` carries its own channel matrix ``W[:, :, j]``, and matrix-valued filters commute only when those matrices commute with each other, which random ones do not. So the sweep order is a real modelling choice even for a linear CNN — and ``ScanPlan`` is not only for the selective models. The depthwise case is the instructive one: the depthwise convolution alone *would* commute, and the pointwise channel mix that follows it is what breaks it. """ lat = td.Lattice(shape=(4, 5), names=("h", "w")) x = torch.randn(2, 4, 5, D, dtype=torch.float64) torch.manual_seed(0) assert _order_gap(factory().double(), lat, x) > 1e-6, name def test_direction_is_free_for_a_centred_convolution_and_not_for_a_recurrence(): """The second consequence. A centred convolution swept backwards is the same operator with a mirrored kernel, so bidirectionality buys nothing it did not already have. A causal mixer genuinely sees a different signal.""" lat = td.Lattice(shape=(7,), names=("h",)) x = torch.randn(2, 7, D, dtype=torch.float64) def swept(mixer, reverse): scan = td.AxialScan( mixer=mixer, plan=td.ScanPlan.from_list([("h", reverse)]), lattice=lat, d_model=D, norm=False, ) with torch.no_grad(): return scan.double().eval()(x) torch.manual_seed(0) # A symmetric-by-construction kernel: forward and reverse must agree # exactly, which is the sharp version of "direction is meaningless here". conv = ConvMixer(D, activation=None, bias=False).double() with torch.no_grad(): w = conv.convs[0].weight w.copy_((w + w.flip(-1)) / 2) assert (swept(conv, False) - swept(conv, True)).abs().max().item() < 1e-12 torch.manual_seed(0) lstm = LSTMMixer(D).double() assert (swept(lstm, False) - swept(lstm, True)).abs().max().item() > 1e-3