Skip to content

test: remove jax from unit tests (moved to autogalaxy_workspace_test)#384

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

test: remove jax from unit tests (moved to autogalaxy_workspace_test)#384
Jammy2211 merged 1 commit intomainfrom
feature/jax-assertions-sweep

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Two changes to keep test_autogalaxy/ numpy-only so the no-jax CI matrix stays green and local pytest runs don't pay the jax import cost.

Changes

  • test_autogalaxy/conftest.py: drop the pytest_configure hook that did jnp.sum(jnp.array([0.0])) to force JAX backend init at session start. This eagerly imported jax even when no jax tests would run.
  • test_autogalaxy/operate/test_deflections.py: delete the two np-vs-jnp parity tests (test__hessian_from__np_richardson_matches_jax_jacfwd_to_float64, test__magnification_2d_via_hessian_from__np_jnp_agree_to_float64). Both moved to scripts/jax_assertions/hessian_parity.py in autogalaxy_workspace_test. The other 27 tests in this file are pure numpy and stay.

No library code changes.

Merge dependency

Depends on autogalaxy_workspace_test#23 merging first, so the moved assertions exist on main before the source tests are removed.

Test plan

  • Local pytest with jax installed: 843 passed (was 845, exactly 2 removed)
  • CI on this PR

Two changes to keep test_autogalaxy/ numpy-only so the no-jax CI matrix
stays green and local pytest runs don't pay the jax import cost:

- conftest.py: drop the pytest_configure hook that did
  `jnp.sum(jnp.array([0.0]))` to force JAX backend init at session
  start. This eagerly imported jax even when no jax tests would run.
- test_deflections.py: delete the two np-vs-jnp parity tests
  (test__hessian_from__np_richardson_matches_jax_jacfwd_to_float64,
  test__magnification_2d_via_hessian_from__np_jnp_agree_to_float64).
  Both moved to scripts/jax_assertions/hessian_parity.py in
  autogalaxy_workspace_test. The other 27 tests in this file are pure
  numpy and stay.

No library code changes. Local pytest: 843 passed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 merged commit 732f934 into main May 1, 2026
5 checks passed
@Jammy2211 Jammy2211 deleted the feature/jax-assertions-sweep branch May 1, 2026 08:02
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