xp module not class attribute

Author: Jammy2211

autolens/interferometer/fit_interferometer.py

@@ -79,7 +79,16 @@ def __init__(
         )
 
         self.preloads = 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 profile_visibilities(self) -> aa.Visibilities:

autolens/point/fit/abstract.py

@@ -72,7 +72,15 @@ def __init__(
                 f"in the tracer (make sure your tracer's point source name is the same the dataset name."
             )
 
-        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):

autolens/point/fit/dataset.py

@@ -125,7 +125,15 @@ def __init__(
         except exc.PointExtractionException:
             self.time_delays = None
 
-        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 model_obj(self):

autolens/point/solver/shape_solver.py

@@ -47,7 +47,15 @@ def __init__(
 
         self.initial_triangles = initial_triangles
 
-        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
 
     # noinspection PyPep8Naming
     @classmethod