TGraphX Insights Writing Source-Grounded Documentation for a GNN Library
← Back to Insights

Writing Source-Grounded Documentation for a GNN Library

Target keyword: GNN library documentation source grounded

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:

bash
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: exports env_report, recommended_device,
    estimate_message_memory — document each with signature and purpose
  • tgraphx.reproducibility: exports set_seed — document the
    deterministic=True behavior 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:

markdown
**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:

  1. Include version numbers in examples where features have version notes
  2. Use the library's own shape validation — let the Graph constructor
    catch errors in examples, not separate assertion logic
  3. Run examples in CI if possible (doctest or a test file per tutorial)
python
# 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:

python
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.