JAX improvements + Fast CPU Pixelizations support + Delaunay
This release continues to build stability for JAX + GPU support:
https://github.com/Jammy2211/PyAutoLens/releases/tag/2025.11.18.1
This includes many fixes to small errors and bugs, for example more light and mass profiles support JAX after small fixes to the source code.
Fast CPU Pixelizations
Before this release, pixelized source reconstructions could only be computed using JAX, either via GPU or CPU.
There were two important factors in run time:
1. GPU VRAM Limitations
JAX only provides significant acceleration on GPUs with large VRAM (≥16 GB).
To avoid excessive VRAM usage, examples often restrict pixelization meshes (e.g. 20 × 20).
On consumer GPUs with limited memory, JAX may be slower than CPU execution.
2. Sparse Matrix Performance
Pixelized source reconstructions require operations on very large, highly sparse matrices.
- JAX currently lacks sparse-matrix support and must compute using dense matrices, which scale poorly.
This release restore support for PyAutoLens’s previous CPU implementation (via numba) which fully exploits sparsity, providing large speed gains at high image resolution (e.g. pixel_scales <= 0.03).
CPU execution can outperform JAX, even on powerful GPUs, for high-resolution datasets or when many CPU cores are used.
Development is actively working on how to get better performance from JAX that exploits sparsity on GPU, but this is proving to be a very challenging problem.
Delaunay
Support for the Delaunay mesh which was the main pixelized source reconstruction has been restored in this release, albeit it only currently works using the numba implementation above and therefore only supports CPU.
Development is actively working on having JAX support for Delaunay source reconstructions, with this expected to be available in the short-term.