chore: import upstream snapshot with attribution

python/dgl/sparse/unary_op.py

@@ -0,0 +1,61 @@
+"""DGL unary operators for sparse matrix module."""
+from .sparse_matrix import diag, SparseMatrix, val_like
+
+
+def neg(A: SparseMatrix) -> SparseMatrix:
+    """Returns a new sparse matrix with the negation of the original nonzero
+    values, equivalent to ``-A``.
+
+    Returns
+    -------
+    SparseMatrix
+        Negation of the sparse matrix
+
+    Examples
+    --------
+
+    >>> indices = torch.tensor([[1, 1, 3], [1, 2, 3]])
+    >>> val = torch.tensor([1., 1., 2.])
+    >>> A = dglsp.spmatrix(indices, val)
+    >>> A = -A
+    SparseMatrix(indices=tensor([[1, 1, 3],
+                                 [1, 2, 3]]),
+                 values=tensor([-1., -1., -2.]),
+                 shape=(4, 4), nnz=3)
+    """
+    return val_like(A, -A.val)
+
+
+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)
+
+
+SparseMatrix.neg = neg
+SparseMatrix.__neg__ = neg
+SparseMatrix.inv = inv