gigl.distributed.utils.degree#

Utility functions for computing node degrees in distributed graph settings.

This module provides functions to compute node out-degrees from graph partitions and aggregate them across distributed machines. Degrees are computed from the CSR (Compressed Sparse Row) topology stored in GraphLearn-Torch Graph objects.

For homogeneous graphs, callers receive a single torch.Tensor. For heterogeneous graphs, degrees are accumulated per anchor node type (summing across all edge types incident to that node type) before the distributed all-reduce, so callers receive dict[NodeType, torch.Tensor].

Requirements#

torch.distributed must be initialized before calling these functions.

Usage#

Access dataset.degree_tensor to lazily compute and cache the degree tensor.

Over-counting correction is handled automatically in _all_reduce_degrees by detecting how many processes share the same machine (and thus the same data).

Heterogeneous partitioned graphs are expected to materialize all registered non-label edge types on every rank, even when a rank has no local edges for a type. This keeps the per-node-type all-reduce order consistent across ranks.

Degree tensors are stored as int32 to stay aligned with the PPR C++ sampler’s total-degree tensor requirement while keeping memory lower than int64. We avoid int16 because it has caused compatibility issues in this path: during the C++ PPR sampler migration, torch.distributed.all_reduce on an int16 tensor produced RuntimeError: Invalid scalar type. Values above the int32 maximum are clamped before casting to avoid wraparound.

Attributes#

logger

Functions#

compute_and_broadcast_degree_tensor(graph, edge_dir)

Compute node degrees from a graph and aggregate across all machines.

Module Contents#

gigl.distributed.utils.degree.compute_and_broadcast_degree_tensor(graph, edge_dir)[source]#

Compute node degrees from a graph and aggregate across all machines.

For each non-label edge type, degrees are derived from the CSR row pointers (indptr). For heterogeneous graphs, degrees are summed across all edge types incident to each anchor node type locally before the all-reduce, so the per-edge-type tensor is only a transient intermediate and is never stored, returned, or transmitted over RPC.

Over-counting correction (for processes sharing the same data) is handled automatically by detecting the distributed topology.

Parameters:

  • graph (Union[graphlearn_torch.data.Graph, dict[torch_geometric.typing.EdgeType, graphlearn_torch.data.Graph]]) – A Graph (homogeneous) or dict[EdgeType, Graph] (heterogeneous). For heterogeneous graphs, label edge types are automatically excluded — they are supervision edges and should not contribute to node degree for graph traversal algorithms like PPR.

  • edge_dir (str) – Sampling direction — "in" or "out". Determines which end of each edge is the anchor node type for degree accumulation.

Returns:

Aggregated degree

tensors. For homogeneous graphs, returns an int32 tensor of shape [num_nodes]. For heterogeneous graphs, returns int32 tensors keyed by node type with shape [num_nodes_of_that_type].

Return type:

Union[torch.Tensor, dict[NodeType, torch.Tensor]]

Raises:

  • RuntimeError – If torch.distributed is not initialized.

  • ValueError – If topology is unavailable.

gigl.distributed.utils.degree.logger[source]#