From 28fe9862d37a40b7c25333d137a106a58db80c82 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 9 Aug 2026 12:40:04 -0400 Subject: [PATCH] fix JAX tracing for message and compound prior xp paths --- autofit/mapper/prior/abstract.py | 6 ++- autofit/mapper/prior/arithmetic/compound.py | 24 ++++++++++- autofit/messages/beta.py | 18 ++++++-- autofit/messages/gamma.py | 24 ++++++++--- .../mapper/prior/test_arithmetic_jax_trace.py | 42 +++++++++++++++++++ test_autofit/messages/test_jax_trace.py | 31 ++++++++++++++ 6 files changed, 132 insertions(+), 13 deletions(-) create mode 100644 test_autofit/mapper/prior/test_arithmetic_jax_trace.py diff --git a/autofit/mapper/prior/abstract.py b/autofit/mapper/prior/abstract.py index e2492d8e5..6ab3061d0 100644 --- a/autofit/mapper/prior/abstract.py +++ b/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): diff --git a/autofit/mapper/prior/arithmetic/compound.py b/autofit/mapper/prior/arithmetic/compound.py index 79af83feb..764cff6cd 100644 --- a/autofit/mapper/prior/arithmetic/compound.py +++ b/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, ) ) diff --git a/autofit/messages/beta.py b/autofit/messages/beta.py index 1bb44bcdf..da1344c48 100644 --- a/autofit/messages/beta.py +++ b/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) diff --git a/autofit/messages/gamma.py b/autofit/messages/gamma.py index af406a2d9..8414e0c38 100644 --- a/autofit/messages/gamma.py +++ b/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() diff --git a/test_autofit/mapper/prior/test_arithmetic_jax_trace.py b/test_autofit/mapper/prior/test_arithmetic_jax_trace.py new file mode 100644 index 000000000..b759ae216 --- /dev/null +++ b/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) diff --git a/test_autofit/messages/test_jax_trace.py b/test_autofit/messages/test_jax_trace.py index 2a2060b65..47992d935 100644 --- a/test_autofit/messages/test_jax_trace.py +++ b/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)