cugraph-pyg API Reference

cugraph-pyg

Graph Storage

cugraph_pyg.data.dask_graph_store.DaskGraphStore(F, ...)	Duck-typed version of PyG's GraphStore and FeatureStore that uses Dask to distribute the graph structure across GPUs and a cugraph.gnn.FeatureStore to store node/edge features.
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.dask_node_loader.DaskNeighborLoader(data)	Duck-typed version of the PyG NeighborLoader interface that uses Dask to sample nodes using the uniform neighbor sampling algorithm.
cugraph_pyg.loader.dask_node_loader.BulkSampleLoader(...)	Iterator that executes sampling using Dask and cuGraph and loads sampled minibatches from disk.
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(...)	Methods
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.