diff --git a/autofit/messages/beta.py b/autofit/messages/beta.py index 3716ffcee..1bb44bcdf 100644 --- a/autofit/messages/beta.py +++ b/autofit/messages/beta.py @@ -223,7 +223,7 @@ def calc_natural_parameters( ------- Natural parameters [alpha - 1, beta - 1]. """ - return xp.array([alpha - 1, beta - 1]) + return xp.stack([alpha - 1, beta - 1]) @staticmethod def invert_natural_parameters( @@ -276,7 +276,7 @@ def to_canonical_form(cls, x: np.ndarray, xp=np) -> np.ndarray: ------- Canonical sufficient statistics [log(x), log(1 - x)]. """ - return xp.array([xp.log(x), xp.log1p(-x)]) + return xp.stack([xp.log(x), xp.log1p(-x)]) @cached_property def mean(self) -> Union[np.ndarray, float]: diff --git a/autofit/messages/gamma.py b/autofit/messages/gamma.py index 63db00630..af406a2d9 100644 --- a/autofit/messages/gamma.py +++ b/autofit/messages/gamma.py @@ -41,7 +41,7 @@ def natural_parameters(self, xp=np) -> np.ndarray: @staticmethod def calc_natural_parameters(alpha, beta, xp=np): - return xp.array([alpha - 1, -beta]) + return xp.stack([alpha - 1, -beta]) @staticmethod def invert_natural_parameters(natural_parameters): @@ -50,7 +50,7 @@ def invert_natural_parameters(natural_parameters): @staticmethod def to_canonical_form(x, xp=np): - return xp.array([np.log(x), x]) + return xp.stack([xp.log(x), x]) @classmethod def invert_sufficient_statistics(cls, suff_stats): diff --git a/autofit/messages/normal.py b/autofit/messages/normal.py index 77f222fc3..1dbf1c927 100644 --- a/autofit/messages/normal.py +++ b/autofit/messages/normal.py @@ -196,7 +196,7 @@ def calc_natural_parameters(mu : Union[float, np.ndarray], sigma : Union[float, η₂ = -1 / (2σ²) """ precision = 1 / sigma**2 - return xp.array([mu * precision, -precision / 2]) + return xp.stack([mu * precision, -precision / 2]) @staticmethod def invert_natural_parameters(natural_parameters : np.ndarray) -> Tuple[float, float]: @@ -235,7 +235,7 @@ def to_canonical_form(x : Union[float, np.ndarray], xp=np) -> np.ndarray: ------- The sufficient statistics [x, x²]. """ - return xp.array([x, x**2]) + return xp.stack([x, x**2]) @classmethod def invert_sufficient_statistics(cls, suff_stats: Tuple[float, float]) -> np.ndarray: @@ -591,7 +591,7 @@ def calc_natural_parameters(eta1: float, eta2: float, xp=np) -> np.ndarray: eta2 The second natural parameter. """ - return xp.array([eta1, eta2]) + return xp.stack([eta1, eta2]) def natural_parameters(self, xp=np) -> np.ndarray: """ diff --git a/autofit/messages/truncated_normal.py b/autofit/messages/truncated_normal.py index 2687d90e9..4c31a111a 100644 --- a/autofit/messages/truncated_normal.py +++ b/autofit/messages/truncated_normal.py @@ -204,7 +204,7 @@ def calc_natural_parameters(mu : Union[float, np.ndarray], sigma : Union[float, η₂ = -1 / (2σ²) """ precision = 1 / sigma**2 - return xp.array([mu * precision, -precision / 2]) + return xp.stack([mu * precision, -precision / 2]) @staticmethod def invert_natural_parameters(natural_parameters : np.ndarray) -> Tuple[float, float]: @@ -248,7 +248,7 @@ def to_canonical_form(x : Union[float, np.ndarray], xp=np) -> np.ndarray: ------- The sufficient statistics [x, x²]. """ - return xp.array([x, x**2]) + return xp.stack([x, x**2]) @classmethod def invert_sufficient_statistics(cls, suff_stats: Tuple[float, float]) -> np.ndarray: @@ -722,7 +722,7 @@ def calc_natural_parameters( eta2 The second natural parameter. """ - return xp.array([eta1, eta2]) + return xp.stack([eta1, eta2]) def natural_parameters(self, xp=np) -> np.ndarray: """ diff --git a/test_autofit/messages/test_jax_trace.py b/test_autofit/messages/test_jax_trace.py new file mode 100644 index 000000000..2a2060b65 --- /dev/null +++ b/test_autofit/messages/test_jax_trace.py @@ -0,0 +1,82 @@ +import numpy as np +import pytest + +from autofit.messages.beta import BetaMessage +from autofit.messages.gamma import GammaMessage +from autofit.messages.normal import NaturalNormal, NormalMessage +from autofit.messages.truncated_normal import ( + TruncatedNaturalNormal, + TruncatedNormalMessage, +) + +jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") + + +MESSAGE_ARRAY_CASES = [ + pytest.param( + lambda value, xp: NormalMessage.calc_natural_parameters( + value, value + 1.0, xp=xp + ), + id="normal-natural-parameters", + ), + pytest.param( + lambda value, xp: NormalMessage.to_canonical_form(value, xp=xp), + id="normal-canonical-form", + ), + pytest.param( + lambda value, xp: NaturalNormal.calc_natural_parameters(value, -value, xp=xp), + id="natural-normal-natural-parameters", + ), + pytest.param( + lambda value, xp: TruncatedNormalMessage.calc_natural_parameters( + value, value + 1.0, xp=xp + ), + id="truncated-normal-natural-parameters", + ), + pytest.param( + lambda value, xp: TruncatedNormalMessage.to_canonical_form(value, xp=xp), + id="truncated-normal-canonical-form", + ), + pytest.param( + lambda value, xp: TruncatedNaturalNormal.calc_natural_parameters( + value, -value, xp=xp + ), + id="truncated-natural-normal-natural-parameters", + ), + pytest.param( + lambda value, xp: BetaMessage.calc_natural_parameters( + value + 1.0, value + 2.0, xp=xp + ), + id="beta-natural-parameters", + ), + pytest.param( + lambda value, xp: BetaMessage.to_canonical_form(value, xp=xp), + id="beta-canonical-form", + ), + pytest.param( + lambda value, xp: GammaMessage.calc_natural_parameters( + value + 1.0, value + 2.0, xp=xp + ), + id="gamma-natural-parameters", + ), + pytest.param( + lambda value, xp: GammaMessage.to_canonical_form(value, xp=xp), + id="gamma-canonical-form", + ), +] + + +@pytest.mark.parametrize("message_array", MESSAGE_ARRAY_CASES) +@pytest.mark.parametrize( + "value", + [pytest.param(0.25, id="scalar"), pytest.param([0.25, 0.5], id="batched")], +) +def test_message_array_construction_is_jittable_and_matches_numpy(message_array, value): + numpy_value = np.asarray(value) + expected = message_array(numpy_value, np) + + actual = jax.jit(lambda traced: message_array(traced, jnp))(jnp.asarray(value)) + + assert actual.shape == expected.shape + np.testing.assert_allclose(np.asarray(actual), expected)