Skip to content

Feature/jax numba mix#187

Merged
Jammy2211 merged 17 commits intofeature/jax_wrapperfrom
feature/jax_numba_mix
Oct 6, 2025
Merged

Feature/jax numba mix#187
Jammy2211 merged 17 commits intofeature/jax_wrapperfrom
feature/jax_numba_mix

Conversation

@Jammy2211
Copy link
Owner

Sections off all functionality which is fast in numba but has proven difficult to translate to JAX efficiently.

This is primarily:

  • The w_tilde formalism for inversion (imaging and interferometry).
  • Delaunay and Voronoi mesh calculations.

The mapping formalism for inversions will next be updated to use JAX FFT convolutions and become the standard pixelization method for JAX. However, this does not exploit sparsity so for CPU hardware is slower than using numba.

After this PR PyAutoLens will therefore offer dual support for JAX and numba, albeit the latter will only be used for lens modeling using pixelizations efficiently on CPU.

@Jammy2211 Jammy2211 requested a review from CKrawczyk October 1, 2025 18:11
@Jammy2211 Jammy2211 changed the base branch from main to feature/jax_wrapper October 1, 2025 18:13
@Jammy2211 Jammy2211 merged commit 06ba8e3 into feature/jax_wrapper Oct 6, 2025
0 of 8 checks passed
@Jammy2211 Jammy2211 deleted the feature/jax_numba_mix branch November 30, 2025 16:31
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