From ffa46d19647dc2db26901564a9db9e9d54726d09 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 22 Aug 2026 13:47:30 +0000 Subject: [PATCH] fix: message shape/size/ndim read the real JAX broadcast shape (#1510) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit jax 0.11 changed jnp.broadcast_arrays to return a tuple instead of a list (NumPy 2 alignment), so the isinstance(..., list) JAX branch in MessageInterface.shape stopped matching and fell through to .shape on a tuple — every jnp-backed message construction raised AttributeError. Rather than widening the isinstance to (list, tuple), which would preserve the () shape sentinel, shape now returns the real broadcast shape read off the container's first element, and size/ndim derive from it instead of attribute-accessing the container (they raised AttributeError for any JAX-backed message on jax 0.10 too). The sentinel was itself a live bug: _broadcast_natural_parameters matched the shape[1:] branch for batched JAX messages, inserted a spurious axis, and logpdf returned an (n, n) matrix of wrong values where NumPy returns the correct (n,) vector. The new parity test asserts NumPy/JAX equality of shape/size/ndim and of batched logpdf values and shape over Normal/Beta/Gamma — it fails 6 ways against the sentinel on jax 0.10 alone, so the () sentinel cannot come back silently. Verified: full suite 2030 passed on jax 0.10.2 and 0.11.1 (Python 3.12, [optional] extras installed); the ten autofit_workspace_test jax_assertions scripts unchanged (9/10 pass on both versions, the priors_xp_dispatch float32-tolerance failure is pre-existing and identical in every cell). Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01EauVX6vD9k1N4PXaEf2wvo --- autofit/messages/interface.py | 21 +++++++--- test_autofit/messages/test_jax_trace.py | 56 +++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 5 deletions(-) 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)