add __eq__

test/test_tensor.py

@@ -37,3 +37,19 @@ def test_to_symmetric(device):
         [6, 0, 5],
         [3, 5, 0],
     ]
+
+
+def test_equal():
+    row = torch.tensor([0, 0, 0, 1, 1])
+    col = torch.tensor([0, 1, 2, 0, 2])
+    value = torch.arange(1, 6)
+    matA = SparseTensor(row=row, col=col, value=value)
+    matB = SparseTensor(row=row, col=col, value=value)
+    col = torch.tensor([0, 1, 2, 0, 1])
+    matC = SparseTensor(row=row, col=col, value=value)
+
+    assert id(matA) != id(matB)
+    assert matA == matB
+
+    assert id(matA) != id(matC)
+    assert matA != matC

torch_sparse/tensor.py

@@ -197,6 +197,28 @@ class SparseTensor(object):
         self.storage.clear_cache_()
         return self
 
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+
+        if self.sizes() != other.sizes():
+            return False
+
+        rowptrA, colA, valueA = self.csr()
+        rowptrB, colB, valueB = other.csr()
+
+        if valueA is None and valueB is not None:
+            return False
+        if valueA is not None and valueB is None:
+            return False
+        if not torch.equal(rowptrA, rowptrB):
+            return False
+        if not torch.equal(colA, colB):
+            return False
+        if valueA is None and valueB is None:
+            return True
+        return torch.equal(valueA, valueB)
+
     # Utility functions #######################################################
 
     def fill_value_(self, fill_value: float, dtype: Optional[int] = None):