xp module not class attribute

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -73,7 +73,16 @@ def __init__(self, array, xp=np):
         while isinstance(array, AbstractNDArray):
             array = array.array
         self._array = array
-        self._xp = xp
+
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     def invert(self):
         new = self.copy()

autoarray/fit/fit_dataset.py

@@ -158,7 +158,15 @@ def __init__(
         self.dataset.grids.blurring
         self.dataset.grids.border_relocator
 
-        self._xp = xp
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     @property
     def mask(self) -> Mask2D:

autoarray/inversion/inversion/abstract.py

@@ -75,7 +75,15 @@ def __init__(
 
         self.preloads = preloads or Preloads()
 
-        self._xp = xp
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     @property
     def data(self):

autoarray/inversion/linear_obj/linear_obj.py

@@ -27,7 +27,16 @@ def __init__(self, regularization: Optional[AbstractRegularization], xp=np):
             The regularization scheme which may be applied to this linear object in order to smooth its solution.
         """
         self.regularization = regularization
-        self._xp = xp
+
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     @property
     def params(self) -> int:

autoarray/mask/abstract_mask.py

@@ -57,7 +57,16 @@ def __init__(
 
         self.pixel_scales = pixel_scales
         self.origin = origin
-        self._xp = xp
+
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     @property
     def mask(self):

autoarray/mask/derive/grid_2d.py

@@ -61,7 +61,16 @@ def __init__(self, mask: Mask2D, xp=np):
             print(derive_grid_2d.border)
         """
         self.mask = mask
-        self._xp = xp
+
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     def tree_flatten(self):
         return (self.mask,), ()

autoarray/mask/derive/indexes_2d.py

@@ -62,7 +62,16 @@ def __init__(self, mask: Mask2D, xp=np):
             print(derive_indexes_2d.edge_native)
         """
         self.mask = mask
-        self._xp = xp
+
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
 
     def tree_flatten(self):
         return (self.mask,), ()

autoarray/structures/decorators/abstract.py

@@ -52,10 +52,19 @@ def __init__(self, func, obj, grid, xp=np, *args, **kwargs):
         self.func = func
         self.obj = obj
         self.grid = grid
-        self._xp = xp
         self.args = args
         self.kwargs = kwargs
 
+        self.use_jax = xp is not np
+
+    @property
+    def _xp(self):
+        if self.use_jax:
+            import jax.numpy as jnp
+
+            return jnp
+        return np
+
     @property
     def mask(self) -> Union[Mask1D, Mask2D]:
         return self.grid.mask