cugraph-pyg API Reference

cugraph-pyg

Graph Storage

cugraph_pyg.data.graph_store.GraphStore([...])

cuGraph-backed PyG GraphStore implementation that distributes the graph across workers.

Feature Storage

cugraph_pyg.data.feature_store.TensorDictFeatureStore()	A basic implementation of the PyG FeatureStore interface that stores feature data in a single TensorDict.
cugraph_pyg.data.feature_store.WholeFeatureStore([...])	A basic implementation of the PyG FeatureStore interface that stores feature data in WholeGraph WholeMemory.

Data Loaders

cugraph_pyg.loader.node_loader.NodeLoader(...)	Duck-typed version of torch_geometric.loader.NodeLoader.
cugraph_pyg.loader.neighbor_loader.NeighborLoader(...)	Duck-typed version of torch_geometric.loader.NeighborLoader

Samplers

cugraph_pyg.sampler.sampler.BaseSampler(...)
cugraph_pyg.sampler.sampler.SampleReader(...)	Iterator that processes results from the cuGraph distributed sampler.
cugraph_pyg.sampler.sampler.HomogeneousSampleReader(...)	Subclass of SampleReader that reads homogeneous output samples produced by the cuGraph distributed sampler.
cugraph_pyg.sampler.sampler.SampleIterator(...)	Iterator that combines output graphs with their features to produce final output minibatches that can be fed into a GNN model.