Permutation Equivariance in Tensor GNNs: What TGraphX Guarantees and Tests

A graph has no canonical ordering of its nodes. Whether you list a node first or last is an accident of bookkeeping, not a property of the data. So the central correctness requirement for a graph neural network is a symmetry: if you relabel the nodes, the network's per-node outputs should be relabelled the same way, and any graph-level output should not change at all. These two conditions are permutation equivariance (for node-level outputs) and permutation invariance (for pooled, graph-level outputs). This note states them precisely and then shows the specific TGraphX test that checks them — because a framework that claims this symmetry should be able to point at the assertion that enforces it.

The symmetry, stated

Let P be a permutation of the N nodes, acting on a node-feature tensor X by reordering the first axis, and acting on an edge list edge_index by relabelling its entries. Write a one-layer GNN as a function f(X, edge_index). The layer is permutation-equivariant if, for every permutation P,

f( P·X , P·edge_index )  =  P· f( X , edge_index )
        

In words: permute the inputs, and the output is the same as permuting the output of the original inputs. A graph-level readout g is permutation-invariant if g(P·X, P·edge_index) = g(X, edge_index) — relabelling changes nothing.

Where does equivariance come from? Entirely from the aggregation step. As covered in Inside Tensor Message Passing, TGraphX aggregates neighbour messages with order-independent reductions (sum, mean, max) implemented via index_add and scatter_reduce_ in tgraphx/layers/base.py. Because these reductions do not depend on the order edges appear in, the layer cannot "see" the node numbering except through the graph structure itself — which is exactly the condition for equivariance. The tensor shape of the node features is irrelevant to this argument: a [C, H, W] node permutes along the node axis just like a [D] vector does.

What TGraphX actually tests

Symmetry claims are cheap to make and easy to break with a careless implementation detail (an accidental sort, a stateful buffer). TGraphX ships these as executable checks in tests/test_math_invariants_v030.py, described in its own docstring as "mathematical invariants that v0.3.0 ships with stable confidence." The suite covers four families:

The first row is the headline. The test permutes node features and edge indices with a random permutation, runs the layer, un-permutes the output, and asserts it matches the layer run on the original ordering to within 1e-5. It does this for the tensor layers (ConvMessagePassing, TensorGraphSAGELayer, TensorGINLayer) and the vector model-zoo layers alike. The third row matters for attention specifically: a TensorGATLayer normalises attention with a softmax over each destination's incoming edges, and the test confirms those weights sum to one — the property that makes attention a proper weighted average rather than an arbitrary reweighting.

Why the chunked-equals-unchunked test belongs here

It is easy to overlook, but the "chunked matches unchunked" assertion is a correctness guarantee with practical teeth. Attention layers can process edges in groups to bound memory (the chunk_size argument). A memory optimisation that changed the numerical result would be a silent bug. By asserting allclose(out_full, out_chunked, atol=1e-4), the suite makes the memory path a behaviour-preserving optimisation rather than a quiet approximation — which is the kind of guarantee you want before trusting a result you could only obtain with chunking.

Honest framing

Three honest qualifications. First, these are numerical unit tests on toy graphs, not formal proofs; they demonstrate the invariants hold for the tested configurations to a stated tolerance, which is the standard, pragmatic level of assurance for research software. Second, 1e-5 and 1e-4 tolerances reflect floating-point reality — exact equality is not expected. Third, equivariance is a property of the architecture, not a promise about accuracy; a permutation-equivariant model can still be a poor fit for your task. What the tests give you is confidence that the framework is not accidentally encoding node identity, which is the foundation everything else rests on. For the broader reproducibility story these tests support, see GNN Research Reproducibility.

A check you can run yourself

You do not have to take the symmetry on faith — the same check the test suite runs is a few lines:

import torch
        perm = torch.randperm(g.num_nodes)
        # permute node features and relabel edge_index by `perm`, run the layer,
        # invert the permutation on the output, then compare to the original run.
        

The recipe is: apply a random permutation to the node features and relabel edge_index accordingly, run the layer, invert the permutation on the output, and compare to the layer run on the original ordering with torch.allclose(..., atol=1e-5). If the two match, the layer is equivariant for that input; if they diverge, something is reading node identity it should not be. Running this on your own custom layer — not just the shipped ones — is the cheapest way to catch an accidental sort or a stateful buffer before it corrupts a training run. The source's tests/test_math_invariants_v030.py is a working template you can copy and adapt to a new layer in minutes.

Related guides

• Inside Tensor Message Passing: The φ → AGG → ψ Pipeline
• Topology vs Feature Geometry in Tensor Graphs
• Research Reproducibility in GNN Projects

Conclusion

Permutation equivariance is the symmetry that makes a graph network a graph network, and in TGraphX it follows from order-independent aggregation in tgraphx/layers/base.py. The reassuring part is that you do not have to take this on faith: tests/test_math_invariants_v030.py checks equivariance, edge-order invariance, attention normalisation, and chunk-consistency for the tensor and vector layers at explicit tolerances. When a framework can name the test that guards its core symmetry, you can build on it with more confidence.