Use xp.stack in messages so JAX can trace fixed arrays

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]:

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):

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:
         """

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:
         """

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)