Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Jul 27, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions test_autolens/potential_correction/test_fit_interferometer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import importlib.util

import numpy as np
import pytest

Expand All @@ -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(
Expand Down Expand Up @@ -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()

Expand Down
17 changes: 17 additions & 0 deletions test_autolens/potential_correction/test_iterative_interferometer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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(
Expand Down Expand Up @@ -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,
):
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
Loading