Function from_torch_sparse

python/dgl/sparse/sparse_matrix.py:1284–1339

Creates a sparse matrix from a torch sparse tensor, which can have coo, csr, or csc layout. Parameters ---------- torch_sparse_tensor : torch.Tensor Torch sparse tensor Returns ------- SparseMatrix Sparse matrix Examples -------- >>> indice

(torch_sparse_tensor: torch.Tensor)

def from_torch_sparse(torch_sparse_tensor: torch.Tensor) -> SparseMatrix:
    """Creates a sparse matrix from a torch sparse tensor, which can have coo,
    csr, or csc layout.

    Parameters
    ----------
    torch_sparse_tensor : torch.Tensor
        Torch sparse tensor

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------

    >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]])
    >>> val = torch.ones(3)
    >>> torch_coo = torch.sparse_coo_tensor(indices, val)
    >>> dglsp.from_torch_sparse(torch_coo)
    SparseMatrix(indices=tensor([[1, 1, 2],
                                 [2, 4, 3]]),
                 values=tensor([1., 1., 1.]),
                 shape=(3, 5), nnz=3)
    """
    assert torch_sparse_tensor.layout in (
        torch.sparse_coo,
        torch.sparse_csr,
        torch.sparse_csc,
    ), (
        f"Cannot convert Pytorch sparse tensor with layout "
        f"{torch_sparse_tensor.layout} to DGL sparse."
    )
    if torch_sparse_tensor.layout == torch.sparse_coo:
        # Use ._indices() and ._values() to access uncoalesced indices and
        # values.
        return spmatrix(
            torch_sparse_tensor._indices(),
            torch_sparse_tensor._values(),
            torch_sparse_tensor.shape[:2],
        )
    elif torch_sparse_tensor.layout == torch.sparse_csr:
        return from_csr(
            torch_sparse_tensor.crow_indices(),
            torch_sparse_tensor.col_indices(),
            torch_sparse_tensor.values(),
            torch_sparse_tensor.shape[:2],
        )
    else:
        return from_csc(
            torch_sparse_tensor.ccol_indices(),
            torch_sparse_tensor.row_indices(),
            torch_sparse_tensor.values(),
            torch_sparse_tensor.shape[:2],
        )