remove autofit imports

Author: Jammy2211

autoarray/geometry/geometry_2d.py

@@ -13,8 +13,6 @@
 from autoarray import type as ty
 from autoarray.geometry import geometry_util
 
-from autofit.jax_wrapper import use_jax
-
 logging.basicConfig()
 logger = logging.getLogger(__name__)
 

autoarray/operators/over_sampling/over_sampler.py

@@ -10,7 +10,7 @@
 
 from autoarray.operators.over_sampling import over_sample_util
 
-from autofit.jax_wrapper import register_pytree_node_class
+from autoarray.numpy_wrapper import register_pytree_node_class
 
 
 @register_pytree_node_class

autoarray/structures/vectors/uniform.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 import jax.numpy as jnp
-# from autofit.jax_wrapper import numpy as np, use_jax
 from typing import List, Optional, Tuple, Union
 
 from autoarray.structures.arrays.uniform_2d import Array2D

test_autoarray/test_jax_changes.py

@@ -1,8 +1,10 @@
-import autoarray as aa
+import jax.numpy as jnp
 import pytest
 
+
+import autoarray as aa
+
 from autoarray import Grid2D, Mask2D
-from autofit.jax_wrapper import numpy as np
 
 
 @pytest.fixture(name="array")
@@ -33,4 +35,4 @@ def test_boolean_issue():
         mask=Mask2D.all_false((10, 10), pixel_scales=1.0),
     )
     values, keys = Grid2D.instance_flatten(grid)
-    np.array(Grid2D.instance_unflatten(keys, values))
+    jnp.array(Grid2D.instance_unflatten(keys, values))