err

Author: Jammy2211

autofit/jax_wrapper.py

@@ -8,14 +8,9 @@
 
 from autoconf import conf
 
-DISABLE_JAX = os.environ.get("DISABLE_JAX", 0)
-
 use_jax = conf.instance["general"]["jax"]["use_jax"]
 
-if DISABLE_JAX != 0:
-    use_jax = False
-
-if use_jax and DISABLE_JAX == 0:
+if use_jax:
 
     from jax import numpy
 

autofit/mapper/prior/vectorized.py

@@ -139,10 +139,8 @@ def __call__(self, cube: np.ndarray) -> np.ndarray:
 
         if len(cube.shape) == 1:
             cube = cube[None, :]
-
             cube_reshaped = True
 
-
         out = np.empty_like(cube)
 
         # 2) Batch‐process all UniformPriors

autofit/non_linear/fitness.py

@@ -140,13 +140,18 @@ def call(self, parameters):
         -------
         The figure of merit returned to the non-linear search, which is either the log likelihood or log posterior.
         """
+        print(parameters)
+        print(type(parameters))
+        print(parameters.shape)
 
         # Get instance from model
         instance = self.model.instance_from_vector(vector=parameters)
 
         # Evaluate log likelihood (must be side-effect free and exception-free)
         log_likelihood = self.analysis.log_likelihood_function(instance=instance)
 
+        print(log_likelihood)
+
         # Penalize NaNs in the log-likelihood
         log_likelihood = xp.where(xp.isnan(log_likelihood), self.resample_figure_of_merit, log_likelihood)
 
@@ -196,8 +201,6 @@ def call_wrap(self, parameters):
 
         figure_of_merit = self._call(parameters)
 
-        print(figure_of_merit)
-
         if self.fom_is_log_likelihood:
             log_likelihood = figure_of_merit
         else: