Skip to content

fix: skip the JAX-only sparse-operator tests where jax is unavailable - #658

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/potential-correction-jax-skip
Jul 27, 2026
Merged

fix: skip the JAX-only sparse-operator tests where jax is unavailable#658
Jammy2211 merged 1 commit into
mainfrom
feature/potential-correction-jax-skip

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Fixes the unit_tests (3.9, PyAutoLens) and (3.10, PyAutoLens) legs of PyAutoHands/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:

ModuleNotFoundError: No module named 'jax'
  ../PyAutoArray/autoarray/dataset/interferometer/dataset.py:280: in apply_sparse_operator
  ../PyAutoArray/.../inversion_interferometer_util.py:654: in from_nufft_precision_operator
      import jax.numpy as jnp

Why this is a test-placement problem, not a library bug

The sparse-operator subsystem is JAX-only by design:

  • InterferometerSparseOperator builds its FFT kernel with jax.numpy (from_nufft_precision_operator), applies it with jax.numpy (apply_operator), and projects with jax.ops.segment_sum + jax.lax (curvature_matrix_diag_from) — three jax sites, no NumPy equivalent.
  • The imaging counterpart has 8+ jax sites and types a field literally as "jax.Array".
  • autonerves[jax] gates jax to python_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_sumnp.add.at/bincount, dropping lax), a substantial new feature with real numerical-parity risk, to serve Pythons the feature is not gated for.

Change

A find_spec-based skipif marker applied to exactly those 5 cases, matching the pytest.importorskip idiom already used in test_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

jax present -> 9 passed, 0 skipped
jax absent  -> 4 passed, 5 skipped

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_matrix being red. The first — nufftax on the 3.11 autolens_workspace smoke leg — was already fixed on autolens_workspace main by #351 earlier today; a verification dispatch is in flight.

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>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 27, 2026
@Jammy2211
Jammy2211 merged commit 88eb46b into main Jul 27, 2026
4 checks passed
@Jammy2211
Jammy2211 deleted the feature/potential-correction-jax-skip branch July 27, 2026 19:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant