Source code for torch_geometric.utils.sparse

import torch


[docs]def dense_to_sparse(adj): r"""Converts a dense adjacency matrix to a sparse adjacency matrix defined by edge indices and edge attributes. Args: adj (Tensor): The dense adjacency matrix. :rtype: (:class:`LongTensor`, :class:`Tensor`) Examples: >>> # Forr a single adjacency matrix >>> adj = torch.tensor([[3, 1], ... [2, 0]]) >>> dense_to_sparse(adj) (tensor([[0, 0, 1], [0, 1, 0]]), tensor([3, 1, 2])) >>> # For two adjacency matrixes >>> adj = torch.tensor([[[3, 1], ... [2, 0]], ... [[0, 1], ... [0, 2]]]) >>> dense_to_sparse(adj) (tensor([[0, 0, 1, 2, 3], [0, 1, 0, 3, 3]]), tensor([3, 1, 2, 1, 2])) """ assert adj.dim() >= 2 and adj.dim() <= 3 assert adj.size(-1) == adj.size(-2) index = adj.nonzero(as_tuple=True) edge_attr = adj[index] if len(index) == 3: batch = index[0] * adj.size(-1) index = (batch + index[1], batch + index[2]) return torch.stack(index, dim=0), edge_attr