Tensor Shapes as an API Contract in Graph Learning

Shape mismatches in neural networks fail in one of two ways: loudly, with an
error that points you to the right line of code; or silently, producing results
that look plausible but are computed from wrong data. Graph neural networks
have more moving parts than standard feed-forward networks — node features,
edge indices, batching offsets — and the silent failures are correspondingly
harder to diagnose.

This article argues that treating tensor shapes as a first-class API contract,
and enforcing them eagerly at object construction, is worth the engineering cost.
TGraphX's Graph class is used as a concrete example of how this looks in
practice.

The Silent Failure Problem

Consider a standard error that occurs in hand-rolled GNN code: you pass an
edge index with shape [E, 2] instead of [2, E]. In many message-passing
implementations, this doesn't raise an error. Instead, the indexing operation
silently uses the wrong dimension, producing a graph where edges reference
nonexistent or wrong nodes. The model trains — loss decreases — but it's
learning from a corrupted graph structure.

This class of error is particularly insidious because:

1. The error is silent — no exception is raised
2. The symptom (strange accuracy or slow convergence) doesn't point at the cause
3. The bug may only surface when you change batch size or graph size

Shape contracts prevent this by making the wrong construction impossible.

What a Shape Contract Looks Like

A shape contract is a specification that says: "for this object to be valid,
its components must have these shapes, in this order, with this consistency."

For TGraphX's Graph class:

node_features: [N, ...]  — first dim is node count
        edge_index:    [2, E]    — exactly 2 rows (src, dst); E columns
        edge_features: [E, ...]  — first dim must match edge_index columns
        node_labels:   [N, ...]  — first dim must match node_features first dim
        

These constraints are validated at construction. If you build a Graph and
the shapes don't satisfy these constraints, you get an error immediately —
before any forward pass, before any optimization step, before any silent
corruption can propagate.

The ASCII Shape Contract Diagram

  Shape Contract for Graph Construction
          ──────────────────────────────────────
        
          node_features  [N, *]  ──── N must equal ────►  node_labels [N, *]
                │
                │  N = number of nodes
                │
          edge_index  [2, E]  ──── E must equal ────►  edge_features [E, *]
                │
                │  row 0: source node indices ∈ [0, N)
                │  row 1: target node indices ∈ [0, N)
                │
          All tensors must be on the same device.
          ──────────────────────────────────────
        

Why Enforce at Construction, Not at Forward Pass?

An alternative approach is to validate shapes during the forward pass of a
layer. This is common in PyTorch layers, where shape errors surface when
F.linear receives mismatched inputs.

The problem with forward-pass validation for graph objects is that the graph
object is often constructed once and reused across many forward passes, possibly
in multiple workers. If validation only happens during the forward pass:

• Errors surface in worker processes, producing hard-to-read stack traces
• You may not catch the error until a specific batch triggers it
• The error message is about a low-level tensor operation, not about your
  graph's structure

Construction-time validation catches the error at the point where you have the
most context: you just built the Graph, you know where the tensors came from,
and the stack trace points directly at your construction code.

Device Contracts

Shapes are one dimension of the API contract; devices are another. A graph
where node_features is on CPU and edge_index is on CUDA will fail at the
first scatter operation, producing an error about device mismatch deep in the
message-passing layer.

TGraphX enforces device consistency at construction: all four tensors must be on
the same device. The .to(device) method moves all tensors together, making
it impossible to end up with a partially-moved graph.

from tgraphx.performance import recommended_device
        from tgraphx import Graph
        import torch
        
        device = recommended_device()
        
        g = Graph(
            node_features=node_features.to(device),
            edge_index=edge_index.to(device),
            edge_features=edge_features.to(device),
            node_labels=node_labels.to(device),
        )
        # Or more concisely:
        g = Graph(...).to(device)  # moves all tensors at once
        

Tensor-Valued Features and Shape Generality

The [N, ...] notation for node_features means the API contract is general:
the first dimension is always nodes, but the trailing dimensions can be
anything. This matters for research workflows where nodes carry structured
data: images [N, C, H, W], time series [N, T, D], or simple vectors
[N, d] all satisfy the contract.

The contract does not change based on what the trailing dimensions are. A layer
that operates on node features handles the trailing dimensions according to its
own design. The Graph class is responsible only for ensuring N is consistent
and devices match.

The Cost of Enforcement

Construction-time validation adds a small overhead: shape checks and device
checks are run once per Graph object. For graphs constructed in a data
loading pipeline, this cost is paid in the worker process, not on the GPU.

In practice, this overhead is negligible compared to the cost of debugging a
silent shape error that has been hiding in your experiment for three days.

When Shape Contracts Are Not Enough

Shape contracts catch structural inconsistencies. They do not catch:

• Semantically wrong tensors with correct shapes (a label tensor with the
  right shape but wrong class indices)
• Numerically degenerate inputs (all-zero features, NaN values)
• Logical errors in edge index construction (self-loops when you didn't want
  them, disconnected graphs when connectivity is required)

For these, you need additional validation logic or unit tests. Shape contracts
are a necessary but not sufficient condition for correctness.

Practical Recommendation

When writing graph learning code, treat shape specification as documentation
that is also checked at runtime. Write out the expected shape as a comment
alongside your construction code, then verify it with an assertion or by
relying on the library's constructor validation:

# node_features: [N, C, H, W] where N=num_patches, C=3, H=W=16
        node_features = extract_patches(image)  # [N, 3, 16, 16]
        assert node_features.ndim == 4, f"Expected 4D, got {node_features.ndim}D"
        
        # edge_index: [2, E] from k-NN graph
        edge_index = build_knn_graph(node_features, k=4)  # [2, E]
        
        # Graph constructor validates everything
        g = Graph(node_features=node_features, edge_index=edge_index, ...)
        

This pattern makes your shape assumptions explicit, catches violations early,
and leaves a documentation trail for anyone reading the code later — including
future you.