test jitting

Jammy2211 · Jan 8, 2025 · 07f27f6 · 07f27f6
1 parent 62f7c54
commit 07f27f6
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/test_autolens/point/model/test_andrew_implementation.py b/test_autolens/point/model/test_andrew_implementation.py
@@ -1,3 +1,10 @@
+try:
+    import jax
+
+    JAX_INSTALLED = True
+except ImportError:
+    JAX_INSTALLED = False
+
 import numpy as np
 import pytest
 
@@ -14,23 +21,23 @@ def noise_map():
     return np.array([1.0, 1.0])
 
 
-def test_andrew_implementation(
-    data,
-    noise_map,
-):
+@pytest.fixture
+def fit(data, noise_map):
     model_positions = np.array(
         [
             (-1.0749, -1.1),
             (1.19117, 1.175),
         ]
     )
 
-    fit = Fit(
+    return Fit(
         data=data,
         noise_map=noise_map,
         model_positions=model_positions,
     )
 
+
+def test_andrew_implementation(fit):
     assert np.allclose(
         fit.all_permutations_log_likelihoods(),
         [
@@ -41,6 +48,11 @@ def test_andrew_implementation(
     assert fit.log_likelihood() == -4.40375330990644
 
 
+@pytest.mark.skipif(not JAX_INSTALLED, reason="JAX is not installed")
+def test_jax(fit):
+    assert jax.jit(fit.log_likelihood)() == -4.40375330990644
+
+
 def test_nan_model_positions(
     data,
     noise_map,