diff --git a/autoarray/util/fnnls.py b/autoarray/util/fnnls.py index c4985d85..01ea1704 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 91879391..e8c96220 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)