fix: skip the JAX-only sparse-operator tests where jax is unavailable - #658
Merged
Conversation
unit_tests (3.9, PyAutoLens) and (3.10, PyAutoLens) have been red in
PyAutoHands/python_matrix with 5 failures, all:
ModuleNotFoundError: No module named 'jax'
.../inversion_interferometer_util.py:654: in from_nufft_precision_operator
import jax.numpy as jnp
reached via Interferometer.apply_sparse_operator (dataset.py:280). The whole
sparse-operator subsystem is JAX-only by design — InterferometerSparseOperator
builds its FFT kernel with jax.numpy and projects with jax.ops.segment_sum /
jax.lax, and the imaging counterpart even types a field as 'jax.Array'. There
is no NumPy equivalent, and autonerves[jax] gates jax to Python >= 3.11.
So these are JAX-feature tests sitting in matrix legs that have no jax: a test
placement problem, not a library bug. Marking exactly the 5 cases that call
apply_sparse_operator() with a find_spec-based skipif, matching the
pytest.importorskip idiom already used in test_autolens/interop/test_coolest.py.
This also restores the standing 'library unit tests are numpy-only' rule for
these files: the dense-route cases stay NumPy-only and keep running on 3.9/3.10,
which is exactly what those legs exist to prove.
Verified both ways:
jax present -> 9 passed, 0 skipped
jax absent -> 4 passed, 5 skipped (the 5 CI failures, and only those)
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
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.
Fixes the
unit_tests (3.9, PyAutoLens)and(3.10, PyAutoLens)legs ofPyAutoHands/python_matrix, red on the schedule (pre-existing — present in the 2026-07-20 and 2026-07-27 runs).The failure
5 failures, all the same:
Why this is a test-placement problem, not a library bug
The sparse-operator subsystem is JAX-only by design:
InterferometerSparseOperatorbuilds its FFT kernel withjax.numpy(from_nufft_precision_operator), applies it withjax.numpy(apply_operator), and projects withjax.ops.segment_sum+jax.lax(curvature_matrix_diag_from) — three jax sites, no NumPy equivalent."jax.Array".autonerves[jax]gates jax topython_version >= '3.11', so 3.9/3.10 have no jax by construction.The 5 failing cases are exactly those calling
apply_sparse_operator()— i.e. tests of the JAX feature, running in matrix legs that deliberately have no jax.A NumPy fallback was considered and rejected: it would mean reimplementing the whole sparse-operator path (
segment_sum→np.add.at/bincount, droppinglax), a substantial new feature with real numerical-parity risk, to serve Pythons the feature is not gated for.Change
A
find_spec-basedskipifmarker applied to exactly those 5 cases, matching thepytest.importorskipidiom already used intest_autolens/interop/test_coolest.py.Deliberately per-test, not module-level: the dense-route cases in the same files (
test__dense_route__end_to_end_evidence_is_finite,test__sparse_route__requires_sparse_operator,test__evidence_terms__match_hand_computed_dense_formulation,test__requires_sparse_operator) are NumPy-only and must keep running on 3.9/3.10 — that is what those legs exist to prove. A module-level skip would have silently dropped them.This also restores the standing "library unit tests are NumPy-only" rule for these files.
Verification
The absent case was produced by shimming
importlib.util.find_spec("jax") -> None, and the 5 skips are exactly the 5 CI failures — no more, no less.Scope note
This is the second cause of
python_matrixbeing red. The first —nufftaxon the 3.11autolens_workspacesmoke leg — was already fixed onautolens_workspacemain by #351 earlier today; a verification dispatch is in flight.