The TGraphX Graph and GraphBatch: How Tensor Graphs Are Stored and Batched

If you work with TGraphX for more than a few minutes, two classes do most of the heavy lifting: Graph, which holds a single graph, and GraphBatch, which packs several graphs into one super-graph for efficient training. Understanding their fields and their batching rules removes most of the confusion that newcomers hit, because the framework is deliberately strict: it would rather raise an error at construction time than silently mangle your data. This tutorial walks through both objects as they exist in tgraphx/core/graph.py.

For the conceptual background on why nodes carry tensors at all, see Reading the TGraphX Source.

The Graph object

A Graph is constructed from node features plus an optional edge structure and labels:

import torch
        import tgraphx as tgx
        
        x = torch.randn(100, 3, 8, 8)                 # 100 nodes, each a [3, 8, 8] tensor
        edge_index = torch.randint(0, 100, (2, 300))  # 300 directed edges
        labels = torch.randint(0, 10, (100,))
        
        g = tgx.Graph(x=x, edge_index=edge_index, labels=labels)
        print(g.num_nodes, g.num_edges, g.feature_shape)   # 100 300 (3, 8, 8)
        

The constructor accepts PyG-style aliases so code generated from muscle memory tends to work: x resolves to node_features, y or labels to node_labels, and edge_attr to edge_features. Providing two conflicting aliases (say both y and labels) raises a ValueError rather than guessing. The same object exposes those aliases back as read/write properties — g.x, g.y, g.edge_attr — alongside descriptive ones:

Train/validation/test masks are not loose globals — when you pass train_mask=, val_mask=, or test_mask=, they are stored inside metadata['masks'] so they travel with the graph through moves and serialization. That single decision prevents a common bug where a split silently detaches from the data it describes.

Validation happens at construction

Every optional tensor is checked the moment you build the graph. In the source, the constructor calls validate_edge_index, validate_edge_weight, validate_edge_features, and a label checker before the object exists. A mask whose length does not match the node count, an edge_index that points past the last node, or labels with the wrong leading dimension all fail immediately with a message that names the offending tensor. This is the same philosophy covered in Shape-Aware Validation: catch the shape error where it is cheap to fix, not three layers deep in a training loop.

Batching: the disjoint union

GNN training batches graphs by placing them side by side as one large graph with no edges between the pieces — a disjoint union. GraphBatch does this in tgraphx/core/graph.py:

batch = tgx.GraphBatch([g1, g2, g3])
        print(batch.num_graphs)        # 3
        print(batch.node_features.shape)   # [N1 + N2 + N3, 3, 8, 8]
        print(batch.batch.shape)       # [N1 + N2 + N3]
        

Two things make this usable. First, edge_index is re-based: graph 2's edges are offset by N1, graph 3's by N1 + N2, so a single [2, E_total] tensor still points at the right nodes. Second, the object carries a batch vector of length N_total that maps each node back to its source graph index — exactly what pooling layers and graph-level readouts need to know which nodes belong together.

The strictness is a feature

GraphBatch validates compatibility before concatenating. All graphs must share the same per-node feature shape (node_features.shape[1:]); if graph 2 has [3, 8, 8] nodes and graph 0 has [3, 16, 16], you get a ValueError telling you to resize or pad first. The same rule applies to edge features. The source comment is blunt about why: "mixing-some-with-none is rejected rather than silently dropped, because silently dropping per-edge data is a footgun." If you have ever lost edge attributes to a quiet broadcast, you will appreciate the explicit failure.

For iterating over many graphs, tgraphx/core/dataloader.py provides GraphDataset and GraphDataLoader, which collate lists of Graph objects into GraphBatch instances so you can write an ordinary PyTorch training loop.

A small debugging tip

When a batch refuses to build, print the per-graph shapes before guessing: [tuple(g.feature_shape) for g in graphs]. Nine times out of ten the culprit is one graph whose node tensors were resized differently upstream. Because the failure is raised at GraphBatch construction with the offending index named, you rarely need a debugger.

Related guides

• Getting Started with Tensor-Valued Nodes in TGraphX
• Shape-Aware Validation in TGraphX
• Reading the TGraphX Source: Tensor-Valued Graph Representation

Conclusion

Graph stores one tensor-valued graph and validates it on construction; GraphBatch packs many into a disjoint union with re-based edges and a batch vector, refusing to proceed when shapes disagree. Both behaviours come straight from tgraphx/core/graph.py. Once you internalise the disjoint-union model and the shared-shape rule, batching tensor graphs stops being mysterious and the framework's error messages start reading like helpful advice rather than obstacles.