Tensor-Valued Nodes in Graph Neural Networks: Why Shape Matters

The standard assumption in graph neural network (GNN) frameworks is that every node is described by a flat feature vector. You get a node feature matrix of shape [N, D], message passing operates on that matrix, and everything is comfortably linear-algebra friendly. This works very well for citation networks, social graphs, and most molecular graphs.

It does not work when your nodes have real structure. An image patch is [C, H, W]. A volumetric block from a medical scan is [C, D, H, W]. A temporal sequence is [T, D]. Flattening these into vectors is technically possible, but it discards spatial or temporal correlation that the model could have used.

TGraphX (pip install tgraphx) is designed around the assumption that node features should be allowed to keep their natural tensor shape. This article explains what that means in practice and where the design helps.

The flat-vector assumption

A typical PyG or DGL pipeline looks like this:

import torch
        x = torch.randn(100, 32)          # 100 nodes, 32 features each
        edge_index = torch.randint(0, 100, (2, 300))
        

Every layer expects x to be rank-2. Message passing aggregates neighbor features into the same shape. Convolutions, attention, all aggregation primitives are defined on [N, D].

This is fine when the features are genuinely vector-shaped. A bag-of-words representation of a citation node, a one-hot encoding of an atom type, a learned embedding of a user — these are all naturally vectors. The framework's assumption matches the data.

Where the assumption breaks

Consider a graph of image regions. Each node represents a cropped patch from an image, and edges link adjacent or visually similar patches. The natural representation of a node feature is [3, 8, 8] — three color channels times spatial dimensions.

Two options exist:

1. Flatten it. Reshape every patch to a [192] vector before building the graph. Lose the spatial layout entirely. Any model that relies on convolutional inductive bias has to relearn that bias from scratch — or you give up and use a simpler MLP-style layer.

2. Preserve the shape. Keep x as [N, 3, 8, 8] and define message passing layers that apply convolution-style operations during aggregation.

Option two is what TGraphX is built around.

How TGraphX preserves shape

TGraphX provides a Graph container that accepts node features of any rank:

import torch
        import tgraphx as tgx
        
        # 100 nodes, each an 8x8 RGB patch
        x = torch.randn(100, 3, 8, 8)
        edge_index = torch.randint(0, 100, (2, 300))
        labels = torch.randint(0, 5, (100,))
        
        g = tgx.Graph(x=x, edge_index=edge_index, labels=labels)
        tgx.validate_graph(g, strict=True)
        

The validation call confirms shapes are consistent and edge indices are within bounds. From here, you can use one of TGraphX's tensor-aware layers — ConvMessagePassing, TensorGATLayer, TensorGraphSAGELayer, or TensorGINLayer — and the spatial dimensions survive through every aggregation step.

For a complete workflow, the one-call API handles graph building, model selection, training, and evaluation:

result = tgx.classify_nodes(
            x=x, edge_index=edge_index, labels=labels,
            model="tensor_gcn", seed=42, device="auto",
        )
        print(result.metrics["val_accuracy"])
        

When does this design actually help

Honest answer: not always. Flat vectors are sufficient for most published GNN benchmarks. Choose tensor-valued nodes when:

• Your nodes represent image patches, volumetric regions, or other inherently spatial data
• You want to apply convolutional operations inside message passing
• You need to preserve channel-wise structure across aggregation
• Multi-modal graphs where different node types carry different tensor modalities

If your nodes are user IDs, atom types, or text embeddings, you do not need this. Use PyG with flat features and skip the complexity.

Trade-offs to acknowledge

Tensor-valued nodes are not free:

• Memory. A [N, C, H, W] feature tensor is much larger than [N, D] for the same node count.
• Layer complexity. Tensor-aware layers do more work per aggregation step.
• Ecosystem. PyG has many more pre-built layers and benchmark integrations. TGraphX is a smaller, newer project.
• Maturity. Several TGraphX subsystems are labeled Experimental — neural generation, heterogeneous GNNs, distributed DDP. Stick to Beta-marked core for production-adjacent work.

The framework's source repository documents stability levels per module in docs/api_stability.md.

Practical guidance

If you are starting a new project where node features have spatial structure, write a 50-line prototype with tgx.classify_nodes() first. If the validation passes and a small model trains sensibly on a subset of your data, the design fits. If you spend more time wrestling with shape errors than building the model, the data probably suits flat vectors better.

The framework is published at arXiv:2504.03953 (Sajjadi & Eramian, 2025) and the code is on PyPI as tgraphx version 1.4.2.

FAQ

Q: Can I mix tensor-valued and flat-vector nodes in the same graph?
A: Heterogeneous graphs are labeled Experimental in TGraphX. Single-graph mixing of feature ranks is not supported at the layer level. If you need this, consider keeping subgraphs separate or projecting all features to a common shape.

Q: Do I need to flatten before any operation?
A: No, not for the tensor-aware layers shipped in TGraphX. You can use ConvMessagePassing or TensorGATLayer on rank-4 inputs directly. Flattening is required only when you mix in a vector-only layer like GCNConv.

Q: Are there pre-trained models that use tensor-valued nodes?
A: Not in TGraphX itself. The framework is research-oriented and ships no pre-trained checkpoints. You train your own models on your own data.

Q: How do I save and load a graph with tensor features?
A: Use the .tgx native format: g.save("file.tgx") and tgx.Graph.load("file.tgx"). Standard graph formats like GraphML cannot represent rank-4 tensors.