fix issues with _use_jax now being private

Author: Jammy2211

autofit/non_linear/analysis/analysis.py

@@ -109,7 +109,7 @@ def compute_latent_samples(self, samples: Samples, batch_size : Optional[int] =
 
             compute_latent_for_model = functools.partial(self.compute_latent_variables, model=samples.model)
 
-            if self.use_jax:
+            if self._use_jax:
                 import jax
                 start = time.time()
                 logger.info("JAX: Applying vmap and jit to likelihood function for latent variables -- may take a few seconds.")
@@ -130,7 +130,7 @@ def batched_compute_latent(x):
                 # batched JAX call on this chunk
                 latent_values_batch = batched_compute_latent(batch)
 
-                if self.use_jax:
+                if self._use_jax:
                     import jax.numpy as jnp
                     latent_values_batch = jnp.stack(latent_values_batch, axis=-1)  # (batch, n_latents)
                     mask = jnp.all(jnp.isfinite(latent_values_batch), axis=0)

autofit/non_linear/fitness.py

@@ -134,7 +134,7 @@ def __init__(
 
     @property
     def _xp(self):
-        if self.analysis.use_jax:
+        if self.analysis._use_jax:
             import jax.numpy as jnp
             return jnp
         return np

autofit/non_linear/search/nest/dynesty/search/abstract.py

@@ -146,7 +146,7 @@ def _fit(
                         "parallel"
                     ].get("force_x1_cpu")
                     or self.kwargs.get("force_x1_cpu")
-                    or analysis.use_jax
+                    or analysis._use_jax
                 ):
                     raise RuntimeError
 

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

@@ -127,7 +127,7 @@ def _fit(self, model: AbstractPriorModel, analysis):
         if (
             self.config_dict.get("force_x1_cpu")
             or self.kwargs.get("force_x1_cpu")
-            or analysis.use_jax
+            or analysis._use_jax
         ):
 
             fitness = Fitness(

test_autofit/graphical/hierarchical/test_optimise.py

@@ -11,13 +11,6 @@ def make_factor(hierarchical_factor):
 def test_optimise(factor):
     search = af.DynestyStatic(maxcall=100, dynamic_delta=False, delta=0.1,)
 
-    print(type(factor.analysis))
-    print(type(factor.analysis))
-    print(type(factor.analysis))
-    print(type(factor.analysis))
-    print(type(factor.analysis))
-
-
     _, status = search.optimise(
         factor.mean_field_approximation().factor_approximation(factor)
     )