From b74452be8a11dcfb230da57bebcaee92720a253d Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 27 Jul 2026 20:24:32 +0100 Subject: [PATCH] fix: skip the JAX-only sparse-operator tests where jax is unavailable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../test_fit_interferometer.py | 14 ++++++++++++++ .../test_iterative_interferometer.py | 17 +++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/test_autolens/potential_correction/test_fit_interferometer.py b/test_autolens/potential_correction/test_fit_interferometer.py index b467ddecb..c74f45964 100644 --- a/test_autolens/potential_correction/test_fit_interferometer.py +++ b/test_autolens/potential_correction/test_fit_interferometer.py @@ -1,3 +1,5 @@ +import importlib.util + import numpy as np import pytest @@ -7,6 +9,17 @@ FitDpsiSrcInterferometer, ) +# `apply_sparse_operator` is a JAX-only code path: `InterferometerSparseOperator` +# builds its FFT kernel with `jax.numpy` and projects with `jax.ops.segment_sum` +# / `jax.lax`, with no NumPy equivalent. `autonerves[jax]` gates jax to +# Python >= 3.11, so the sparse-route case cannot run on the 3.9/3.10 matrix +# legs — skip it there rather than fail. The dense-route cases stay NumPy-only +# and run everywhere, which is what those legs exist to prove. +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, + reason="apply_sparse_operator is a JAX-only path; jax requires Python >= 3.11", +) + def fit_from(dataset, use_sparse_operator): lens = al.Galaxy( @@ -50,6 +63,7 @@ def test__dense_route__end_to_end_evidence_is_finite(interferometer_7): ) +@requires_jax def test__sparse_route__matches_dense_route(interferometer_7): dataset_sparse = interferometer_7.apply_sparse_operator() diff --git a/test_autolens/potential_correction/test_iterative_interferometer.py b/test_autolens/potential_correction/test_iterative_interferometer.py index 1a89a81a3..d15683f08 100644 --- a/test_autolens/potential_correction/test_iterative_interferometer.py +++ b/test_autolens/potential_correction/test_iterative_interferometer.py @@ -1,9 +1,22 @@ +import importlib.util + import numpy as np import pytest import autoarray as aa import autolens as al +# `apply_sparse_operator` is a JAX-only code path: `InterferometerSparseOperator` +# builds its FFT kernel with `jax.numpy` and projects with `jax.ops.segment_sum` +# / `jax.lax`, with no NumPy equivalent. `autonerves[jax]` gates jax to +# Python >= 3.11, so the cases below cannot run on the 3.9/3.10 matrix legs — +# skip them there rather than fail. The dense-route cases stay NumPy-only and +# run everywhere, which is what those legs exist to prove. +requires_jax = pytest.mark.skipif( + importlib.util.find_spec("jax") is None, + reason="apply_sparse_operator is a JAX-only path; jax requires Python >= 3.11", +) + def iter_fit_from(dataset, gauge_constraints=False, n_iter=2): lens = al.Galaxy( @@ -33,6 +46,7 @@ def test__requires_sparse_operator(interferometer_7): iter_fit_from(interferometer_7) +@requires_jax def test__solve_joint_optimization__finite_state_and_decreasing_cost( interferometer_7, ): @@ -56,6 +70,7 @@ def test__solve_joint_optimization__finite_state_and_decreasing_cost( assert cost_opt < 0.5 * fit.data_weighted_norm +@requires_jax def test__cost_identity_matches_direct_visibility_chi2(interferometer_7): """ The normal-equation chi^2 identity (d^H C^-1 d - 2 x^T D + x^T F x) must @@ -88,6 +103,7 @@ def test__cost_identity_matches_direct_visibility_chi2(interferometer_7): assert chi2_half == pytest.approx(chi2_direct, rel=1e-3) +@requires_jax def test__gauge_constraints_are_satisfied(interferometer_7): dataset = interferometer_7.apply_sparse_operator() fit = iter_fit_from(dataset, gauge_constraints=True) @@ -103,6 +119,7 @@ def test__gauge_constraints_are_satisfied(interferometer_7): assert G @ dpsi_opt == pytest.approx(np.zeros(3), abs=1.0e-6) +@requires_jax def test__log_evidence__finite_at_optimum(interferometer_7): dataset = interferometer_7.apply_sparse_operator() fit = iter_fit_from(dataset)