def inv(A: SparseMatrix) -> SparseMatrix:
    """Returns the inverse of the sparse matrix.

    This function only supports square diagonal matrices with scalar nonzero
    values.

    Returns
    -------
    SparseMatrix
        Inverse of the sparse matrix

    Examples
    --------

    >>> val = torch.arange(1, 4).float()
    >>> D = dglsp.diag(val)
    >>> D.inv()
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1., 2., 3.]),
                 shape=(3, 3), nnz=3)
    """
    num_rows, num_cols = A.shape
    assert A.is_diag(), "Non-diagonal sparse matrix does not support inversion."
    assert num_rows == num_cols, f"Expect a square matrix, got shape {A.shape}"
    assert len(A.val.shape) == 1, "inv only supports 1D nonzero val"

    return diag(1.0 / A.val, A.shape)
Returns the inverse of the sparse matrix.

    This function only supports square diagonal matrices with scalar nonzero
    values.

    Returns
    -------
    SparseMatrix
        Inverse of the sparse matrix

    Examples
    --------

    >>> val = torch.arange(1, 4).float()
    >>> D = dglsp.diag(val)
    >>> D.inv()
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1., 2., 3.]),
                 shape=(3, 3), nnz=3)
    
python/dgl/sparse/unary_op.py