Repository: cai4cai/torchsparsegradutils (language: python)

Issue:
Add convenience functions to map sparse tensor to/from cupy and jax
Having conversion functions going across pytorch and cupy/jax for sparse tensors would be useful.

Basic example supporting torch->cupy COO tensors:
```python
# Convenience function to map a torch COO tensor in a cupy one
def coo_torch2cupy(A):
  A = A.data.coalesce() # Maybe this should raise if the input is not coalesced as it may mess up with gradients otherwise
  Avals_cp = cp.asarray(A.values())
  Aidx_cp = cp.asarray(A.indices())
  return cp.sparse.coo_matrix((Avals_cp, Aidx_cp))
```

This relates to https://github.com/cai4cai/sparsitypreservingtorchutils/issues/5


Discussion / hints:
Seems like jax only has built-in support for Batched-COO:
https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html

For Jax, a conversion of the dense tensor might be enough in the first instance. See https://github.com/lucidrains/jax2torch/blob/main/jax2torch/jax2torch.py
```python
import torch
from torch.utils import dlpack as torch_dlpack

import jax
from jax import dlpack as jax_dlpack

def j2t(x_jax):
    x_torch = torch_dlpack.from_dlpack(jax_dlpack.to_dlpack(x_jax))
    return x_torch

def t2j(x_torch):
    x_torch = x_torch.contiguous() # https://github.com/google/jax/issues/8082
    x_jax = jax_dlpack.from_dlpack(torch_dlpack.to_dlpack(x_torch))
    return x_jax
```

Fix patch:
diff --git a/setup.py b/setup.py
index 21b0d8c..2a04615 100644
--- a/setup.py
+++ b/setup.py
@@ -22,12 +22,12 @@ def readme():
     author="CAI4CAI research group",
     author_email="contact@cai4cai.uk",
     license="Apache-2.0",
-    packages=["torchsparsegradutils"],
+    packages=setuptools.find_packages(),
     install_requires=[
         "torch",
     ],
     extras_require={
-        "extras": ["jax"],
+        "extras": ["jax", "cupy"],
     },
     zip_safe=False,
 )
diff --git a/torchsparsegradutils/cupy/__init__.py b/torchsparsegradutils/cupy/__init__.py
new file mode 100644
index 0000000..4cfeddb
--- /dev/null
+++ b/torchsparsegradutils/cupy/__init__.py
@@ -0,0 +1,16 @@
+import importlib
+
+cupy_spec = importlib.util.find_spec("cupy")
+if cupy_spec is None:
+    have_cupy = False
+    import warnings
+
+    warnings.warn(
+        "\n\nAttempting to import an optional module in torchsparsegradutils that depends on cupy but cupy couldn't be imported -> falling back to numpy.\n"
+    )
+else:
+    have_cupy = True
+
+from .cupy_bindings import c2t_coo, t2c_coo, c2t_csr, t2c_csr
+
+__all__ = ["c2t_coo", "t2c_coo", "c2t_csr", "t2c_csr"]
diff --git a/torchsparsegradutils/cupy/cupy_bindings.py b/torchsparsegradutils/cupy/cupy_bindings.py
new file mode 100644
index 0000000..1293f79
--- /dev/null
+++ b/torchsparsegradutils/cupy/cupy_bindings.py
@@ -0,0 +1,75 @@
+import torchsparsegradutils.cupy as tsgucupy
+
+if tsgucupy.have_cupy:
+    import cupy as cp
+    import cupyx.scipy.sparse as csp
+
+import numpy as np
+import scipy.sparse as nsp
+
+import torch
+
+import warnings
+
+
+def _get_array_modules(x):
+    # Get cupy or numpy based on context
+    if tsgucupy.have_cupy:
+        if torch.is_tensor(x):
+            if x.device == torch.device("cpu"):
+                xp = np
+                xsp = nsp
+            else:
+                xp = cp
+                xsp = csp
+        else:
+            xp = cp.get_array_module(x_cupy)
+            xsp = cupyx.scipy.get_array_module(x_cupy)
+    else:
+        xp = np
+        xsp = nsp
+    return xp, xsp
+
+
+def t2c_csr(x_torch):
+    # Convert a torch CSR tensor to a cupy/numpy CSR array
+    xp, xsp = _get_array_modules(x_torch)
+
+    data_c = xp.asarray(x_torch.values())
+    col_idx_c = xp.asarray(x_torch.col_indices())
+    ind_ptr_c = xp.asarray(x_torch.crow_indices())
+    x_cupy = xsp.csr_matrix((data_c, col_idx_c, ind_ptr_c))
+    return x_cupy
+
+
+def c2t_csr(x_cupy):
+    # Convert a cupy/numpy CSR array to a torch CSR tensor
+    data_t = torch.as_tensor(x_cupy.data)
+    idices_t = torch.as_tensor(x_cupy.indices)
+    ind_ptr_t = torch.as_tensor(x_cupy.indptr)
+    x_torch = torch.sparse_csr_tensor(ind_ptr_t, idices_t, data_t, x_cupy.shape)
+    return x_torch
+
+
+def t2c_coo(x_torch):
+    # Convert a torch COO tensor to a cupy/numpy COO array
+    xp, xsp = _get_array_modules(x_torch)
+
+    if not x_torch.is_coalesced():
+        warnings.warn(
+            "Requested a conversion from torch to cupy/numpy on a non-coalesced tensor -> coalescing implicitly"
+        )
+        x_torch = x_torch.coalesce()
+    data_c = xp.asarray(x_torch.values())
+    idx_cp = xp.asarray(x_torch.indices())
+    x_cupy = xsp.coo_matrix((data_c, idx_cp))
+    return x_cupy
+
+
+def c2t_coo(x_cupy):
+    # Convert a cupy COO array to a torch COO tensor
+    data_t = torch.as_tensor(x_cupy.data)
+    row_t = torch.as_tensor(x_cupy.row)
+    col_t = torch.as_tensor(x_cupy.col)
+    x_torch = torch.sparse_coo_tensor(torch.stack([row_t, col_t], dim=0), data_t, x_cupy.shape)
+    return x_torch


Test patch:
diff --git a/tests/test_cupy_bindings.py b/tests/test_cupy_bindings.py
new file mode 100644
index 0000000..888db9e
--- /dev/null
+++ b/tests/test_cupy_bindings.py
@@ -0,0 +1,95 @@
+import torch
+import unittest
+import torchsparsegradutils as tsgu
+import torchsparsegradutils.cupy as tsgucupy
+
+import warnings
+
+if tsgucupy.have_cupy:
+    import cupy as cp
+    import cupyx.scipy.sparse as csp
+else:
+    warnings.warn(
+        "Importing optional cupy-related module failed to find cupy -> cupy-related tests running as numpy only."
+    )
+
+import numpy as np
+import scipy.sparse as nsp
+
+def _c2n(x_cupy):
+    if tsgucupy.have_cupy:
+        return cp.asnumpy(x_cupy)
+    else:
+        return np.asarray(x_cupy)
+        
+
+class C2TIOTest(unittest.TestCase):
+    """IO conversion tests between torch and cupy"""
+
+    def setUp(self) -> None:
+        # The device can be specialised by a daughter class
+        if not hasattr(self, "device"):
+            self.device = torch.device("cpu")
+            self.xp = np
+            self.xsp = nsp
+
+        self.x_shape = (4, 2)
+        self.x_t = torch.randn(self.x_shape, dtype=torch.float64, device=self.device)
+        rng = np.random.default_rng()
+        x_n = rng.standard_normal(self.x_shape, dtype=np.float64)
+
+        self.x_c = self.xp.asarray(x_n)
+        self.A_shape = (4, 4)
+        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device).relu().to_sparse_csr()
+
+    def test_t2c_coo(self):
+        x_t_coo = self.x_t.to_sparse_coo()
+        x_c_coo = tsgucupy.t2c_coo(x_t_coo)
+        self.assertTrue(x_c_coo.shape == x_t_coo.shape)
+        self.assertTrue(np.isclose(_c2n(x_c_coo.todense()), x_t_coo.to_dense().cpu().numpy()).all())
+        x_t2c2t_coo = tsgucupy.c2t_coo(x_c_coo)
+        self.assertTrue(x_t2c2t_coo.shape == x_t_coo.shape)
+        self.assertTrue(np.isclose(x_t2c2t_coo.to_dense().cpu().numpy(), x_t_coo.to_dense().cpu().numpy()).all())
+
+    def test_c2t_coo(self):
+        x_c_coo = self.xsp.coo_matrix(self.x_c)
+        x_t_coo = tsgucupy.c2t_coo(x_c_coo)
+        self.assertTrue(x_c_coo.shape == x_t_coo.shape)
+        self.assertTrue(np.isclose(_c2n(x_c_coo.todense()), x_t_coo.to_dense().cpu().numpy()).all())
+        x_c2t2c_coo = tsgucupy.t2c_coo(x_t_coo)
+        self.assertTrue(x_c2t2c_coo.shape == x_c_coo.shape)
+        self.assertTrue(np.isclose(_c2n(x_c2t2c_coo.todense()), _c2n(x_c_coo.todense())).all())
+
+    def test_t2c_csr(self):
+        x_t_csr = self.x_t.to_sparse_csr()
+        x_c_csr = tsgucupy.t2c_csr(x_t_csr)
+        self.assertTrue(x_c_csr.shape == x_t_csr.shape)
+        self.assertTrue(np.isclose(_c2n(x_c_csr.todense()), x_t_csr.to_dense().cpu().numpy()).all())
+        x_t2c2t_csr = tsgucupy.c2t_csr(x_c_csr)
+        self.assertTrue(x_t2c2t_csr.shape == x_t_csr.shape)
+        self.assertTrue(np.isclose(x_t2c2t_csr.to_dense().cpu().numpy(), x_t_csr.to_dense().cpu().numpy()).all())
+
+    def test_c2t_csr(self):
+        x_c_csr = self.xsp.csr_matrix(self.x_c)
+        x_t_csr = tsgucupy.c2t_csr(x_c_csr)
+        self.assertTrue(x_c_csr.shape == x_t_csr.shape)
+        self.assertTrue(np.isclose(_c2n(x_c_csr.todense()), x_t_csr.to_dense().cpu().numpy()).all())
+        x_c2t2c_csr = tsgucupy.t2c_csr(x_t_csr)
+        self.assertTrue(x_c2t2c_csr.shape == x_c_csr.shape)
+        self.assertTrue(np.isclose(_c2n(x_c2t2c_csr.todense()), _c2n(x_c_csr.todense())).all())
+
+
+class C2TIOTestCUDA(C2TIOTest):
+    """Override superclass setUp to run on GPU"""
+
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+        self.device = torch.device("cuda")
+        self.xp = cp
+        self.xsp = csp
+        super().setUp()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_jax_bindings.py b/tests/test_jax_bindings.py
index 614cde9..30f2ba1 100644
--- a/tests/test_jax_bindings.py
+++ b/tests/test_jax_bindings.py
@@ -33,15 +33,15 @@ def setUp(self) -> None:
     def test_t2j(self):
         x_j = tsgujax.t2j(self.x_t)
         self.assertTrue(x_j.shape == self.x_t.shape)
-        self.assertTrue(np.isclose(np.asarray(x_j), self.x_t.numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(x_j), self.x_t.cpu().numpy()).all())
         x_t2j2t = tsgujax.j2t(x_j)
         self.assertTrue(x_t2j2t.shape == self.x_t.shape)
-        self.assertTrue(np.isclose(x_t2j2t.numpy(), self.x_t.numpy()).all())
+        self.assertTrue(np.isclose(x_t2j2t.cpu().numpy(), self.x_t.cpu().numpy()).all())
 
     def test_j2t(self):
         x_t = tsgujax.j2t(self.x_j)
         self.assertTrue(x_t.shape == self.x_j.shape)
-        self.assertTrue(np.isclose(np.asarray(self.x_j), x_t.numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(self.x_j), x_t.cpu().numpy()).all())
         x_j2t2j = tsgujax.t2j(x_t)
         self.assertTrue(x_j2t2j.shape == self.x_j.shape)
         self.assertTrue(np.isclose(np.asarray(x_j2t2j), np.asarray(self.x_j)).all())
