tests/test_flatten.py

def test_joint_attention_is_permutation_invariant_without_positions():
    """`td.flatten` adds no positional term, so with a set-function mixer the
    whole model is a set function: permuting an axis and un-permuting the
    output changes nothing. This is why a joint-attention model cannot learn
    a cumulative sum along an axis, and why `td.ViT` carries its own
    positional embedding. Documented in `Flatten`'s docstring; pinned here so
    the claim cannot quietly stop being true."""
    import torch

    import torch_dimensions as td

    torch.manual_seed(0)
    dense = td.Lattice(shape=(6, 8), names=("h", "w"), time=True)
    model = td.Transformer(32, 2, dense, d_input=1, method=td.flatten).eval()

    x = torch.randn(1, 3, 6, 8, 1)
    perm = torch.randperm(8)
    with torch.no_grad():
        straight = model(x)
        shuffled = model(x[:, :, :, perm])[:, :, :, torch.argsort(perm)]
    assert torch.allclose(straight, shuffled, atol=1e-5)