jax imports deferred

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -4,7 +4,7 @@
 
 from abc import ABC
 from abc import abstractmethod
-import jax.numpy as jnp
+
 from jax._src.tree_util import register_pytree_node
 
 import numpy as np
@@ -88,7 +88,7 @@ def __init__(self, array, xp=np):
 
     def invert(self):
         new = self.copy()
-        new._array = jnp.invert(new._array)
+        new._array = self._xp.invert(new._array)
         return new
 
     @classmethod
@@ -117,7 +117,7 @@ def instance_unflatten(cls, aux_data, children):
             setattr(instance, key, value)
         return instance
 
-    def with_new_array(self, array: jnp.ndarray) -> "AbstractNDArray":
+    def with_new_array(self, array: np.ndarray) -> "AbstractNDArray":
         """
         Copy this object but give it a new array.
 
@@ -137,10 +137,9 @@ def with_new_array(self, array: jnp.ndarray) -> "AbstractNDArray":
         new_array._array = array
         return new_array
 
-    @staticmethod
-    def flip_hdu_for_ds9(values):
+    def flip_hdu_for_ds9(self, values):
         if conf.instance["general"]["fits"]["flip_for_ds9"]:
-            return jnp.flipud(values)
+            return self._xp.flipud(values)
         return values
 
     def copy(self):
@@ -170,7 +169,7 @@ def __iter__(self):
 
     @to_new_array
     def sqrt(self):
-        return jnp.sqrt(self._array)
+        return self._xp.sqrt(self._array)
 
     @property
     def array(self):
@@ -333,7 +332,10 @@ def __getattr__(self, item):
         )
 
     def __getitem__(self, item):
+
+        import jax.numpy as jnp
         result = self._array[item]
+
         if isinstance(item, slice):
             result = self.with_new_array(result)
         if isinstance(result, jnp.ndarray):
@@ -342,6 +344,7 @@ def __getitem__(self, item):
 
     def __setitem__(self, key, value):
         from jax import Array
+        import jax.numpy as jnp
 
         if isinstance(key, (jnp.ndarray, AbstractNDArray, Array)):
             self._array = jnp.where(key, value, self._array)