Make remaining message and prior xp paths JAX-traceable

autofit/mapper/prior/abstract.py

@@ -185,6 +185,7 @@ def instance_for_arguments(
         self,
         arguments,
         ignore_assertions=False,
+        xp=np,
     ):
         """Look up this prior's value in an arguments dictionary.
 
@@ -194,8 +195,11 @@ def instance_for_arguments(
             A dictionary mapping Prior objects to physical values.
         ignore_assertions
             Unused for priors (present for interface compatibility).
+        xp
+            Unused for direct prior lookup. Accepted so compound prior models can
+            propagate their selected NumPy/JAX backend recursively.
         """
-        _ = ignore_assertions
+        _ = ignore_assertions, xp
         return arguments[self]
 
     def project(self, samples, log_weight_list):

autofit/mapper/prior/arithmetic/compound.py

@@ -149,6 +149,7 @@ def left_for_arguments(
         self,
         arguments: dict,
         ignore_assertions=False,
+        xp=np,
     ):
         """
         Instantiate the left object.
@@ -168,6 +169,7 @@ def left_for_arguments(
             return self._left.instance_for_arguments(
                 arguments,
                 ignore_assertions=ignore_assertions,
+                xp=xp,
             )
         except AttributeError:
             return self._left
@@ -176,6 +178,7 @@ def right_for_arguments(
         self,
         arguments: dict,
         ignore_assertions=False,
+        xp=np,
     ):
         """
         Instantiate the right object.
@@ -195,6 +198,7 @@ def right_for_arguments(
             return self._right.instance_for_arguments(
                 arguments,
                 ignore_assertions=ignore_assertions,
+                xp=xp,
             )
         except AttributeError:
             return self._right
@@ -217,9 +221,11 @@ def _instance_for_arguments(
         return self.left_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         ) + self.right_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
     def __str__(self):
@@ -243,9 +249,11 @@ def _instance_for_arguments(
         return self.left_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         ) * self.right_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
 
@@ -263,9 +271,11 @@ def _instance_for_arguments(
         return self.left_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         ) / self.right_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
 
@@ -283,9 +293,11 @@ def _instance_for_arguments(
         return self.left_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         ) // self.right_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
 
@@ -303,9 +315,11 @@ def _instance_for_arguments(
         return self.left_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         ) % self.right_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
 
@@ -323,9 +337,11 @@ def _instance_for_arguments(
         return self.left_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         ) ** self.right_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
 
@@ -407,6 +423,7 @@ def _instance_for_arguments(
         return -self.prior.instance_for_arguments(
             arguments,
             ignore_assertions=ignore_assertions,
+            xp=xp,
         )
 
 
@@ -425,6 +442,7 @@ def _instance_for_arguments(
             self.prior.instance_for_arguments(
                 arguments,
                 ignore_assertions=ignore_assertions,
+                xp=xp,
             )
         )
 
@@ -440,10 +458,11 @@ def _instance_for_arguments(
         ignore_assertions=False,
         xp=np,
     ):
-        return np.log(
+        return xp.log(
             self.prior.instance_for_arguments(
                 arguments,
                 ignore_assertions=ignore_assertions,
+                xp=xp,
             )
         )
 
@@ -459,9 +478,10 @@ def _instance_for_arguments(
         ignore_assertions=False,
         xp=np,
     ):
-        return np.log10(
+        return xp.log10(
             self.prior.instance_for_arguments(
                 arguments,
                 ignore_assertions=ignore_assertions,
+                xp=xp,
             )
         )

autofit/messages/beta.py

@@ -146,14 +146,21 @@ def __init__(
         id_
             Identifier for the message. Default is None.
         """
-        self.alpha = alpha
-        self.beta = beta
+        if isinstance(alpha, (np.ndarray, float, int, list)):
+            xp = np
+        else:
+            import jax.numpy as jnp
+
+            xp = jnp
+
         super().__init__(
             alpha,
             beta,
             log_norm=log_norm,
-            id_=id_
+            id_=id_,
+            _xp=xp,
         )
+        self.alpha, self.beta = self.parameters
 
     def value_for(self, unit: float) -> float:
         """
@@ -184,7 +191,10 @@ def log_partition(self, xp=np) -> np.ndarray:
         -------
         The value of the log Beta function, i.e. betaln(alpha, beta).
         """
-        from scipy.special import betaln
+        if xp is np:
+            from scipy.special import betaln
+        else:
+            from jax.scipy.special import betaln
 
         return betaln(*self.parameters)
 

autofit/messages/gamma.py

@@ -8,10 +8,15 @@
 class GammaMessage(AbstractMessage):
 
     def log_partition(self, xp=np):
-        from scipy import special
+        if xp is np:
+            from scipy.special import gammaln
+        else:
+            from jax.scipy.special import gammaln
 
-        alpha, beta = GammaMessage.invert_natural_parameters(self.natural_parameters(xp=xp))
-        return special.gammaln(alpha) - alpha * np.log(beta)
+        alpha, beta = GammaMessage.invert_natural_parameters(
+            self.natural_parameters(xp=xp)
+        )
+        return gammaln(alpha) - alpha * xp.log(beta)
 
     log_base_measure = 0.0
     _support = ((0, np.inf),)
@@ -24,14 +29,21 @@ def __init__(
             log_norm=0.0,
             id_=None
     ):
-        self.alpha = alpha
-        self.beta = beta
+        if isinstance(alpha, (np.ndarray, float, int, list)):
+            xp = np
+        else:
+            import jax.numpy as jnp
+
+            xp = jnp
+
         super().__init__(
             alpha,
             beta,
             log_norm=log_norm,
-            id_=id_
+            id_=id_,
+            _xp=xp,
         )
+        self.alpha, self.beta = self.parameters
 
     def value_for(self, unit: float) -> float:
         raise NotImplemented()

test_autofit/mapper/prior/test_arithmetic_jax_trace.py

@@ -0,0 +1,42 @@
+import numpy as np
+import pytest
+
+import autofit as af
+
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
+
+
+COMPOUND_PRIOR_CASES = [
+    pytest.param(lambda prior: af.Log(prior), id="log"),
+    pytest.param(lambda prior: af.Log10(prior), id="log10"),
+    pytest.param(
+        lambda prior: af.Log(prior) + af.Log10(prior),
+        id="compound-children",
+    ),
+    pytest.param(lambda prior: af.Log(af.Log(prior)), id="nested-modifier"),
+]
+
+
+@pytest.mark.parametrize("compound_prior_from", COMPOUND_PRIOR_CASES)
+@pytest.mark.parametrize(
+    "value",
+    [pytest.param(4.0, id="scalar"), pytest.param([4.0, 10.0], id="batched")],
+)
+def test_compound_prior_is_jittable_and_matches_numpy(compound_prior_from, value):
+    prior = af.GaussianPrior(mean=0.0, sigma=1.0)
+    compound_prior = compound_prior_from(prior)
+
+    def evaluate(prior_value, xp):
+        return compound_prior.instance_for_arguments(
+            {prior: prior_value},
+            ignore_assertions=True,
+            xp=xp,
+        )
+
+    numpy_value = np.asarray(value)
+    expected = evaluate(numpy_value, np)
+    actual = jax.jit(lambda traced: evaluate(traced, jnp))(jnp.asarray(value))
+
+    assert actual.shape == np.shape(expected)
+    np.testing.assert_allclose(np.asarray(actual), expected, rtol=1e-6)

test_autofit/messages/test_jax_trace.py

@@ -80,3 +80,34 @@ def test_message_array_construction_is_jittable_and_matches_numpy(message_array,
 
     assert actual.shape == expected.shape
     np.testing.assert_allclose(np.asarray(actual), expected)
+
+
+MESSAGE_LOG_PARTITION_CASES = [
+    pytest.param(
+        lambda value, xp: GammaMessage(value + 1.0, value + 2.0).log_partition(xp=xp),
+        id="gamma",
+    ),
+    pytest.param(
+        lambda value, xp: BetaMessage(value + 1.0, value + 2.0).log_partition(xp=xp),
+        id="beta",
+    ),
+]
+
+
+@pytest.mark.parametrize("message_log_partition", MESSAGE_LOG_PARTITION_CASES)
+@pytest.mark.parametrize(
+    "value",
+    [pytest.param(2.0, id="scalar"), pytest.param([2.0, 3.0], id="batched")],
+)
+def test_message_log_partition_is_jittable_and_matches_numpy(
+    message_log_partition, value
+):
+    numpy_value = np.asarray(value)
+    expected = message_log_partition(numpy_value, np)
+
+    actual = jax.jit(lambda traced: message_log_partition(traced, jnp))(
+        jnp.asarray(value)
+    )
+
+    assert actual.shape == np.shape(expected)
+    np.testing.assert_allclose(np.asarray(actual), expected, rtol=1e-6)