cugraph_pyg.sampler.sampler.HeterogeneousSampleReader

class cugraph_pyg.sampler.sampler.HeterogeneousSampleReader(base_reader: Iterator[Tuple[Dict[str, Tensor], int, int]], src_types: Tensor, dst_types: Tensor, vertex_offsets: Tensor, edge_types: List[Tuple[str, str, str]], vertex_types: List[str])

Subclass of SampleReader that reads heterogeneous output samples produced by the cuGraph distributed sampler.

__init__(base_reader: Iterator[Tuple[Dict[str, Tensor], int, int]], src_types: Tensor, dst_types: Tensor, vertex_offsets: Tensor, edge_types: List[Tuple[str, str, str]], vertex_types: List[str])

Constructs a new HeterogeneousSampleReader

Parameters:

base_reader: Iterator[Tuple[Dict[str, “torch.Tensor”], int, int]]

The iterator responsible for loading saved samples produced by the cuGraph distributed sampler.

src_types: torch.Tensor

Integer source type for each integer edge type.

dst_types: torch.Tensor

Integer destination type for each integer edge type.

vertex_offsets: torch.Tensor

Vertex offsets for each vertex type. Used to de-offset vertices outputted by the cuGraph sampler and return PyG-compliant vertex IDs.

edge_types: List[Tuple[str, str, str]]

List of edge types in the graph in order, so they can be mapped to numeric edge types.

vertex_types: List[str]

List of vertex types, in order so they can be mapped to numeric vertex types.

Methods

__init__(base_reader, src_types, dst_types, ...)

Constructs a new HeterogeneousSampleReader