@@ -56,22 +56,22 @@ def test_spmm_t4j(self):
         x_t = tsgujax.j2t(self.x_j)
         Ax_t = A_t @ x_t
         self.assertTrue(Ax_j.shape == Ax_t.shape)
-        self.assertTrue(np.isclose(np.asarray(Ax_j), Ax_t.numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(Ax_j), Ax_t.cpu().numpy()).all())
 
     def test_t2j_coo(self):
         x_t_coo = self.x_t.to_sparse_coo()
         x_j_coo = tsgujax.t2j_coo(x_t_coo)
         self.assertTrue(x_j_coo.shape == x_t_coo.shape)
-        self.assertTrue(np.isclose(np.asarray(x_j_coo.todense()), x_t_coo.to_dense().numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(x_j_coo.todense()), x_t_coo.to_dense().cpu().numpy()).all())
         x_t2j2t_coo = tsgujax.j2t_coo(x_j_coo)
         self.assertTrue(x_t2j2t_coo.shape == x_t_coo.shape)
-        self.assertTrue(np.isclose(x_t2j2t_coo.to_dense().numpy(), x_t_coo.to_dense().numpy()).all())
+        self.assertTrue(np.isclose(x_t2j2t_coo.to_dense().cpu().numpy(), x_t_coo.to_dense().cpu().numpy()).all())
 
     def test_j2t_coo(self):
         x_j_coo = jax.experimental.sparse.COO.fromdense(self.x_j)
         x_t_coo = tsgujax.j2t_coo(x_j_coo)
         self.assertTrue(x_j_coo.shape == x_t_coo.shape)
-        self.assertTrue(np.isclose(np.asarray(x_j_coo.todense()), x_t_coo.to_dense().numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(x_j_coo.todense()), x_t_coo.to_dense().cpu().numpy()).all())
         x_j2t2j_coo = tsgujax.t2j_coo(x_t_coo)
         self.assertTrue(x_j2t2j_coo.shape == x_j_coo.shape)
         self.assertTrue(np.isclose(np.asarray(x_j2t2j_coo.todense()), np.asarray(x_j_coo.todense())).all())
@@ -80,16 +80,16 @@ def test_t2j_csr(self):
         x_t_csr = self.x_t.to_sparse_csr()
         x_j_csr = tsgujax.t2j_csr(x_t_csr)
         self.assertTrue(x_j_csr.shape == x_t_csr.shape)
-        self.assertTrue(np.isclose(np.asarray(x_j_csr.todense()), x_t_csr.to_dense().numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(x_j_csr.todense()), x_t_csr.to_dense().cpu().numpy()).all())
         x_t2j2t_csr = tsgujax.j2t_csr(x_j_csr)
         self.assertTrue(x_t2j2t_csr.shape == x_t_csr.shape)
-        self.assertTrue(np.isclose(x_t2j2t_csr.to_dense().numpy(), x_t_csr.to_dense().numpy()).all())
+        self.assertTrue(np.isclose(x_t2j2t_csr.to_dense().cpu().numpy(), x_t_csr.to_dense().cpu().numpy()).all())
 
     def test_j2t_csr(self):
         x_j_csr = jax.experimental.sparse.CSR.fromdense(self.x_j)
         x_t_csr = tsgujax.j2t_csr(x_j_csr)
         self.assertTrue(x_j_csr.shape == x_t_csr.shape)
-        self.assertTrue(np.isclose(np.asarray(x_j_csr.todense()), x_t_csr.to_dense().numpy()).all())
+        self.assertTrue(np.isclose(np.asarray(x_j_csr.todense()), x_t_csr.to_dense().cpu().numpy()).all())
         x_j2t2j_csr = tsgujax.t2j_csr(x_t_csr)
         self.assertTrue(x_j2t2j_csr.shape == x_j_csr.shape)
         self.assertTrue(np.isclose(np.asarray(x_j2t2j_csr.todense()), np.asarray(x_j_csr.todense())).all())