Writing Source-Grounded Documentation for a GNN Library
Documentation that overstates a library's capabilities erodes trust faster
than no documentation at all. Research users are skeptical by default — they
will test claims, read the source, and notice when documentation makes
promises the code doesn't keep. This tutorial is for library authors and
technical writers who want documentation that is accurate, verifiable, and
useful for the people who actually read source code before trusting a claim.
The Documentation Accuracy Problem
GNN library documentation commonly fails in several ways:
| Failure mode | Example |
|---|---|
| Speed claims without evidence | "2x faster than PyG on citation networks" |
| Scope overclaiming | "supports distributed training" (with no |
| actual DDP integration) | |
| Missing stability signals | Feature described as production-ready |
| but in Beta with recent API changes | |
| Undocumented assumptions | "pass your graph" without shape spec |
| Outdated examples | Code that worked in v0.9 but not in v1.4 |
Each of these causes a specific kind of user frustration. Source-grounded
documentation prevents most of them by starting from the code rather than
from marketing goals.
Step 1: Audit What Is Actually Exported
The public API of a Python library is what is exported from __init__.py. Start
there:
python -c "import tgraphx; print(dir(tgraphx))"
Or read __init__.py directly. For each exported symbol, ask:
1. Is it documented?
2. Does the documentation match the current signature?
3. Are there stability notes where appropriate?
Everything exported is part of the public API and deserves accurate
documentation. Everything not exported is internal and should be documented
only if users need to understand it for debugging.
Step 2: Check Module Docstrings
Module-level docstrings set reader expectations. A module with no docstring
communicates nothing about its purpose or stability. For TGraphX modules:
tgraphx.performance: exportsenv_report,recommended_device,
estimate_message_memory— document each with signature and purposetgraphx.reproducibility: exportsset_seed— document the
deterministic=Truebehavior explicitly (performance tradeoff exists)tgraphx.distributed: explicitly documents that DDP is not provided —
this is a model for honest scoping in documentation
The distributed module's explicit disclaimer ("TGraphX does not ship a
distributed training framework") is more trustworthy than silence on the topic.
Document what you don't do, not just what you do.
Step 3: Mark Unverified Claims
When writing documentation, mark claims by type:
[VERIFIED FROM SOURCE] — confirmed by reading the relevant source file
[FROM DOCSTRING] — from module/function docstring, not independently verified
[VERSION NOTE] — introduced in specific version, may change
[USER RESPONSIBILITY] — TGraphX does not handle this; user must implement
For example:
**Message passing aggregation** [VERIFIED FROM SOURCE: _scatter.py, layers/]:
TensorMessagePassingLayer uses `index_add_` and `scatter_reduce_` —
no Python loop per node.
**Distributed training** [USER RESPONSIBILITY: distributed.py]:
TGraphX does not provide a distributed training framework. Users must
implement DDP or DataParallel themselves. TGraphX provides rank-zero
print helpers only.
**float16 / AMP** [VERIFIED FROM SOURCE: _scatter.py]:
The attention softmax path upcasts float16 to float32 for numerical
stability. Users running AMP will see float32 in the attention path.
This level of annotation takes more effort but produces documentation readers
can trust.
Step 4: Use Shape Specifications as Documentation
For any function or class that accepts tensors, document the expected shapes
explicitly. The standard notation is:
tensor_name: [dim1, dim2, ...]
Where each dimension is either a symbolic name or a concrete value:
node_features: [N, ...] — N nodes, any trailing dims
edge_index: [2, E] — exactly 2 rows, E columns
edge_features: [E, ...] — E must match edge_index columns
node_labels: [N, ...] — N must match node_features first dim
This is both documentation and specification. Anyone reading it knows exactly
what tensors to pass.
Step 5: Examples That Actually Run
Code examples in documentation drift. An example that worked in v0.9 may
silently produce wrong results in v1.4 if an argument name changed. Best
practices:
- Include version numbers in examples where features have version notes
- Use the library's own shape validation — let the
Graphconstructor
catch errors in examples, not separate assertion logic - Run examples in CI if possible (doctest or a test file per tutorial)
# Example from documentation — verifiable, runnable
from tgraphx import Graph
from tgraphx.performance import recommended_device, env_report
from tgraphx.reproducibility import set_seed
import torch
# Step 1: Environment and seed (always first)
print(env_report(include_hardware=True))
set_seed(42)
# Step 2: Construct graph with documented shapes
device = recommended_device()
g = Graph(
node_features=torch.randn(20, 64), # [N=20, d=64]
edge_index=torch.randint(0, 20, (2, 50)), # [2, E=50]
edge_features=torch.randn(50, 16), # [E=50, de=16]
node_labels=torch.randint(0, 3, (20,)), # [N=20]
).to(device)
# If this runs without error, shapes are correct
print("Graph constructed successfully")
print(f" node_features: {g.node_features.shape}")
print(f" edge_index: {g.edge_index.shape}")
Step 6: Distinguish Layers by Capability
TGraphX provides several layer types. Documentation should describe each
with its specific characteristics:
| Layer | Type | Notes |
|---|---|---|
| TensorGraphSAGELayer | SAGE convolution | Tensor-valued node features |
| TensorGATLayer | GAT (attention) | Attention path upcasts to f32 |
| TensorGINLayer | GIN convolution | Isomorphism-motivated |
| GCNConv | GCN convolution | Spectral-motivated |
| GATv2Conv | GATv2 (attention) | Improved attention mechanism |
| APPNP | APPNP propagation | Personalized PageRank-based |
Each entry should link to the corresponding module, show the input/output
shape contract, and note any version constraints.
Step 7: Document the Experiments Infrastructure
The experiments module contains Runner, ExperimentConfig, GridRunner,
EarlyStopping, and ModelCheckpoint. These are higher-level than layers
and need usage-oriented documentation:
from tgraphx.experiments import Runner, ExperimentConfig, EarlyStopping
config = ExperimentConfig(
model=model,
optimizer=optimizer,
criterion=criterion,
epochs=100,
)
early_stop = EarlyStopping(patience=10, monitor="val_loss")
runner = Runner(config, callbacks=[early_stop])
runner.fit(train_loader, val_loader)
Documentation for these objects needs to explain: what fit returns, when
EarlyStopping triggers, how ModelCheckpoint interacts with Runner, and
what GridRunner does differently from Runner.
Common Documentation Failures to Avoid
| Failure | Prevention |
|---|---|
| "Fast" without definition | Never use speed claims without data |
| "Supports X" for partial support | Say "provides helpers for X" or |
| "users must implement X" | |
| Example code not matching API | Test examples in CI |
| No shape specifications | Document every tensor argument's shape |
| No version notes on Beta features | Tag features with "Beta vX.Y+" |
| Missing error behavior documentation | Document what errors are raised and when |
Documentation written from the source — checking __init__.py, reading module
docstrings, verifying examples — produces a more trustworthy artifact than
documentation written from memory or from marketing intent.