Unverified Commit 97fbd94d authored by Mufei Li's avatar Mufei Li Committed by GitHub
Browse files

[Sparse] neg and inv (#4980)



* Update

* update
Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-36-188.ap-northeast-1.compute.internal>
parent edb14f93
"""dgl sparse class.""" """dgl sparse class."""
import sys
import os import os
import sys
import torch import torch
from .._ffi import libinfo from .._ffi import libinfo
from .sparse_matrix import *
from .diag_matrix import * from .diag_matrix import *
from .elementwise_op import * from .elementwise_op import *
from .sparse_matrix import *
from .unary_op_diag import *
from .unary_op_sp import *
def load_dgl_sparse(): def load_dgl_sparse():
......
"""DGL unary operators for diagonal matrix module."""
# pylint: disable= invalid-name
from .diag_matrix import DiagMatrix, diag
def neg(D: DiagMatrix) -> DiagMatrix:
"""Return a new diagonal matrix with the negation of the original nonzero
values.
Returns
-------
DiagMatrix
Negation of the diagonal matrix
Examples
--------
>>> val = torch.arange(3).float()
>>> mat = diag(val)
>>> mat = -mat
>>> print(mat)
DiagMatrix(val=tensor([-0., -1., -2.]),
shape=(3, 3))
"""
return diag(-D.val, D.shape)
def inv(D: DiagMatrix) -> DiagMatrix:
"""Return the inverse of the diagonal matrix.
This function only supports square matrices with scalar nonzero values.
Returns
-------
DiagMatrix
Inverse of the diagonal matrix
Examples
--------
>>> val = torch.arange(1, 4).float()
>>> mat = diag(val)
>>> mat = mat.inv()
>>> print(mat)
DiagMatrix(val=tensor([1.0000, 0.5000, 0.3333]),
shape=(3, 3))
"""
num_rows, num_cols = D.shape
assert num_rows == num_cols, f"Expect a square matrix, got shape {D.shape}"
assert len(D.val.shape) == 1, "inv only supports 1D nonzero val"
return diag(1.0 / D.val, D.shape)
DiagMatrix.neg = neg
DiagMatrix.__neg__ = neg
DiagMatrix.inv = inv
"""DGL unary operators for sparse matrix module."""
from .sparse_matrix import SparseMatrix, val_like
def neg(A: SparseMatrix) -> SparseMatrix:
"""Return a new sparse matrix with the negation of the original nonzero
values.
Returns
-------
SparseMatrix
Negation of the sparse matrix
Examples
--------
>>> row = torch.tensor([1, 1, 3])
>>> col = torch.tensor([1, 2, 3])
>>> val = torch.tensor([1., 1., 2.])
>>> A = create_from_coo(row, col, val)
>>> A = -A
>>> print(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)
SparseMatrix.neg = neg
SparseMatrix.__neg__ = neg
import pytest
import torch
import sys
import backend as F
from dgl.mock_sparse2 import diag
def test_neg():
ctx = F.ctx()
val = torch.arange(3).float().to(ctx)
D = diag(val)
neg_D = -D
assert D.shape == neg_D.shape
assert torch.allclose(-D.val, neg_D.val)
assert D.val.device == neg_D.val.device
def test_inv():
ctx = F.ctx()
val = torch.arange(1, 4).float().to(ctx)
D = diag(val)
inv_D = D.inv()
assert D.shape == inv_D.shape
assert torch.allclose(1.0 / D.val, inv_D.val)
assert D.val.device == inv_D.val.device
import pytest
import torch
import sys
import backend as F
from dgl.mock_sparse2 import create_from_coo
# FIXME(issue #4818): Skipping tests on win.
if not sys.platform.startswith("linux"):
pytest.skip("skipping tests on win", allow_module_level=True)
def test_neg():
ctx = F.ctx()
row = torch.tensor([1, 1, 3]).to(ctx)
col = torch.tensor([1, 2, 3]).to(ctx)
val = torch.tensor([1.0, 1.0, 2.0]).to(ctx)
A = create_from_coo(row, col, val)
neg_A = -A
assert A.shape == neg_A.shape
assert A.nnz == neg_A.nnz
assert torch.allclose(-A.val, neg_A.val)
assert torch.allclose(torch.stack(A.coo()), torch.stack(neg_A.coo()))
assert A.val.device == neg_A.val.device
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment