From fc00636a6cb490ed9e1882f1c77d8fdfa3826280 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 21 Aug 2026 20:04:44 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20coerce=20fnnls=5Fcholesky=20inputs=20to?= =?UTF-8?q?=20numpy=20=E2=80=94=20JAX=20ZTx=20broke=20the=20numba=20buffer?= =?UTF-8?q?=20kernels?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sparse-operator inversion path hands fnnls_cholesky JAX arrays even when the fit itself runs the numba CPU path. Indexing a JAX array yields another JAX array, which numba maps to a *readonly* buffer — and the in-place buffer kernels introduced by #453 (_cho_solve_buffer via _solve_upper_transposed_buffer) overwrite their vector argument, so kernel compilation fails with "Cannot modify readonly array" the first time the solver runs on such input. The scipy solvers #453 replaced tolerated JAX input by copying internally, which is why this never bit before. Visible failure: HowToLens smoke on main, red since 2026-08-20 22:30 UTC — tutorial_8_adaptive_pixelization and tutorial_11_brightness_adaption (the two tutorials whose fits route through the sparse operator) died with "During: Pass nopython_type_inference" on both CI Pythons, 21 minutes after the #453 merge. This RED is what blocked the 2026-08-21 nightly release. Coerce ZTZ / ZTx / P_initial once at the function boundary: fancy indexing a numpy parent hands every downstream kernel a fresh writeable array. The new regression test calls fnnls_cholesky with jnp arrays (importorskip'd) and is verified to fail without the coercion. test_autoarray: 1063 passed; the 3 test_transformer pynufft failures reproduce identically on unmodified main in this environment (missing [optional] extras) and are unrelated. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_015GsUfbCwPd4XC8kpsiUJp7 --- autoarray/util/fnnls.py | 14 +++++++++++ test_autoarray/util/test_cholesky_inplace.py | 26 ++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/autoarray/util/fnnls.py b/autoarray/util/fnnls.py index c4985d852..01ea17047 100644 --- a/autoarray/util/fnnls.py +++ b/autoarray/util/fnnls.py @@ -32,6 +32,20 @@ def fnnls_cholesky( """ from scipy import linalg as slg + # The buffer kernels below (_cho_solve_buffer / cholinsertlast_inplace) + # overwrite their vector argument in place, so every slice handed to them + # must be a writeable numpy array. A JAX ZTZ / ZTx — the sparse-operator + # inversion path hands one over even when the fit itself runs the numba + # CPU path — breaks that contract: indexing a JAX array yields another + # JAX array, which numba maps to a *readonly* buffer and rejects at + # compile time ("Cannot modify readonly array"). Coerce once at the + # boundary — fancy indexing a numpy parent then hands the kernels fresh + # writeable arrays, exactly as the scipy solvers this replaced tolerated + # by copying internally. + ZTZ = np.asarray(ZTZ) + ZTx = np.asarray(ZTx) + P_initial = np.asarray(P_initial) + lstsq = lambda A, x: slg.solve( A, x, diff --git a/test_autoarray/util/test_cholesky_inplace.py b/test_autoarray/util/test_cholesky_inplace.py index 91879391a..e8c962207 100644 --- a/test_autoarray/util/test_cholesky_inplace.py +++ b/test_autoarray/util/test_cholesky_inplace.py @@ -167,3 +167,29 @@ def test__fnnls_cholesky__warm_start_matches_cold_start(seed): d_warm = fnnls_cholesky(ZTZ, ZTx, P_initial=P_initial) assert d_warm == pytest.approx(d_cold, rel=1e-8, abs=1e-10) + + +@pytest.mark.parametrize("seed", [0, 1]) +def test__fnnls_cholesky__accepts_jax_arrays(seed): + """ + The sparse-operator inversion path hands fnnls_cholesky JAX arrays even + when the fit runs the numba CPU path. Indexing a JAX array yields another + JAX array, which numba maps to a readonly buffer — before the boundary + coercion in fnnls_cholesky this failed kernel compilation with + "Cannot modify readonly array" (HowToLens smoke, 2026-08-20). + """ + jnp = pytest.importorskip("jax.numpy") + + rng = np.random.default_rng(seed) + n = 30 + Z = rng.normal(size=(50, n)) + x = Z @ rng.normal(size=n) + rng.normal(size=50) + + ZTZ = Z.T @ Z + ZTx = Z.T @ x + + d_np = fnnls_cholesky(ZTZ, ZTx) + d_jax = fnnls_cholesky(jnp.asarray(ZTZ), jnp.asarray(ZTx)) + + assert np.all(np.asarray(d_jax) >= 0.0) + assert np.asarray(d_jax) == pytest.approx(d_np, rel=1e-6, abs=1e-8)