feat: enable jax in todense

Author: mrava87

pylops/jaxoperator.py

@@ -24,6 +24,22 @@
 
 
 class JaxOperator(LinearOperator):
+    """Enable JAX backend for PyLops operator.
+
+    This class can be used to wrap a pylops operator to enable the JAX
+    backend. Doing so, users can run all of the methods of a pylops
+    operator with JAX arrays. Moreover, the forward and adjoint
+    are internally just-in-time compiled, and other JAX functionalities
+    such as automatic differentiation and automatic vectorization
+    are enabled.
+
+    Parameters
+    ----------
+    Op : :obj:`pylops.LinearOperator`
+        PyLops operator
+
+    """
+
     def __init__(self, Op: LinearOperator) -> None:
         if not deps.jax_enabled:
             raise NotImplementedError(jax_message)
@@ -43,7 +59,12 @@ def __call__(self, x, *args, **kwargs):
         return self._matvec(x)
 
     def _rmatvecad(self, x: JaxType, y: JaxType) -> JaxType:
-        """Vector-Jacobian products
+        _, f_vjp = jax.vjp(self._matvec, x)
+        xadj = jax.jit(f_vjp)(y)[0]
+        return xadj
+
+    def rmatvecad(self, x: JaxType, y: JaxType) -> JaxType:
+        """Vector-Jacobian product
 
         JIT-compiled Vector-Jacobian product
 
@@ -59,33 +80,12 @@ def _rmatvecad(self, x: JaxType, y: JaxType) -> JaxType:
         xadj : :obj:`jaxlib.xla_extension.ArrayImpl`
             Output array
 
-        """
-        _, f_vjp = jax.vjp(self._matvec, x)
-        xadj = jax.jit(f_vjp)(y)[0]
-        return xadj
-
-    def rmatvecad(self, x: JaxType, y: JaxType) -> JaxType:
-        """Adjoint matrix-vector multiplication with AD
-
-        Parameters
-        ----------
-        x : :obj:`jaxlib.xla_extension.ArrayImpl`
-            Input array
-        y : :obj:`jaxlib.xla_extension.ArrayImpl`
-            Output array (where to store the
-            Vector-Jacobian product)
-
-        Returns
-        -------
-        x : :obj:`numpy.ndarray`
-            Output array of shape (N,) or (N,1)
-
         """
         M, N = self.shape
 
         if x.shape != (M,) and x.shape != (M, 1):
             raise ValueError(
-                f"Dimension mismatch. Got {x.shape}, but expected {(M, 1)} or {(M,)}."
+                f"Dimension mismatch. Got {x.shape}, but expected  ({M},) or ({M}, 1)."
             )
 
         y = self._rmatvecad(x, y)

pylops/linearoperator.py

@@ -442,10 +442,11 @@ def _matmat(self, X: NDArray) -> NDArray:
         Modified version of scipy _matmat to avoid having trailing dimension
         in col when provided to matvec
         """
+        ncp = get_array_module(X)
         if sp.sparse.issparse(X):
-            y = np.vstack([self.matvec(col.toarray().reshape(-1)) for col in X.T]).T
+            y = ncp.vstack([self.matvec(col.toarray().reshape(-1)) for col in X.T]).T
         else:
-            y = np.vstack([self.matvec(col.reshape(-1)) for col in X.T]).T
+            y = ncp.vstack([self.matvec(col.reshape(-1)) for col in X.T]).T
         return y
 
     def _rmatmat(self, X: NDArray) -> NDArray:
@@ -454,10 +455,11 @@ def _rmatmat(self, X: NDArray) -> NDArray:
         Modified version of scipy _rmatmat to avoid having trailing dimension
         in col when provided to rmatvec
         """
+        ncp = get_array_module(X)
         if sp.sparse.issparse(X):
-            y = np.vstack([self.rmatvec(col.toarray().reshape(-1)) for col in X.T]).T
+            y = ncp.vstack([self.rmatvec(col.toarray().reshape(-1)) for col in X.T]).T
         else:
-            y = np.vstack([self.rmatvec(col.reshape(-1)) for col in X.T]).T
+            y = ncp.vstack([self.rmatvec(col.reshape(-1)) for col in X.T]).T
         return y
 
     def _adjoint(self) -> LinearOperator:
@@ -509,7 +511,7 @@ def matvec(self, x: NDArray) -> NDArray:
 
         if x.shape != (N,) and x.shape != (N, 1):
             raise ValueError(
-                f"Dimension mismatch. Got {x.shape}, but expected {(M, 1)} or {(M,)}."
+                f"Dimension mismatch. Got {x.shape}, but expected ({N},) or ({N}, 1)."
             )
 
         y = self._matvec(x)
@@ -545,7 +547,7 @@ def rmatvec(self, x: NDArray) -> NDArray:
 
         if x.shape != (M,) and x.shape != (M, 1):
             raise ValueError(
-                f"Dimension mismatch. Got {x.shape}, but expected {(M, 1)} or {(M,)}."
+                f"Dimension mismatch. Got {x.shape}, but expected ({M},) or ({M}, 1)."
             )
 
         y = self._rmatvec(x)
@@ -795,7 +797,7 @@ def todense(
         Parameters
         ----------
         backend : :obj:`str`, optional
-            Backend used to densify matrix (``numpy`` or ``cupy``). Note that
+            Backend used to densify matrix (``numpy`` or ``cupy`` or ``jax``). Note that
             this must be consistent with how the operator has been created.
 
         Returns
@@ -820,7 +822,7 @@ def todense(
         if Op.shape[1] == shapemin:
             matrix = Op.matmat(identity)
         else:
-            matrix = np.conj(Op.rmatmat(identity)).T
+            matrix = ncp.conj(Op.rmatmat(identity)).T
         return matrix
 
     def tosparse(self) -> NDArray: