fix numba import

Author: Jammy2211

autoarray/config/general.yaml

@@ -13,7 +13,7 @@ inversion:
   reconstruction_vmax_factor: 0.5     # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
 numba:
   use_numba: true
-  cache: false
+  cache: true
   nopython: true
   parallel: false
 pixelization:

autoarray/inversion/inversion/abstract.py

@@ -1,5 +1,5 @@
 import copy
-from jax.scipy.linalg import block_diag
+
 import numpy as np
 from typing import Dict, List, Optional, Type, Union, TYPE_CHECKING
 
@@ -331,6 +331,12 @@ def regularization_matrix(self) -> Optional[np.ndarray]:
         If the `settings.force_edge_pixels_to_zeros` is `True`, the edge pixels of each mapper in the inversion
         are regularized so high their value is forced to zero.
         """
+        if self.xp.__name__.startswith("jax"):
+            from jax.scipy.linalg import block_diag
+            return block_diag(
+                *[linear_obj.regularization_matrix for linear_obj in self.linear_obj_list]
+            )
+        from scipy.linalg import block_diag
         return block_diag(
             *[linear_obj.regularization_matrix for linear_obj in self.linear_obj_list]
         )

autoarray/numba_util.py

@@ -14,7 +14,6 @@
     cache = True
     parallel = False
 
-
 def jit(nopython=nopython, cache=cache, parallel=parallel):
 
     def wrapper(func):

test_autoarray/config/general.yaml

@@ -15,8 +15,8 @@ inversion:
   no_regularization_add_to_curvature_diag_value : 1.0e-8 # The default value added to the curvature matrix's diagonal when regularization is not applied to a linear object, which prevents inversion's failing due to the matrix being singular.
   positive_only_uses_p_initial: false  # If True, the positive-only solver of an inversion's uses an initial guess of the reconstructed data's values as which values should be positive, speeding up the solver.
 numba:
-  cache: true
   nopython: true
+  cache: true
   parallel: false
   use_numba: true
 pixelization:

test_autoarray/inversion/inversion/test_factory.py

@@ -63,6 +63,7 @@ def test__inversion_imaging__via_mapper(
     rectangular_mapper_7x7_3x3,
     delaunay_mapper_9_3x3,
 ):
+
     inversion = aa.Inversion(
         dataset=masked_imaging_7x7_no_blur,
         linear_obj_list=[rectangular_mapper_7x7_3x3],
@@ -114,7 +115,6 @@ def test__inversion_imaging__via_mapper(
     assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(10.6674, 1.0e-4)
     assert inversion.mapped_reconstructed_image == pytest.approx(np.ones(9), 1.0e-4)
 
-    fff
 
 
 def test__inversion_imaging__via_regularizations(