rllm.transforms.graph_transforms.EdgeTransform

class rllm.transforms.graph_transforms.EdgeTransform

Bases: ABC

Base class for edge-wise transformations on graph data.

The transform is applied to adj for homogeneous graphs, each valid store.adj for heterogeneous graphs, or directly to a tensor input.

Shape:

• GraphData: data.adj can be dense or sparse and should follow the adjacency format expected by subclasses.

• HeteroGraphData: store.adj can be dense or sparse and should follow the adjacency format expected by subclasses.

• torch.Tensor: if dense, input must be a square matrix with shape [N, N]; sparse tensors are forwarded as-is.

Examples:

class KeepSelfLoops(EdgeTransform):
    def forward(self, adj):
        return adj

transform = KeepSelfLoops()
out = transform(data)