pyg-lib 0.3.0: PyTorch 2.1 support, METIS partitioning, neighbor sampler improvements

pyg-lib==0.3.0 brings PyTorch 2.1 support, METIS partioning and further neighbor sampling improvements to PyG 🎉🎉🎉

Highlights

PyTorch 2.1 Support

pyg-lib==0.3.0 is fully compatible with PyTorch 2.1 (#256). To install for PyTorch 2.1, simply run

pip install pyg-lib -f https://data.pyg.org/whl/torch-2.1.0+${CUDA}.html

where ${CUDA} should be replaced by either cpu, cu118 or cu121

The following combinations are supported:

PyTorch 2.1	cpu	cu118	cu121
Linux	✅	✅	✅
macOS	✅

Older PyTorch versions like PyTorch 1.12, 1.13 and 2.0.0 are still supported, and can be installed as described in our README.md. PyTorch 1.11 support has been dropped.

METIS partioning

pyg-lib==0.3.0 enables METIS partioning by introducing pyg_lib.partition (#229).

from pyg_lib.partition import metis

cluster = metis(rowptr, col, num_partitions)

Neighbor Sampling Improvements

pyg-lib==0.3.0 brings various improvements to our neighbor sampling routine:

• Support for biased/weighted sampling: pyg_lib.sampler.neighbor_sample and pyg_lib.sampler.hetero_neighbor_sample now support the additional edge_weight argument (#247, #251)
• pyg_lib.sampler.hetero_neighbor_sample now performs neighborhood sampling across edge types in parallel (#211)
• Added low-level support for distributed neighborhood sampling (#246, #252, #253, #254)

Additional Features

• Added dispatch for XPU device in index_sort (#243)
• Updated cutlass version for speed boosts in segment_matmul and grouped_matmul (#235)

Bugfixes

• Fixed vector-based mapping issue in Mapping (#244)
• Fixed performance issues reported by Coverity Tool (#240)
• Fixed TorchScript support in grouped_matmul (#220)

New Contributors

• @yaox12 made their first contribution in #213
• @yanbing-j made their first contribution in #231
• @akihironitta made their first contribution in #248

Full Changelog: 0.2.0...0.3.0