64 bit config in

Author: Jammy2211

autoarray/inversion/inversion/abstract.py

@@ -638,7 +638,7 @@ def log_det_curvature_reg_matrix_term(self) -> float:
         if not self.has(cls=AbstractRegularization):
             return 0.0
 
-        return 2.0 * np.sum(
+        return 2.0 * jnp.sum(
             jnp.log(jnp.diag(jnp.linalg.cholesky(self.curvature_reg_matrix_reduced)))
         )
 
@@ -659,7 +659,7 @@ def log_det_regularization_matrix_term(self) -> float:
         if not self.has(cls=AbstractRegularization):
             return 0.0
 
-        return 2.0 * np.sum(
+        return 2.0 * jnp.sum(
             jnp.log(jnp.diag(jnp.linalg.cholesky(self.regularization_matrix_reduced)))
         )
 

test_autoarray/conftest.py

@@ -1,9 +1,10 @@
+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