$ show pyg-team/pytorch_geometric
PyTorch Geometric (PyG) is a library for Graph Neural Networks (GNNs) built on PyTorch. It provides a modular framework for building and training GNNs, covering the full pipeline from data handling to model evaluation. The core architecture includes a message-passing base class (MessagePassing) that underpins dozens of convolution layers (e.g., GCN, GAT, SAGE, GIN, PNA, GPS), a composable aggregation system (e.g., sum, mean, attention, set2set, transformers), pooling operators (e.g., TopK, SAGPool, DiffPool), normalization layers, and high-level model implementations (e.g., SchNet, DimeNet, LightGCN, TGN, SGFormer). The library supports heterogeneous graphs, temporal graphs, point clouds, knowledge graphs, and distributed training. Key entry points include the `torch_geometric.nn`, `torch_geometric.data`, `torch_geometric.loader`, `torch_geometric.datasets`, and `torch_geometric.transforms` modules. External dependencies are primarily PyTorch itself, with optional packages (torch-scatter, torch-sparse, pyg-lib, etc.) for optimized operations. The repository also contains extensive examples, benchmarks, and an experiment management system (GraphGym). Although primarily a software library, it serves as a research platform with many paper reproductions and dataset loaders, making it a hybrid between a production-grade library and a research codebase.
The LLM identified 3 files worth highlighting for reuse:
Provides a generic benchmark function that measures forward and backward pass execution time with warmup, optional backward, and per-step reporting. Accepts a callable and its arguments, runs multiple iterations, and returns detailed timing statistics. The implementation is self-contained (only depends on torch) and can be extracted for benchmarking any PyTorch model or function.
Implements a dynamic resolver pattern that maps string names to callables (e.g., activation functions, aggregators) using a registry. Provides utility functions like `activation_resolver`, `aggregation_resolver`, and `normalization_resolver`. The design is clean, extensible, and decoupled from the rest of PyG, making it reusable for any project that needs configuration-driven component instantiation.
A modular aggregation framework with a base class `Aggregation` and implementations for sum, mean, max, min, var, std, attention, set2set, LSTMAggregation, GraphMultisetTransformer, and more. Aggregators can be combined via `MultiAggregation` with configurable reduction modes. The design demonstrates clean separation of concerns, extensibility, and reuse – the core abstractions (base class, simple aggregators, composition) could be extracted as a standalone library for tensor reduction operations.