all tests pass

Author: Jammy2211

autofit/jax_wrapper.py

@@ -12,6 +12,9 @@
 
 use_jax = conf.instance["general"]["jax"]["use_jax"]
 
+if DISABLE_JAX != 0:
+    use_jax = False
+
 if use_jax and DISABLE_JAX == 0:
 
     from jax import numpy

autofit/mapper/prior/vectorized.py

@@ -135,7 +135,14 @@ def __call__(self, cube: np.ndarray) -> np.ndarray:
             Transformed parameters of shape (n_samples, n_priors).
         """
 
-        n_samples, n_priors = cube.shape
+        cube_reshaped = False
+
+        if len(cube.shape) == 1:
+            cube = cube[None, :]
+
+            cube_reshaped = True
+
+
         out = np.empty_like(cube)
 
         # 2) Batch‐process all UniformPriors
@@ -177,4 +184,8 @@ def __call__(self, cube: np.ndarray) -> np.ndarray:
                 + subcube * (self.loguniform_log_uppers - self.loguniform_log_lowers)
             )
 
+        if cube_reshaped:
+
+            return out[0]
+
         return out

autofit/non_linear/fitness.py

@@ -1,5 +1,6 @@
 import jax
 import logging
+import numpy as np
 import os
 import time
 
@@ -108,9 +109,6 @@ def __init__(
         self.convert_to_chi_squared = convert_to_chi_squared
         self.store_history = store_history
 
-        if self.paths is not None:
-            self.check_log_likelihood(fitness=self)
-
         self.parameters_history_list = []
         self.log_likelihood_history_list = []
 
@@ -122,6 +120,9 @@ def __init__(
             if self.use_jax_vmap:
                 self._call = self._vmap
 
+        if self.paths is not None:
+            self.check_log_likelihood(fitness=self)
+
     def call(self, parameters):
         """
         A private method that calls the fitness function with the given parameters and additional keyword arguments.
@@ -189,8 +190,14 @@ def call_wrap(self, parameters):
             the log-likelihood itself or another objective function value,
             depending on configuration.
         """
+        if self.use_jax_vmap:
+            if len(np.array(parameters).shape) == 1:
+                parameters = np.array(parameters)[None, :]
+
         figure_of_merit = self._call(parameters)
 
+        print(figure_of_merit)
+
         if self.fom_is_log_likelihood:
             log_likelihood = figure_of_merit
         else:
@@ -224,7 +231,7 @@ def __call__(self, parameters, *kwargs):
         -------
         The figure of merit returned to the non-linear search, which is either the log likelihood or log posterior.
         """
-        return self._call(parameters)
+        return self.call_wrap(parameters)
 
     # def __getstate__(self):
     #     state = self.__dict__.copy()

autofit/non_linear/search/abstract_search.py

@@ -22,6 +22,8 @@
 
 from autoconf.output import should_output
 
+from autofit.jax_wrapper import numpy as xp
+
 from autofit import exc, jax_wrapper
 from autofit.database.sqlalchemy_ import sa
 from autofit.graphical import (
@@ -913,7 +915,6 @@ def perform_update(
         )
         self.paths.save_samples(samples=samples_save)
 
-        latent_samples = None
 
         if (during_analysis and conf.instance["output"]["latent_during_fit"]) or (
             not during_analysis and conf.instance["output"]["latent_after_fit"]
@@ -978,7 +979,7 @@ def perform_update(
                 parameters = samples.max_log_likelihood(as_instance=False)
 
                 start = time.time()
-                figure_of_merit = fitness(parameters)
+                figure_of_merit = fitness.call_wrap(parameters)
 
                 # account for asynchronous JAX calls
                 np.array(figure_of_merit)

autofit/non_linear/search/nest/nautilus/search.py

@@ -145,28 +145,21 @@ def _fit(self, model: AbstractPriorModel, analysis):
                 analysis=analysis,
             )
         else:
-            if not self.using_mpi:
 
-                fitness = Fitness(
-                    model=model,
-                    analysis=analysis,
-                    paths=self.paths,
-                    fom_is_log_likelihood=True,
-                    resample_figure_of_merit=-1.0e99,
-                )
+            fitness = Fitness(
+                model=model,
+                analysis=analysis,
+                paths=self.paths,
+                fom_is_log_likelihood=True,
+                resample_figure_of_merit=-1.0e99,
+            )
+
+            search_internal = self.fit_multiprocessing(
+                fitness=fitness,
+                model=model,
+                analysis=analysis,
+            )
 
-                search_internal = self.fit_multiprocessing(
-                    fitness=fitness,
-                    model=model,
-                    analysis=analysis,
-                )
-            else:
-                search_internal = self.fit_mpi(
-                    fitness=fitness,
-                    model=model,
-                    analysis=analysis,
-                    checkpoint_exists=checkpoint_exists,
-                )
         return search_internal, fitness
 
     @property
@@ -229,7 +222,6 @@ def fit_x1_cpu(self, fitness, model, analysis):
             prior=PriorVectorized(model=model),
             likelihood=fitness.call_wrap,
             n_dim=model.prior_count,
-         #   prior_kwargs={"model": model},
             filepath=self.checkpoint_file,
             pool=None,
             vectorized=True,
@@ -263,7 +255,6 @@ def fit_multiprocessing(self, fitness, model, analysis):
             prior=PriorVectorized(model=model),
             likelihood=fitness.call_wrap,
             n_dim=model.prior_count,
-            prior_kwargs={"model": model},
             filepath=self.checkpoint_file,
             pool=self.number_of_cores,
             **self.config_dict_search,

test_autofit/non_linear/test_dict.py

@@ -34,7 +34,6 @@ def make_dynesty_dict():
             "sample": "auto",
             "slices": 5,
             "unique_tag": None,
-            "use_gradient": False,
             "walks": 5,
         },
         "class_path": "autofit.non_linear.search.nest.dynesty.search.static.DynestyStatic",