minor: add explicit and clinear to JaxOperator

mrava87 · Jun 24, 2024 · f5205ab · f5205ab
1 parent 72c30d1
commit f5205ab
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 3 deletions.
diff --git a/docs/source/gpu.rst b/docs/source/gpu.rst
@@ -1,7 +1,7 @@
 .. _gpu:
 
-GPU and TPU Support
-===================
+GPU / TPU Support
+=================
 
 Overview
 --------
@@ -238,6 +238,43 @@ Signal processing:
      - |:white_check_mark:|
      - |:white_check_mark:|
      - |:white_check_mark:|
+   * - :class:`pylops.signalprocessing.DWT`
+     - |:white_check_mark:|
+     - |:red_circle:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.DWT2D`
+     - |:white_check_mark:|
+     - |:red_circle:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.DCT`
+     - |:white_check_mark:|
+     - |:red_circle:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.Seislet`
+     - |:white_check_mark:|
+     - |:red_circle:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.Radon2D`
+     - |:white_check_mark:|
+     - |:red_circle:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.Radon3D`
+     - |:white_check_mark:|
+     - |:red_circle:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.ChirpRadon2D`
+     - |:white_check_mark:|
+     - |:white_check_mark:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.ChirpRadon3D`
+     - |:white_check_mark:|
+     - |:white_check_mark:|
+     - |:red_circle:|
+   * - :class:`pylops.signalprocessing.Sliding1D`
+     - |:white_check_mark:|
+     - |:white_check_mark:|
+     - |:white_check_mark:|
+
 
 
 .. warning::

diff --git a/pylops/jaxoperator.py b/pylops/jaxoperator.py
@@ -20,7 +20,15 @@
 
 class JaxOperator(LinearOperator):
     def __init__(self, Op: LinearOperator) -> None:
-        super().__init__(dtype=Op.dtype, dims=Op.dims, dimsd=Op.dimsd, name=Op.name)
+        super().__init__(
+            dtype=Op.dtype,
+            dims=Op.dims,
+            dimsd=Op.dimsd,
+            clinear=Op.clinear,
+            explicit=Op.explicit,
+            forceflat=Op.forceflat,
+            name=Op.name,
+        )
         self._matvec = jax.jit(Op._matvec)
         self._rmatvec = jax.jit(Op._rmatvec)