PyTorch Geometric (PyG)

PyG is the standard library for Graph Neural Networks built on PyTorch. It provides data structures for graphs, 60+ GNN layer implementations, scalable mini-batch training, and support for heterogeneous graphs.

Install: uv add torch_geometric (or uv pip install torch_geometric; requires PyTorch). Optional: pyg-lib, torch-scatter, torch-sparse, torch-cluster for accelerated ops.

Core Concepts

Graph Data: Data and HeteroData

A graph lives in a Data object. The key attributes:

from torch_geometric.data import Data

data = Data(
    x=node_features,          # [num_nodes, num_node_features]
    edge_index=edge_index,     # [2, num_edges] — COO format, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]