fix: coerce fnnls_cholesky inputs to numpy — JAX ZTx broke the numba buffer kernels - #463
Merged
Conversation
…buffer kernels 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015GsUfbCwPd4XC8kpsiUJp7
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
fnnls_choleskynow coercesZTZ/ZTx/P_initialto numpy once at the function boundary, with a regression test that feeds itjax.numpyarrays.Why
The sparse-operator inversion path hands
fnnls_choleskyJAX arrays even when the fit itself runs the numba CPU path (PYAUTO_DISABLE_JAX=1). 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→_solve_upper_transposed_buffer) overwrite their vector argument, so kernel compilation fails withCannot modify readonly array of type: readonly buffer(float64, 1d, C)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_pixelizationandtutorial_11_brightness_adaption(the two tutorials whose fits route through the sparse operator) failed withDuring: Pass nopython_type_inferenceon both CI Pythons, 21 minutes after #453 merged (failing run). This is theHowToLens: Smoke Tests failure on mainRED that blocked the 2026-08-21 nightly release. Timeline pinned it to #453: no other library merge landed in the pass→fail window and both runs resolved identical dependency versions (numba 0.67.0 / numpy 2.5.2).With numpy parents, fancy indexing hands every downstream kernel a fresh writeable array, so the in-place contract holds for all callers.
Validation
test__fnnls_cholesky__accepts_jax_arrays(importorskip'd on jax) reproduces the failure without the coercion and passes with it; result cross-checked against the numpy-input solve.test_autoarray/: 1063 passed. (3test_transformer.pypynufft failures reproduce identically on unmodifiedmainin the validation env — missing[optional]extras — and are unrelated.)tutorial_8_adaptive_pixelization+tutorial_11_brightness_adaptionpass locally against current library mains with this fix (plus the HowToLens-side mesh-rename API update, PR'd separately).🤖 Generated with Claude Code
https://claude.ai/code/session_015GsUfbCwPd4XC8kpsiUJp7
Generated by Claude Code