rllm.data.EdgeStorage

class rllm.data.EdgeStorage(initialdata: Dict[str, Any] | None = None, **kwargs)[source]

Bases: BaseStorage

A storage class for edge attributes in a graph.

Parameters:

  • initialdata (Optional[Dict[str, Any]]) – Initial data to populate the storage.

  • **kwargs – Additional keyword arguments to populate the storage.

num_edges

The number of edges in the storage.

Type:

int

is_edge_attr(key: str) → bool[source]

Edge attributes should be: 1. List, tuple, or TableData with length equal to the number of edges. 2. Tensor with the first dimension equal to the number of edges. 3. Numpy array with the first dimension equal to the number of edges.

to_csc(device: device | None = None, num_nodes: int | None = None, share_memory: bool = False, is_sorted: bool = False, src_node_time: Tensor | None = None, edge_time: str | Tensor | None = None) → Tuple[Tensor, Tensor, Tensor | None][source]

Convert the edge storage to a CSC format. This can be applied to both homogeneous and heterogeneous graphs.

Parameters:

  • device (torch.device, optional) – The desired device of the returned tensors. If None, use the current device. (default: None)

  • num_nodes (int, optional) – The number of nodes. If None, infer from destination node index of edge_index. (default: None)

  • share_memory (bool, optional) – If set to True, will share memory among returned tensors. This can accelerate process when using multiple processes. (default: False)

  • is_sorted (bool, optional) – If set to True, will not sort the edge index. (default: False)

  • src_node_time (Tensor, optional) – The source node time. If not None, will sort the edge index by src_node_time. (default: None)

  • edge_time (Union[str, Tensor], optional) – The edge time attribute. If not None, will sort the edge index by edge_time_attr. (default: None)

Returns:

The column indices,

row indices, and the permutation index.

Return type:

Tuple[Tensor, Tensor, Optional[Tensor]]