Inside Tensor Message Passing: The φ → AGG → ψ Pipeline in TGraphX

Message passing is the computational core of most graph neural networks, and almost every framework expresses it as the same three-step recipe: compute a message along each edge, aggregate the incoming messages at each node, then update the node's state. What changes between frameworks is the type of object that flows through those three steps. In TGraphX, that object is a multi-dimensional tensor — a node feature can be a [C, H, W] feature map rather than a flat [D] vector — and the three steps are written so that spatial structure survives every hop. This note works through the math of that decomposition and shows exactly where it lives in the source.

This article is written for readers who already know what a tensor-valued node is. If that idea is new, start with Tensor-Valued Nodes in Graph Neural Networks and then come back.

The three-function decomposition

Let a graph have node states X_i for each node i, and let E_{i→j} denote an optional feature on the edge from i to j. One round of message passing is:

M_{i→j} = φ( X_i , X_j , E_{i→j} )          # message along each edge
        A_j     = AGG_{ i ∈ N(j) }  M_{i→j}          # aggregate over the neighbourhood
        X'_j    = ψ( X_j , A_j )                     # update the node state
        

Three symbols carry all the modelling choices. φ (the message function) decides what one neighbour tells another. AGG is a permutation-invariant reduction over the neighbour set N(j) — order must not matter, because a graph has no canonical neighbour ordering. ψ (the update function) combines a node's old state with the aggregated evidence to produce its new state. This is the standard message-passing template described in the PyTorch Geometric documentation; TGraphX's contribution is the choice to let X_i, E_{i→j}, and the message M_{i→j} all retain their tensor shape.

In a flat-vector GNN, X_i ∈ ℝ^D and M_{i→j} ∈ ℝ^{D'}. In TGraphX, X_i ∈ ℝ^{C×H×W} and the message keeps the same rank: M_{i→j} ∈ ℝ^{C'×H×W}. The spatial axes H and W are never flattened away, so φ can apply convolution-style operators and ψ can keep reasoning about local spatial neighbourhoods inside each node.

Implementation anchor

In the TGraphX source, this decomposition is not a metaphor — it is the literal method layout of the base class. tgraphx/layers/base.py defines TensorMessagePassingLayer, and its forward() runs exactly the three steps:

# tgraphx/layers/base.py (paraphrased control flow)
        messages   = self.message(src_features, dest_features, edge_features)   # φ
        # optional per-edge scalar weighting broadcasts across the message tensor
        aggregated = self.aggregate(messages, edge_index, num_nodes)            # AGG
        updated    = self.update(node_features, aggregated)                     # ψ
        

The mapping from math to code is one-to-one:

aggregate() is where permutation invariance is enforced. For sum it scatters edge messages into a zeroed [N, ...] buffer with index_add; for mean it divides by per-node in-degree; for max it uses torch.scatter_reduce_(reduce='amax'). None of these depend on the order edges are listed in, which is precisely the invariance the math requires.

Where the tensor shape actually matters

The base class is generic over shape, but the concrete layers are where "tensor-aware" becomes real. ConvMessagePassing in tgraphx/layers/conv_message.py builds its message with a 1×1 convolution: self.conv = nn.Conv2d(conv_in_channels, out_channels, kernel_size=1). A 1×1 kernel mixes channels at every spatial location independently, so a message map [C', H, W] is produced for each edge while H and W pass through untouched. That is the concrete sense in which, as the TGraphX README puts it, "every neighbourhood aggregation step acts like a miniature CNN across neighbouring feature maps."

AttentionMessagePassing (tgraphx/layers/attention_message.py) follows the same skeleton but computes a gated message. It projects source and destination maps with 1×1 convolutions to query/key/value tensors, forms a score by summing the query–key product over the channel axis and scaling by 1/√d, applies sigmoid(leaky_relu(·)), and multiplies the value map by that spatial coefficient. The score tensor has shape [E, 1, H, W]: one coefficient per spatial location per edge. It is worth being precise here — this is a sigmoid spatial gate, not a softmax-normalised attention distribution over neighbours, and the surviving aggregation step is still the permutation-invariant reduction above.

What this design does not promise

A faithful description has to include the costs. Materialising a message tensor [E, C', H, W] for every edge is more memory-hungry than a [E, D'] vector message, so ConvMessagePassing ships a _chunked_forward path that processes edges in batches to bound peak memory. The package does not claim a universal accuracy improvement from tensor messages; whether preserving spatial structure helps is task-dependent and should be measured, not assumed. If you care about runtime or memory, profile your specific graph — the source exposes the hooks, but it does not ship a benchmark that would let anyone claim a speedup in the abstract.

How to describe this accurately

A conservative, citation-ready summary: TGraphX implements message passing as a message/aggregate/update decomposition in which node and edge features may be rank-3 or rank-4 tensors; concrete layers such as ConvMessagePassing use 1×1 convolutions so that spatial dimensions are preserved through aggregation, and aggregation operators (sum, mean, max) are permutation-invariant. That sentence is fully backed by tgraphx/layers/base.py and tgraphx/layers/conv_message.py, and it is the level of claim the framework supports.

Related guides

• Tensor-Valued Nodes in Graph Neural Networks: Why Shape Matters
• Reading the TGraphX Source: Tensor-Valued Graph Representation
• Tensor Shapes as an API Contract in Graph Learning

Conclusion

The φ → AGG → ψ template is shared across the GNN world; TGraphX's specific choice is to run it over tensors instead of vectors, and to make that choice visible in the method names of TensorMessagePassingLayer. Reading message, aggregate, and update side by side with the math is the most direct way to understand what the framework guarantees — shape-preserving messages and order-independent aggregation — and what it leaves to the practitioner to measure.