fitness passed around due to JAX

Author: Jammy2211

autofit/non_linear/mock/mock_search.py

@@ -103,7 +103,7 @@ def __call__(self, vector):
         fitness = Fitness(model.instance_from_vector, result=self.result)
         fitness([prior.mean for prior in model.priors_ordered_by_id])
 
-        return fitness.result
+        return fitness.result, fitness
 
     def _fit(self, model, analysis):
         if self.fit_fast:
@@ -141,9 +141,9 @@ def _fit(self, model, analysis):
         return analysis.make_result(
             samples_summary=samples_summary,
             paths=self.paths,
-        )
+        ), None
 
-    def perform_update(self, model, analysis, during_analysis, search_internal=None):
+    def perform_update(self, model, analysis, during_analysis, fitness=None, search_internal=None):
         if self.samples_summary is not None and not self.return_sensitivity_results:
             self.paths.save_samples_summary(self.samples_summary)
 

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

@@ -14,7 +14,6 @@
 from autofit.non_linear.search.nest import abstract_nest
 from autofit.non_linear.samples.sample import Sample
 from autofit.non_linear.samples.nest import SamplesNest
-from autogalaxy_workspace_test.jax_examples.func_grad import fitness
 
 logger = logging.getLogger(__name__)
 
@@ -225,7 +224,7 @@ def fit_x1_cpu(self, fitness, model, analysis):
             **self.config_dict_search,
         )
 
-        return self.call_search(search_internal=search_internal, model=model, analysis=analysis)
+        return self.call_search(search_internal=search_internal, model=model, analysis=analysis, fitness=fitness)
 
     def fit_multiprocessing(self, fitness, model, analysis):
         """
@@ -259,9 +258,9 @@ def fit_multiprocessing(self, fitness, model, analysis):
             **self.config_dict_search,
         )
 
-        return self.call_search(search_internal=search_internal, model=model, analysis=analysis)
+        return self.call_search(search_internal=search_internal, model=model, analysis=analysis, fitness=fitness)
 
-    def call_search(self, search_internal, model, analysis):
+    def call_search(self, search_internal, model, analysis, fitness):
         """
         The x1 CPU and multiprocessing searches both call this function to perform the non-linear search.