fix: make message array construction JAX-traceable

[Jammy2211]

Overview

Several autofit.messages methods construct fixed-shape arrays with xp.array([a, b]). NumPy accepts this, but JAX tracing can try to coerce traced scalar inputs through __array__, raising TracerArrayConversionError in hierarchical factor likelihoods.

This is a focused library-source follow-up to the closed JAX-native-priors work in #1262. It also overlaps the still-open priors/messages audit hub #1331, which should be checked for adjacent findings without expanding this repair into unrelated statistical changes.

High-level plan

• Add a direct JAX JIT regression test for the affected Normal, TruncatedNormal, Beta and Gamma message methods.
• Replace the ten fixed-shape xp.array([...]) constructions with xp.stack([...]).
• Replace np.log(x) with xp.log(x) in GammaMessage.to_canonical_form.
• Verify scalar and batched NumPy/JAX shape and numerical parity.
• Audit the prior/message xp boundary and existing prior-related issues; file unrelated findings separately.
• Run focused message tests, the full PyAutoFit suite, and the hierarchical integration reproducer with JAX enabled.

Detailed implementation plan

Affected repository

• PyAutoFit

Suggested branch

feature/messages-xp-stack-jax-trace

Source changes

Update these fixed-shape constructors:

• autofit/messages/normal.py: three sites
• autofit/messages/truncated_normal.py: three sites
• autofit/messages/beta.py: two sites
• autofit/messages/gamma.py: two sites

For GammaMessage.to_canonical_form, use xp.log(x) so a traced JAX value never enters NumPy.

Regression coverage

Add parameterized tests under test_autofit/messages/ that:

• execute every changed constructor through jax.jit;
• cover traced scalar inputs and representative batched inputs;
• compare JAX output shape and values against the NumPy path;
• skip cleanly when the optional JAX extra is unavailable;
• prove the original TracerArrayConversionError path is closed rather than merely testing eager JAX.

Prior/message audit

While touching the message layer:

• compare current autofit/mapper/prior/ and autofit/messages/ code against feat: JAX-native priors — xp dispatch on value_for / log_prior_from_value / vector_from_unit_vector #1262 and Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331;
• search for analogous backend leaks such as literal NumPy operations inside methods accepting xp;
• keep corrections in this PR only when they are the same traced-array root cause;
• record separate bugs or stale issue findings rather than broadening this PR into prior-statistics redesign.

Validation

NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib   python -m pytest test_autofit/messages -x

NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib   python -m pytest test_autofit -x

Then run the z_projects/concr toy hierarchical fit with JAX enabled and confirm it completes without TracerArrayConversionError while remaining consistent with the NumPy run.

CI must pass on Python 3.12 and 3.13.

Scope discipline

This is one PyAutoFit repair PR. Any wider autofit/ JAX audit or independent prior/message statistical bug becomes a separate tracked task. No workspace script should be changed to mask the library defect.

[Jammy2211]

Shipped in #1460 (squash commit be9c6c9204df87c987694313f460c37764d6963a).

Validation before merge:

• GitHub Docs: passed
• GitHub Tests: passed
• focused message suite: 32 passed
• full local PyAutoFit suite: 1689 passed, 4 skipped
• untouched-main regression proof: two expected Gamma canonical-form JIT failures; all 20 cases passed on the branch

The adjacent prior/message audit found four separate JAX backend leaks, tracked independently in #1459.