Skip to content

test: remove jax-using unit tests; assertions moved to autolens_workspace_test#295

Merged
Jammy2211 merged 1 commit intomainfrom
feature/jax-assertions-move
May 1, 2026
Merged

test: remove jax-using unit tests; assertions moved to autolens_workspace_test#295
Jammy2211 merged 1 commit intomainfrom
feature/jax-assertions-move

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

  • Remove all 9 jax-using unit tests from test_autoarray/ (7 functions across 5 files + 2 from the deleted test_jax_pytree.py).
  • Library code unchanged.
  • Conftest's collect_ignore_glob for test_jax_*.py removed since no such files remain.
  • Renamed test_jax_changes.pytest_changes.py (it never used jax).

Why

Two problems with the prior layout:

  1. No-jax CI matrix breaks on Python 3.9 / 3.10 — tests like test__data_vector_via_weighted_data_two_methods_agree, test__inversion_imaging__via_mapper_knn, etc. trigger eager import jax inside library helpers (ImagingSparseOperator, data_vector_via_psf_weighted_data_from, KNN interpolator) and crash with ModuleNotFoundError: No module named 'jax'.
  2. Local dev runs pay the jax import cost (~1-3s cold) on every pytest invocation even when those tests would self-skip via pytest.importorskip.

Per the project rule "library unit tests stay numpy-only; cross-xp checks live in workspace_test", the assertions are relocated as runnable scripts in autolens_workspace_test#PR.

Tests removed → destination

Removed test New script
test_inversion_imaging_util.py::test__data_vector_via_weighted_data_two_methods_agree scripts/jax_assertions/sparse_operators.py
test_inversion_imaging_util.py::test__curvature_matrix_via_psf_weighted_noise_two_methods_agree scripts/jax_assertions/sparse_operators.py
test_inversion_interferometer_util.py::test__curvature_matrix_via_psf_precision_operator_from scripts/jax_assertions/sparse_operators.py
test_factory.py::test__inversion_imaging__via_mapper_knn scripts/jax_assertions/knn_interpolator.py
test_inversion_util.py::test__reconstruction_positive_only_from__jax_ill_conditioned_grad_is_finite scripts/jax_assertions/nnls.py
test_inversion_util.py::test__reconstruction_positive_only_from__jax_matches_unpreconditioned_primal scripts/jax_assertions/nnls.py
test_jax_pytree.py (3 tests) scripts/jax_assertions/pytree.py
test_irregular_2d.py jax half of test__grid_2d_via_deflection_grid_from__propagates_xp scripts/jax_assertions/grid_irregular.py

Test plan

  • Local pytest with jax installed: 739 passed (was 748 — exactly 9 removed)
  • Local pytest in no-jax venv (numba+pynufft installed): zero jax-related failures (the 31 pinv2 failures are unrelated pynufft↔scipy incompat, not in CI)
  • All 5 new scripts execute successfully on the autolens_workspace_test side
  • PyAutoBuild python_matrix.yml matrix passes on 3.9 / 3.10 / 3.11 / 3.12 / 3.13 after merge

…pace_test

Unit tests in test_autoarray/ that imported jax (directly via importorskip
or indirectly via library helpers like ImagingSparseOperator and the KNN
interpolator) caused two problems: they broke the no-jax CI matrix on
Python 3.9 / 3.10, and they paid the jax import cost on every local run
even when self-skipping. Per the project rule "library unit tests stay
numpy-only; cross-xp checks live in workspace_test", move all 9 jax
assertions into autolens_workspace_test/scripts/jax_assertions/.

Library code unchanged. Conftest's collect_ignore_glob for test_jax_*.py
removed since no such files remain (test_jax_changes.py renamed to
test_changes.py since it never used jax in the first place).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 merged commit 25a988d into main May 1, 2026
4 checks passed
@Jammy2211 Jammy2211 deleted the feature/jax-assertions-move branch May 1, 2026 07:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant