fix: message shape/size/ndim read the real JAX broadcast shape (jax 0.11 compat)

autofit/messages/interface.py

@@ -24,18 +24,29 @@ def broadcast(self):
     @property
     def shape(self) -> Tuple[int, ...]:
 
-        # JAX behaviour
-        if isinstance(self.broadcast, list):
-            return ()
-
-        return self.broadcast.shape
+        # jnp.broadcast_arrays returns a list on jax <= 0.10 and a tuple on
+        # jax >= 0.11 (mirroring the NumPy 2 change to np.broadcast_arrays),
+        # so the container is matched on (list, tuple) and the broadcast shape
+        # read off its first element — every element already carries the
+        # common broadcast shape.
+        broadcast = self.broadcast
+        if isinstance(broadcast, (list, tuple)):
+            if not broadcast:
+                return ()
+            return np.shape(broadcast[0])
+
+        return broadcast.shape
 
     @property
     def size(self) -> int:
+        if isinstance(self.broadcast, (list, tuple)):
+            return int(np.prod(self.shape, dtype=int))
         return self.broadcast.size
 
     @property
     def ndim(self) -> int:
+        if isinstance(self.broadcast, (list, tuple)):
+            return len(self.shape)
         return self.broadcast.ndim
 
     def __eq__(self, other):

test_autofit/messages/test_jax_trace.py

@@ -111,3 +111,59 @@ def test_message_log_partition_is_jittable_and_matches_numpy(
 
     assert actual.shape == np.shape(expected)
     np.testing.assert_allclose(np.asarray(actual), expected, rtol=1e-6)
+
+
+MESSAGE_PARITY_CASES = [
+    pytest.param(
+        lambda params, xp: NormalMessage(xp.asarray(params), xp.asarray(params) + 1.0),
+        0.5,
+        id="normal",
+    ),
+    pytest.param(
+        lambda params, xp: BetaMessage(
+            xp.asarray(params) + 1.0, xp.asarray(params) + 2.0
+        ),
+        0.25,
+        id="beta",
+    ),
+    pytest.param(
+        lambda params, xp: GammaMessage(
+            xp.asarray(params) + 1.0, xp.asarray(params) + 2.0
+        ),
+        0.5,
+        id="gamma",
+    ),
+]
+
+
+@pytest.mark.parametrize("make_message, x_offset", MESSAGE_PARITY_CASES)
+@pytest.mark.parametrize(
+    "params",
+    [pytest.param(1.0, id="scalar"), pytest.param([1.0, 2.0], id="batched")],
+)
+def test_message_shape_and_logpdf_match_numpy(make_message, x_offset, params):
+    """
+    A JAX-backed message must report the same shape/size/ndim as its NumPy
+    twin, and batched logpdf must return the same values and shape.
+
+    Guards the `()` shape sentinel regression: with `shape` hard-wired to `()`
+    for the JAX branch, `_broadcast_natural_parameters` matched the
+    `shape[1:] == self.shape` branch for batched messages, inserted a spurious
+    axis and returned an (n, n) matrix of wrong values instead of the (n,)
+    vector NumPy produces (#1510).
+    """
+    numpy_message = make_message(np.asarray(params), np)
+    jax_message = make_message(jnp.asarray(params), jnp)
+
+    assert jax_message.shape == numpy_message.shape
+    assert jax_message.size == numpy_message.size
+    assert jax_message.ndim == numpy_message.ndim
+
+    x = np.asarray(params) * 0.0 + x_offset
+    expected = numpy_message.logpdf(x, xp=np)
+
+    actual = jax_message.logpdf(jnp.asarray(x), xp=jnp)
+
+    assert np.shape(actual) == np.shape(expected)
+    # rtol reflects float32 accumulation in the JAX natural_logpdf path.
+    np.testing.assert_allclose(np.asarray(actual), expected, rtol=1e-4)