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

[Jammy2211]

Implements the fix specified in #1510 (merge this before the PyAutoNerves cap widen; the issue stays open until both legs land).

Root cause

jax 0.11 changed jnp.broadcast_arrays to return a tuple instead of a list (NumPy 2 alignment). MessageInterface.shape type-tested the container with isinstance(..., list), so on 0.11 the JAX branch stopped matching, fell through to .shape on a tuple, and every jnp-backed message construction raised AttributeError — the four test_jax_trace log-partition failures blocking the autonerves jax cap widen.

Why not the one-line isinstance widen

The JAX branch returned a hard-wired () shape, which is a live numerical bug on jax 0.10 today: for batched JAX messages, _broadcast_natural_parameters matched the shape[1:] == self.shape branch, inserted a spurious axis, and logpdf returned an (n, n) matrix of wrong values where NumPy returns the correct (n,) vector. size/ndim had no JAX branch at all and raised AttributeError for any jnp-backed message on both jax versions. Widening the isinstance would ship all of that unchanged.

Instead, shape now returns the real broadcast shape read off the container's first element (correct for both the list and tuple forms), and size/ndim derive from it.

Reviewer note: this changes numerical output for batched JAX-backed messages — from the incorrect broadcast-matrix result to the NumPy-matching vector. That is the intended correction. The size/ndim repair goes beyond the 0.11 regression (they were equally broken on 0.10) and is included deliberately: it is the same three-property cluster reading the same container, covered by the new test.

Regression test

test_message_shape_and_logpdf_match_numpy asserts NumPy/JAX parity of shape/size/ndim and of logpdf values and shape, scalar and batched, over Normal/Beta/Gamma. Mutation-checked: against the old () sentinel it fails 6 ways on jax 0.10 alone, so the sentinel cannot return silently.

Verification

Two Python 3.12 venvs (jax/jaxlib 0.10.2 and 0.11.1), [optional] extras installed (blackjax, nautilus-sampler — without them 18 tests silently skip):

tree	jax 0.10.2	jax 0.11.1
main, unfixed	2024 passed, 3 skipped	4 failed, 2020 passed, 3 skipped
this branch	2030 passed, 3 skipped	2030 passed, 3 skipped

• autofit_workspace_test/scripts/jax_assertions/: 9/10 pass in all four version-by-fix cells; only the pre-existing priors_xp_dispatch.py float32-tolerance failure, identical everywhere.
• Downstream, against this branch: test_autogalaxy 1103 passed and test_autolens 532 passed, identical under jax 0.10.2 and 0.11.1; test_autonerves 157 passed under 0.11.1.

🤖 Generated with Claude Code

https://claude.ai/code/session_01EauVX6vD9k1N4PXaEf2wvo

---

Generated by Claude Code