Reading the TGraphX Source: How a Tensor-Valued Graph Is Represented
Most graph learning libraries store node features as 2D matrices — one row per
node, one column per feature. TGraphX takes a different stance: node features
can be arbitrary tensors. A node can carry an image patch, a time-series slice,
or a multi-channel feature map. This article walks through what that design
choice looks like in the source code, and why the validation logic at
construction time is worth understanding before you write your first experiment.
The Four Core Fields
The Graph class in TGraphX is constructed from four tensors:
| Field | Shape | Meaning |
|---|---|---|
node_features |
[N, ...] |
One tensor per node; trailing dims |
| can be anything: C, H, W, T, etc. | ||
edge_index |
[2, E] |
Source and target node indices |
edge_features |
[E, ...] |
One tensor per edge; trailing dims |
| must be consistent across edges | ||
node_labels |
[N, ...] |
Per-node labels for supervised tasks |
The N dimension is the number of nodes, E is the number of edges. The ...
notation means any additional dimensions are allowed, as long as they are
consistent across all nodes (or edges).
Why Tensor-Valued Nodes Matter
Consider a graph where each node represents an image patch. In a standard GNN
library, you would flatten the patch to a 1D vector: a 3×16×16 patch becomes a
768-dimensional vector. This works, but it discards spatial structure before
the network ever sees it.
With TGraphX, the node feature for that patch can stay as shape [3, 16, 16].
The node_features tensor is then [N, 3, 16, 16]. Layers that understand
spatial structure (convolutions, pooling) can operate on those trailing
dimensions before or during message passing.
Construction Validation
TGraphX validates shapes and devices at construction time. This means that if
you create a Graph with mismatched devices — for example, node_features on
CPU and edge_index on CUDA — you get an error immediately, not at some
unpredictable point during the forward pass.
Similarly, if edge_index has shape [E, 2] instead of [2, E], the error
surfaces at construction rather than producing silently wrong results.
import torch
from tgraphx import Graph
N, E, C, H, W = 8, 12, 3, 16, 16
node_features = torch.randn(N, C, H, W) # [8, 3, 16, 16]
edge_index = torch.randint(0, N, (2, E)) # [2, 12]
edge_features = torch.randn(E, 64) # [12, 64]
node_labels = torch.randint(0, 4, (N,)) # [8]
g = Graph(
node_features=node_features,
edge_index=edge_index,
edge_features=edge_features,
node_labels=node_labels,
)
print(g.node_features.shape) # torch.Size([8, 3, 16, 16])
print(g.edge_index.shape) # torch.Size([2, 12])
ASCII Tensor Layout Diagram
Graph object
┌───────────────────────────────────────────────────────────┐
│ │
│ node_features [N, C, H, W] │
│ ┌────┬────┬────┬────┬────┬────┬────┬────┐ │
│ │ n0 │ n1 │ n2 │ n3 │ n4 │ n5 │ n6 │ n7 │ ← N=8 nodes │
│ └────┴────┴────┴────┴────┴────┴────┴────┘ │
│ each slot holds a [C, H, W] tensor │
│ │
│ edge_index [2, E] │
│ ┌──────────────────────────────────────────────────┐ │
│ │ src: [0, 1, 1, 2, 3, 4, 4, 5, 6, 7, 0, 3] │ │
│ │ dst: [1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 4, 6] │ │
│ └──────────────────────────────────────────────────┘ │
│ row 0 = source indices, row 1 = target indices │
│ │
│ edge_features [E, d] │
│ ┌──────────────────────────────────────────────────┐ │
│ │ e0 │ e1 │ e2 │ e3 │ ... │ e11 │ │ │
│ └──────────────────────────────────────────────────┘ │
│ each slot holds a [d] feature vector │
│ │
│ node_labels [N] or [N, K] │
│ ┌────┬────┬────┬────┬────┬────┬────┬────┐ │
│ │ 0 │ 2 │ 1 │ 3 │ 1 │ 0 │ 2 │ 3 │ │
│ └────┴────┴────┴────┴────┴────┴────┴────┘ │
└───────────────────────────────────────────────────────────┘
Device Consistency
When you call g.to(device), all four tensors are moved together. This
eliminates a common bug in hand-rolled code: moving model weights to GPU but
forgetting to move the edge index, or vice versa.
from tgraphx.performance import recommended_device
device = recommended_device() # 'cuda' if available, else 'cpu'
g = g.to(device)
# All tensors — node_features, edge_index, edge_features,
# node_labels — are now on the same device.
Batching Multiple Graphs
GraphBatch.from_graphs([g1, g2, ...]) combines a list of Graph objects into
a single batched graph. The edge_index tensors are offset so that node indices
from different graphs do not collide. A batch vector tracks which node belongs
to which original graph.
from tgraphx import GraphBatch
batch = GraphBatch.from_graphs([g1, g2, g3])
# batch.node_features shape: [N1+N2+N3, C, H, W]
# batch.edge_index: offsets applied per graph
This is the standard "disjoint union" batching approach used across graph
learning libraries. TGraphX applies it to the tensor-valued setting, preserving
the trailing dimensions of node and edge features.
GraphML I/O
For loading graphs from files, TGraphX provides read_graphml and
write_graphml (Beta v1.2+). These handle graphs stored in the GraphML format,
which is common in network science tools. The resulting objects are TGraphX
Graph instances with the same tensor-valued structure.
from tgraphx.io import read_graphml, write_graphml
g = read_graphml("my_graph.graphml")
# Modify or augment g...
write_graphml(g, "output.graphml")
Note that GraphML files store numerical attributes; loading a GraphML graph
gives you tensors with scalar or vector features, not image-shaped tensors.
Tensor-valued nodes with spatial structure come from constructing Graph
objects programmatically.
What Validation Does Not Cover
Construction-time validation catches device mismatches and shape inconsistencies
between fields. It does not validate:
- Whether the edge index contains valid node indices (no out-of-bounds check
on the node dimension) - Whether node labels are appropriate for your loss function
- Whether the graph is connected, has self-loops, or is directed vs undirected
These are application-level concerns that TGraphX leaves to the user. The
library's validation is a safety net for the most common structural errors, not
a semantic validator for graph topology.
Key Takeaway
The Graph class is a thin, validated container around four tensors. Its value
is not in adding algorithmic complexity but in enforcing a consistent contract:
shapes and devices are checked at construction, all tensors move together, and
batching handles the bookkeeping of combining multiple graphs. Understanding
this structure is foundational to using any layer or utility in TGraphX
effectively.