pyg-lib Roadmap

pyg-lib.md

The first pyg-lib releases focuses on unifying the implementations from torch-sparse and torch-cluster into a single package in order to reduce the number of external low-level library dependencies of PyG. In addition, implementations will be improved, e.g., by out-sourcing common routines into re-usable building blocks, unifiying the interfaces, supporting various data types, biased sampling, etc. New functionality will be integrated for temporal-based learning and GNN acceleration.

Samplers

• Unify common routines behind re-usable functions (e.g., sampling with/without replacement)
• Deterministic sampling routines
• Full support for different data types, e.g., int32, int64, etc
• Integration of cugraph dependency
• Neighbor Sampling neighbor_sample(rowptr, col, seed, num_neighbors):
  • Support for homogeneous graphs (CPU+GPU)
  • Support for heterogeneous graphs (CPU-only)
  • replace: sampling with or without replacement
  • directed: sub-tree vs sub-graph sampling (CPU-only)
  • disjoint: disjoint subtrees for every seed node (CPU-only)
  • temporal: temporal sampling (CPU-only)
  • weighted: Support for biased sampling (CPU-only)
  • temporal_weighted: Support for biased temporal sampling (CPU-only)
  • return_edge_id: Support for returning edge IDs (CPU-only)
• Subgraph Sampling subgraph_sample(rowptr, col, nodes):
  • Support for homogeneous graphs (CPU+GPU)
  • Support for heterogeneous graphs (CPU-only)
  • return_edge_id: Support for returning edge IDs (CPU-only)
• Random Walk Sampling random_walk(rowptr, col, nodes):
  • Support for homogeneous graphs (CPU+GPU)
  • Support for heterogeneous graphs (CPU-only)
  • weighted: Support for biased sampling (CPU-only)
  • node2vec-based sampling (CPU-only)
  • return_edge_id: Support for returning edge IDs (CPU-only)
• Heterogeneous Graph Transformer Sampling hgt_sample(rowptr, dict, seed):
  • weighted: Support for biased sampling (CPU-only)

Operators

• segment_matmul(src, ptr, other):
  • forward (CPU+GPU)
  • backward (CPU+GPU)
• sparse_softmax(src, index):
  • forward (CPU+GPU)
  • backward (CPU+GPU)
• Variadic EmbeddingBag:
  • forward (CPU+GPU)
  • backward (CPU+GPU)

Others

• Full support for different data types, e.g., int32, int64, etc
• METIS graph partitioning
• Farthest point sampling fps
• k-NN graph generation
• radius graph generation/ball query