jax wrapper moved up to autoconf

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -5,11 +5,11 @@
 from abc import ABC
 from abc import abstractmethod
 import jax.numpy as jnp
+from jax._src.tree_util import register_pytree_node
+from jax import Array
 
 from autoconf.fitsable import output_to_fits
 
-from autoarray.numpy_wrapper import register_pytree_node, Array
-
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:

autoarray/mask/derive/indexes_2d.py

@@ -4,7 +4,7 @@
 
 from autoconf import cached_property
 
-from autoarray.numpy_wrapper import register_pytree_node_class
+from jax._src.tree_util import register_pytree_node_class
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:

autoarray/numpy_wrapper.py

@@ -1,6 +1,7 @@
 import numpy as np
 import jax.numpy as jnp
 import jax
+from jax._src.tree_util import register_pytree_node_class
 from typing import Union
 
 from autoconf import conf
@@ -11,7 +12,6 @@
 
 from autoarray.operators.over_sampling import over_sample_util
 
-from autoarray.numpy_wrapper import register_pytree_node_class
 
 
 @register_pytree_node_class

autoarray/operators/over_sampling/over_sampler.py

@@ -2,11 +2,11 @@
 
 import numpy as np
 import jax.numpy as jnp
+from jax._src.tree_util import register_pytree_node_class
 
 from autoarray.structures.triangles.abstract import HEIGHT_FACTOR
 from autoarray.structures.triangles.abstract import AbstractTriangles
 from autoarray.structures.triangles.array import ArrayTriangles
-from autoarray.numpy_wrapper import register_pytree_node_class
 
 
 @register_pytree_node_class

autoarray/structures/triangles/coordinate_array.py

@@ -1,10 +1,9 @@
 from abc import ABC, abstractmethod
+from jax._src.tree_util import register_pytree_node_class
 from typing import List, Tuple
 
 import numpy as np
 
-from autoarray.numpy_wrapper import register_pytree_node_class
-
 
 class Shape(ABC):
     """

autoarray/structures/triangles/shape.py

@@ -29,8 +29,6 @@ dependencies = [
     "astropy>=5.0,<=6.1.2",
     "decorator>=4.0.0",
     "dill>=0.3.1.1",
-    "jax==0.4.28",
-    "jaxlib==0.4.28",
     "jaxnnls==1.0.1",
     "matplotlib>=3.7.0",
     "scipy<=1.14.0",

pyproject.toml

@@ -1,11 +1,8 @@
-import jax
 import jax.numpy as jnp
 
 def pytest_configure():
     _ = jnp.sum(jnp.array([0.0]))  # Force backend init
 
-jax.config.update("jax_enable_x64", True)
-
 import os
 from os import path
 import pytest

test_autoarray/conftest.py

@@ -1,4 +1,4 @@
-from autoarray.numpy_wrapper import np
+from autoconf.jax_wrapper import np
 from autoarray.structures.triangles.array import ArrayTriangles
 from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles