fix: message shape/size/ndim break on jax 0.11 (broadcast_arrays now returns a tuple)

[Jammy2211]

Overview

jax 0.11 changed jnp.broadcast_arrays to return a tuple instead of a list (aligning with NumPy 2, which made the same change to np.broadcast_arrays). MessageInterface.shape type-tests that container with isinstance(..., list), so on jax 0.11 the JAX branch stops matching, control falls through to self.broadcast.shape, and every BetaMessage/GammaMessage construction under jax.jit dies with AttributeError: 'tuple' object has no attribute 'shape'.

This is what blocks widening the jax/jaxlib cap in autonerves from <0.11.0 to <0.12.0 — the widen was reverted in PyAutoNerves#150 (848a254) with a comment pointing at this bug. The cap matters because jax is now a base dependency (PyAutoLens#702) and <0.11 conflicts with e.g. Colab's preinstalled jax.

An adversarial review of the first plan changed the fix. The obvious one-line repair (widen the isinstance to accept a tuple) makes the four tests pass while preserving a silent numerical bug that is live on jax 0.10 today. See "Why the minimal fix is wrong" below. This issue now specifies the semantic fix.

Reproduction

Under jax 0.11.1, pytest test_autofit gives exactly four failures, all the same AttributeError:

test_autofit/messages/test_jax_trace.py::test_message_log_partition_is_jittable_and_matches_numpy[scalar-gamma]
test_autofit/messages/test_jax_trace.py::test_message_log_partition_is_jittable_and_matches_numpy[scalar-beta]
test_autofit/messages/test_jax_trace.py::test_message_log_partition_is_jittable_and_matches_numpy[batched-gamma]
test_autofit/messages/test_jax_trace.py::test_message_log_partition_is_jittable_and_matches_numpy[batched-beta]

Traceback tail:

autofit/messages/beta.py:156: in __init__
    super().__init__(
autofit/messages/abstract.py:65: in __init__
    if self.shape:
autofit/messages/interface.py:31: in shape
    return self.broadcast.shape
AttributeError: 'tuple' object has no attribute 'shape'

test_autofit/graphical/functionality/test_messages.py::test_beta — the fifth failure originally reported — now passes: the np.generic xp-dispatch fix (PyAutoFit 19c679583) already cleared it. The compat surface is these four tests only.

Verified container-type change:

jax 0.10.2 → type(jnp.broadcast_arrays(a, b)) is list
jax 0.11.1 → type(jnp.broadcast_arrays(a, b)) is tuple

jax.scipy.special.betaln / gammaln are not implicated — a standalone repro of the same maths jits cleanly on 0.11.1.

Why the minimal fix is wrong

shape returns () for the JAX branch unconditionally, including batched parameters. That falsehood propagates: MessageInterface._broadcast_natural_parameters compares xp.shape(x) against self.shape, misses the equal-shape branch, matches shape[1:] == self.shape instead, and inserts a spurious axis. On jax 0.10.2, current main, no jax 0.11 involved:

mnp = NormalMessage(np.array([1.0, 2.0]),  np.array([1.0, 1.0]))
mj  = NormalMessage(jnp.array([1.0, 2.0]), jnp.array([1.0, 1.0]))
x   = np.array([0.5, 1.5])

mnp.logpdf(x, xp=np)              # [-1.04393853 -1.04393853]      shape (2,)   correct
mj.logpdf(jnp.asarray(x), xp=jnp) # [[-1.1689 -0.6689]
                                  #  [-1.2939  0.4561]]            shape (2, 2) WRONG

size and ndim are worse: they call .size / .ndim on the broadcast container with no JAX branch at all, so they raise AttributeError for any JAX-backed message on 0.10 as well as 0.11.

Widening the isinstance to (list, tuple) turns the four tests green and ships all of that unchanged. Returning the real broadcast shape instead fixes the regression and the silent wrongness, and passes the same suites.

Plan

• Replace the () sentinel in MessageInterface.shape with the actual broadcast shape, read off the first element of the JAX broadcast container — correct for both the list (jax <= 0.10) and tuple (jax >= 0.11) forms.
• Derive size and ndim from that shape for the JAX branch instead of attribute-accessing the container.
• Add a regression test asserting NumPy/JAX parity of .shape / .size / .ndim and of batched logpdf output — the parity assertion is what stops the () sentinel coming back.
• Once this merges, widen the jax / jaxlib cap to <0.12.0 in PyAutoNerves and drop the "Cap stays <0.11" comment (library-first: PyAutoNerves must not widen ahead of this fix).

Detailed implementation plan

Work Classification

Library (no workspace change). Two library repos, merged in order.

Affected Repositories

• PyAutoLabs/PyAutoFit (primary) — the shape / size / ndim fix
• PyAutoLabs/PyAutoNerves — the dependency cap widen, merged after PyAutoFit

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./PyAutoNerves	main	clean

Suggested branch: feature/message-log-partition-tuple-shape
Worktree root: ~/Code/PyAutoLabs-wt/message-log-partition-tuple-shape/ (created later by /start_library)

worktree_check_conflict message-log-partition-tuple-shape PyAutoFit PyAutoNerves → exit 0 (no conflict).

Implementation Steps

1. PyAutoFit/autofit/messages/interface.py — shape returns the real broadcast shape for the JAX branch (np.shape(broadcast[0]), () for an empty container), replacing the () sentinel. Replace the stale # JAX behaviour comment with the actual reason: jnp.broadcast_arrays returns a list on jax <= 0.10 and a tuple on jax >= 0.11, mirroring the NumPy 2 change, so the container is matched on (list, tuple) and never on identity with one of them.
2. PyAutoFit/autofit/messages/interface.py — size and ndim take the same container branch and derive from self.shape (prod(shape) and len(shape)) rather than .size / .ndim on the container.
3. PyAutoFit/test_autofit/messages/test_jax_trace.py — add a NumPy/JAX parity case over BetaMessage / GammaMessage / NormalMessage asserting equal .shape, .size, .ndim and equal batched logpdf values and shape. The existing test_message_log_partition_is_jittable_and_matches_numpy already covers the four regressions.
4. Verification must span the whole JAX surface, not just test_autofit:
   • pytest test_autofit with the [optional] extras installed (blackjax and nautilus-sampler in particular — without them 18 tests skip, 13 of which are the blackjax JAX samplers, the heaviest JAX consumers in the repo).
   • the ten scripts in autofit_workspace_test/scripts/jax_assertions/ (JAX unit coverage was moved out of test_autofit/ in test: delete jax-using unit tests (moved to autofit_workspace_test) #1247, so the library suite alone is not the JAX gate).
   • both of the above under jax 0.10.2 and jax 0.11.1, requiring identical results.
5. PyAutoNerves/pyproject.toml — after the PyAutoFit PR merges, set jax>=0.7.0,<0.12.0 and jaxlib>=0.7.0,<0.12.0, delete the three-line "Cap stays <0.11" comment, and leave the Intel-macOS environment markers untouched.

Key Files

• autofit/messages/interface.py — MessageInterface.shape / .size / .ndim; the defect site (lines 24-39).
• autofit/messages/abstract.py — lines 61-65; the np.broadcast vs xp.broadcast_arrays split that produces the container, and the if self.shape: that raises.
• test_autofit/messages/test_jax_trace.py — the four failing parametrisations.
• autofit_workspace_test/scripts/jax_assertions/ — the out-of-repo JAX assertion surface.
• PyAutoNerves/pyproject.toml — lines 35-39; the cap and its comment.

Evidence gathered while planning

Two Python 3.12 virtualenvs pinned to jax 0.10.2 and 0.11.1, [optional] extras installed (blackjax 1.6.2, nautilus-sampler 1.0.5, getdist, astropy). The candidate fix was applied and reverted in a scratch clone; the clone is clean.

pytest test_autofit:

tree	jax 0.10.2	jax 0.11.1
main, unfixed	2024 passed, 3 skipped	4 failed, 2020 passed, 3 skipped
candidate fix	2024 passed, 3 skipped	2024 passed, 3 skipped

autofit_workspace_test/scripts/jax_assertions/ — 8 of 10 pass on every combination. The two that fail (priors_xp_dispatch.py, multi_start_gradient_auto_convergence.py) fail identically in all four version-by-fix cells, so they are pre-existing and unrelated to this change; they are noted separately below.

Batched JAX logpdf parity with NumPy: fails on main (returns a (2, 2) matrix), passes with the fix, on both jax versions.

Blast-radius audit for the cap widen

• Full public-API diff, jax 0.10.2 → 0.11.1 across jax, jax.numpy, jax.scipy.special, jax.scipy.linalg, jax.numpy.linalg, jax.lax, jax.tree_util, jax.nn, jax.random, jax.scipy.stats: the only removal is jax.lax.dce_sink_p (unused in this stack); everything else is additive.
• Container return-type sweep over 22 jnp/jax entry points: two changed list → tuple — broadcast_arrays (this bug) and meshgrid. All nine meshgrid call sites across PyAutoFit and PyAutoArray tuple-unpack the result, so they are unaffected — but meshgrid is a second instance of the same class of change, and any new isinstance/.append handling of a jnp container is a latent repeat of this bug.
• PyAutoArray: pytest test_autoarray under both jax versions gives identical results (1084 passed, 54 skipped, 11 failed — the 11 are pre-existing missing-dependency failures in test_transformer.py / test_factory.py, unchanged between versions). PyAutoArray is jax-0.11-clean.
• jaxnnls==1.0.1 (hard-pinned alongside the cap): imports and solves under jax 0.11.1, bit-identical output to 0.10.2.
• Python floor: PyAutoFit, PyAutoNerves and PyAutoArray all declare requires-python = ">=3.12", and jax 0.11 requires >=3.12. There is therefore no cohort of users who stay on jax 0.10 after the cap widens — every supported environment resolves to 0.11. The widen is an all-users change, which is the reason for the ecosystem checks above rather than autofit-only evidence.
• Not yet verified: PyAutoGalaxy and PyAutoLens under jax 0.11. Their suites should be green before the PyAutoNerves cap widen merges.

Adjacent findings (out of scope, worth their own prompts)

1. autofit_workspace_test/scripts/jax_assertions/priors_xp_dispatch.py fails on library main under both jax versions — a float32/float64 tolerance mismatch (max relative difference 1.5e-7 against rtol=1e-7).
2. autofit_workspace_test/scripts/jax_assertions/multi_start_gradient_auto_convergence.py fails on library main under both jax versions — recovers normalization = 12.32 against an assertion of 25.0 +/- 3.0, after auto-convergence stops at 1 of 300 steps.
3. NormalMessage(1.0, jnp.array([1.0, 2.0])) dispatches to the NumPy backend: __init__ selects xp from the first parameter's type only, so a leading Python float wins over a trailing JAX array.

Notes on scope

• Prior art: fix: make message array construction JAX-traceable #1458 / Use xp.stack in messages so JAX can trace fixed arrays #1460 (xp.stack message constructors) and fix: make remaining message and compound-prior xp paths JAX-traceable #1459 / Make remaining message and prior xp paths JAX-traceable #1461 (Beta/Gamma log_partition xp dispatch, which introduced the current code). This bug is a container-type regression in the shared shape property, not a backend-dispatch leak.
• 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, not a side effect.

Autonomy

Prompt declares Autonomy: supervised. This run was launched without --auto, so every checkpoint presents and waits; merge stays a human act.

Original Prompt

Click to expand starting prompt

jax 0.11 breaks beta/gamma message log_partition under jit ('tuple' object has no attribute 'shape')

Type: bug
Target: PyAutoFit
Repos:

• PyAutoFit
• PyAutoNerves
  Difficulty: small
  Autonomy: supervised
  Priority: medium
  Status: draft

Found during the JAX-default-dependency arc (PyAutoLens#702): widening the jax
cap in autonerves from <0.11.0 to <0.12.0 let CI resolve jax/jaxlib 0.11.1,
which fails five autofit tests on both Python legs
(run 32285606183; local jax 0.10.2 passes):

• test_autofit/graphical/functionality/test_messages.py::test_beta
• test_autofit/messages/test_jax_trace.py::test_message_log_partition_is_jittable_and_matches_numpy[{scalar,batched}-{gamma,beta}]

all with AttributeError: 'tuple' object has no attribute 'shape'.

The cap widen was reverted in PyAutoNerves#150 (commit 848a254) with a comment
pointing here — the promotion shipped with the cap still <0.11.0.

Task: find what jax 0.11 changed in the beta/gamma log_partition trace path
(likely a jax.scipy.special return-shape/tuple change or a shape-polymorphism
change under jit), fix autofit's message code to be compatible with both 0.10
and 0.11, then widen the autonerves cap to <0.12.0 in the same arc
(@PyAutoNerves pyproject — remove the "Cap stays <0.11" comment). The cap
widen matters because jax is now a base dependency and the <0.11 cap
conflicts with e.g. Colab's preinstalled jax.

Note (2026-08-19, later same day): the no-jax CI leg exposed that
beta/gamma/normal message xp dispatch misrouted NumPy scalars
(np.int64/np.float64 are not int/float under NumPy 2) into the JAX branch —
fixed on the same branch (PyAutoFit 19c6795, np.generic added). test_beta
was one of the five jax-0.11 failures, so re-test under 0.11 AFTER that fix
lands: the remaining failures are probably only the deliberate jax-trace
tests (test_jax_trace.py), which narrows the compat surface.

[Jammy2211]

Implementation pushed after an independent adversarial review of the plan (all core claims reproduced from scratch: the list→tuple container change, the exact 4-test failure matrix, the silent (2,2) batched-logpdf bug on jax 0.10, and the size/ndim AttributeErrors).

Branches (same name in both repos, per the merge-order plan):

• PyAutoFit@feature/message-log-partition-tuple-shape — the semantic shape/size/ndim fix in autofit/messages/interface.py, plus a NumPy/JAX parity regression test in test_jax_trace.py. Mutation-checked: against the old () sentinel the new test fails 6 ways on jax 0.10 alone, so the sentinel cannot come back silently.
• PyAutoNerves@feature/message-log-partition-tuple-shape — the jax/jaxlib cap widen to <0.12.0 with the stale "Cap stays <0.11" comment replaced by a note that jax 0.11 raises the effective floors to numpy>=2.1 / scipy>=1.15. Merges only after the PyAutoFit PR, and only once PyAutoGalaxy/PyAutoLens are verified under 0.11 (see below).

Verification (two Python 3.12 venvs, jax/jaxlib 0.10.2 and 0.11.1, [optional] extras installed):

tree	jax 0.10.2	jax 0.11.1
main, unfixed	2024 passed, 3 skipped	4 failed, 2020 passed, 3 skipped
fix + new test	2030 passed, 3 skipped	2030 passed, 3 skipped

jax_assertions scripts: 9/10 pass in all four version-by-fix cells; only priors_xp_dispatch.py fails, identically everywhere (pre-existing float32 tolerance, unrelated). Note: multi_start_gradient_auto_convergence.py passes in all four cells in this environment — its reported failure is environment-sensitive, not a stable main fact; the adjacent-finding prompt should say so. test_autonerves also passes under 0.11.1 (157 passed).

Two review corrections to the plan, both acted on:

1. The breakage surface is broader than "Beta/Gamma under jit": any jnp-backed message construction (Normal included, scalar or batched, jit or not) raised on 0.11. The fix covers all of it; the parity test now guards Normal too.
2. "PyAutoGalaxy/PyAutoLens CI green first" cannot gate the cap widen: their CI (PyAutoHeart lib-tests.yml) resolves jax from autonerves' cap, so pre-widen CI runs 0.10 and proves nothing about 0.11. The workable gate is the reusable workflow's matching-branch checkout — push feature/message-log-partition-tuple-shape branches in PyAutoGalaxy/PyAutoLens so their CI resolves the widened cap from the Nerves branch — or forced pip install --no-deps jax==0.11.1 runs, as already done for PyAutoFit/PyAutoArray. The blast-radius audit also missed three 0.11 changes only the changelog reveals (jnp.empty now uninitialized, take_along_axis wrap_negative_indices default flip, ogrid container change) — all verified harmless for this stack, but future cap-widen audits should include a changelog pass.

---

Generated by Claude Code

[Jammy2211]

All three follow-ups are done; the task is now awaiting human merge.

• Downstream gate verified: test_autogalaxy (1103 passed, 1 skipped) and test_autolens (532 passed, 1 skipped) are identical under jax 0.10.2 and 0.11.1, run as forced jax==0.11.1 installs against the fix branch (their CI cannot pre-verify this — it resolves jax from the autonerves cap). test_autonerves: 157 passed under 0.11.1.
• PRs open: fix: message shape/size/ndim read the real JAX broadcast shape (jax 0.11 compat) #1513 (the fix — merges first) and build: widen jax/jaxlib cap to &lt;0.12.0 (PyAutoFit#1510) PyAutoNerves#152 (the cap widen — gated on fix: message shape/size/ndim read the real JAX broadcast shape (jax 0.11 compat) #1513, gate table on the PR).
• Mind ledger updated on claude/message-log-partition-tuple-shape-p7uylt: branch claims recorded, status → PRs open.

---

Generated by Claude Code

[Jammy2211]

Shipped. #1513 merged first (CI green on both Python legs, running under jax 0.11.1 via the matching-branch cap resolution), then PyAutoLabs/PyAutoNerves#152 (cap now <0.12.0). Mind completion record: complete/2026/08/message-log-partition-tuple-shape.md on the task branch. Closing as completed.

---

Generated by Claude Code