diff --git a/autofit/messages/interface.py b/autofit/messages/interface.py index 038bd44c0..3a0410624 100644 --- a/autofit/messages/interface.py +++ b/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): diff --git a/test_autofit/messages/test_jax_trace.py b/test_autofit/messages/test_jax_trace.py index 47992d935..1afa8e952 100644 --- a/test_autofit/messages/test_jax_trace.py +++ b/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)