Shape Algebra Through GNN Layers: Tracking Tensor Dimensions in TGraphX

In a flat-vector GNN you rarely think about shapes: everything is [N, D] and the only number that changes is D. In a tensor-aware framework like TGraphX, node features carry spatial axes, so a layer transforms [N, C, H, W] into [N, C', H, W], a pooling step collapses [N, …] into [B, …], and getting any of those wrong produces an error several lines later. The good news is that the shape transformations follow a small, predictable algebra. This tutorial gives you that algebra and the tools to debug it.

If you have read Tensor Shapes as an API Contract, this is the hands-on companion that traces shapes through a concrete pipeline.

The two-shape convention

TGraphX layers declare their input and output shapes explicitly. In tgraphx/layers/base.py, TensorMessagePassingLayer.__init__ takes in_shape and out_shape — the per-node trailing shapes, not including the node-count axis. So a convolutional message-passing layer built as ConvMessagePassing(in_shape=(16, 8, 8), out_shape=(32, 8, 8)) consumes node features [N, 16, 8, 8] and emits [N, 32, 8, 8]. The leading N (node count) and, after batching, the disjoint-union total are handled by the framework; you reason only about the trailing shape.

This convention is the single most useful thing to internalise. Every shape question reduces to: what is the per-node trailing shape here, and which axis is the layer changing?

A worked pipeline

Consider 1000 image-patch nodes, each [3, 8, 8], classified into 10 classes. Here is the shape at every stage:

Two kinds of shape change appear. Message-passing layers keep the rank and edit the channel axis (3 → 32 → 64) while preserving H and W — that is the 1×1-convolution behaviour explained in Inside Tensor Message Passing. The transitions that drop axes are the deliberate ones: a spatial reduction collapses H, W, and global_mean_pool (from tgraphx/layers/pooling.py) collapses the node axis to a per-graph vector using the batch vector. Knowing which step removes which axis is the whole game.

Pooling: node axis vs spatial axis

It is easy to conflate two different reductions. Spatial pooling reduces H and W inside each node, turning a feature map into a vector while keeping N intact. Graph pooling — global_mean_pool, global_sum_pool, global_max_pool — reduces the node axis to one row per graph, using the batch vector to know which nodes belong together. A node-classification model usually does spatial pooling then a per-node head; a graph-classification model does spatial pooling then graph pooling then a per-graph head. Mixing these up is the most common shape bug in graph-level models.

Debugging shape mismatches

When shapes disagree, TGraphX gives you a guard rather than a cryptic CUDA error. tgx.assert_tensor_native(g, min_rank=3) (in tgraphx/ux/validation.py) asserts that node features have at least the rank you expect — useful at the top of a function to fail fast if someone hands you flattened [N, D] data where [N, C, H, W] was required:

import tgraphx as tgx
        tgx.validate_graph(g, strict=True)
        tgx.assert_tensor_native(g, min_rank=3)   # raise if nodes are not tensor-valued
        

The discipline that prevents most mismatches is simply to write the expected trailing shape in a comment beside each layer, then check it with feature_shape on a sample graph. Because layers declare in_shape/out_shape, a mismatch is usually a typo in one of those tuples, not a deep bug. For the validation philosophy behind these guards, see Shape-Aware Validation.

Honest notes

Two caveats. First, the in_shape/out_shape convention means you cannot rely on lazy shape inference the way some frameworks allow — you state the shapes, which is more verbose but also more auditable. Second, spatial dimensions must be consistent across a message-passing layer (the convolution preserves H, W); if you need to change spatial resolution, do it explicitly with a pooling or resizing step rather than expecting a message-passing layer to do it implicitly.

Two habits that prevent shape bugs

Two small habits remove most shape errors before they happen. First, annotate every layer with its expected trailing shape in a comment, exactly as the worked pipeline above does — when a mismatch occurs, the comment tells you instantly which in_shape/out_shape tuple is wrong. Second, probe a single sample graph before wiring a full training loop: construct one Graph, print g.feature_shape, push it through the model, and check the output shape against what you expect.

Because TGraphX validates on construction and layers declare their shapes, a probe like this surfaces a typo in seconds rather than after a long data-loading step. If you ever find that runtime is sensitive to a reshape you added, profile that step against the source rather than guessing — a stray view or permute between layers is a common and cheap-to-find culprit. The discipline is unglamorous, but tracing the trailing shape stage by stage is what turns tensor GNNs from fiddly into routine.

Related guides

• Tensor Shapes as an API Contract in Graph Learning
• Shape-Aware Validation in TGraphX
• Inside Tensor Message Passing: The φ → AGG → ψ Pipeline

Conclusion

Tensor GNNs add a spatial dimension to the bookkeeping, but the rules are simple: message-passing layers edit the channel axis and preserve H, W; pooling steps deliberately drop the spatial or node axes. Declare in_shape/out_shape, trace the per-node trailing shape stage by stage, and guard the entry points with assert_tensor_native. Do that and shape errors become a five-second fix instead of a debugging session.