Why Explicit Graph APIs Help LLM-Assisted Research Coding
Research code written with AI assistance (GitHub Copilot, Claude, GPT-4, etc.)
has a specific failure mode: the generated code looks syntactically correct but
makes wrong assumptions about tensor shapes. In graph neural network code,
this is especially common because the conventions — edge index orientation,
batching, feature dimensions — vary across libraries and are rarely spelled out
in the code itself.
Explicit APIs help here in a concrete way: they give the language model (and
the human reviewer) a contract to read.
The Implicit Convention Problem
Consider this function signature in a hand-rolled GNN:
def message_pass(node_feats, adj, edge_feats):
# ... implementation
An LLM generating the body of this function has to guess:
- Is adj an adjacency matrix [N, N] or an edge list [2, E]?
- Is it dense or sparse?
- Is node_feats [N, d] or [d, N]?
- Are edge_feats aligned with rows or columns of adj?
Different GNN libraries use different conventions. An LLM trained on code from
PyG, DGL, and hand-rolled implementations will generate code that mixes these
conventions in plausible but incorrect ways.
What Explicit Contracts Provide
TGraphX's Graph class makes the conventions explicit in the object's type:
from tgraphx import Graph
# The type itself documents the contract:
# node_features: [N, *]
# edge_index: [2, E] — row 0 is source, row 1 is target
# edge_features: [E, *]
# node_labels: [N, *]
g = Graph(
node_features=feats, # [N, d] or [N, C, H, W] etc.
edge_index=edge_index, # [2, E]
edge_features=e_feats, # [E, de]
node_labels=labels, # [N] or [N, K]
)
An LLM generating code that uses g.edge_index knows the shape is [2, E]
from documentation and type annotations. It knows g.node_features[i] gives
the features of node i. The contract is readable from the object's attribute
names and documented shapes, not implied by argument order.
Workflow Diagram
LLM-Assisted GNN Development Workflow
───────────────────────────────────────────────────────────────
Human: "Build a 2-layer SAGE model for node classification
on graphs with image-patch nodes [N, 3, 16, 16]"
↓
LLM reads: Graph API contract
node_features: [N, C, H, W]
edge_index: [2, E]
Available layers: TensorGraphSAGELayer
↓
LLM generates: Model code with correct tensor shapes
↓
Human: Runs skeleton with shape assertions:
assert out.shape == (N, num_classes)
↓
If assertion fails → explicit error with shape info
If assertion passes → proceed to training
───────────────────────────────────────────────────────────────
Concrete Example: Where Implicit Code Fails
Suppose you ask an LLM to write a message passing aggregation for a custom
layer, and your graph is stored as raw tensors:
# Raw tensor approach — conventions unclear
node_feats = torch.randn(100, 64) # could be [N, d] or [d, N]
edge_index = torch.randint(0, 100, (200, 2)) # could be [E, 2] or [2, E]
The LLM may generate:
# Generated code — wrong convention assumed
agg = torch.zeros_like(node_feats)
agg.index_add_(0, edge_index[:, 1], node_feats[edge_index[:, 0]])
This code assumes edge_index is [E, 2]. If your convention is [2, E],
this silently computes the wrong aggregation.
With TGraphX:
g = Graph(node_features=node_feats, edge_index=edge_index_2E, ...)
# LLM knows edge_index is [2, E]
# Generates:
agg = torch.zeros_like(g.node_features)
agg.index_add_(0, g.edge_index[1], g.node_features[g.edge_index[0]])
# row 0 = source, row 1 = target — contract is readable
The explicit g.edge_index[0] (sources) and g.edge_index[1] (targets) are
self-documenting in a way that edge_index[:, 0] vs edge_index[0, :] is not.
Shape Errors Surface Earlier
When LLM-generated code makes a shape mistake, the earlier it surfaces the
better. TGraphX's construction-time validation catches device and shape
inconsistencies immediately:
# This raises an error at construction, not during the forward pass
g = Graph(
node_features=torch.randn(100, 64),
edge_index=torch.randint(0, 100, (200, 2)), # wrong! should be [2, E]
...
)
# → Error: edge_index must have shape [2, E], got [200, 2]
This error message points directly at the construction code, which is where
the LLM generated the wrong shape. The human can immediately see and fix it.
Readable Code for Code Review
LLM-generated code needs human review. Explicit contracts make this review
faster. Compare:
Hard to review:
def agg(x, ei, ef):
m = x[ei[0]] + ef
return scatter_mean(m, ei[1], dim_size=x.shape[0])
Easier to review:
def agg(graph: Graph) -> torch.Tensor:
src_feats = graph.node_features[graph.edge_index[0]] # [E, d]
messages = src_feats + graph.edge_features # [E, d]
return scatter_mean(
messages,
graph.edge_index[1], # target node indices
dim_size=graph.node_features.shape[0],
)
The second version is longer, but every variable name and index communicates
intent. A reviewer (human or AI) can check correctness from the code itself.
Using write_experiment_config for Audit Trails
When using LLM-assisted code in research, saving the model and data
configuration explicitly matters for reproducibility:
from tgraphx import write_experiment_config
write_experiment_config({
"model": "LLM-generated 2-layer SAGE",
"node_feature_shape": list(g.node_features.shape),
"edge_index_shape": list(g.edge_index.shape),
"generated_by": "claude-3.5-sonnet",
"reviewed_by": "human",
"shape_assertions": "passed",
}, path="./logs/config.json")
This creates an audit trail: what shapes were used, whether the shape checks
passed, and who reviewed the generated code. For research reproducibility, this
metadata is as important as the hyperparameters.
Summary
Explicit graph APIs do not prevent LLMs from generating wrong code. But they
reduce the frequency of shape convention errors by making contracts readable
in the code, surface errors earlier through construction-time validation, and
make human review faster by using self-documenting attribute names. For
research coding workflows where AI assistance is common, these properties
compound into meaningfully fewer debugging sessions.