pith. sign in

arxiv: 2606.01547 · v1 · pith:BRHN7INZnew · submitted 2026-06-01 · 🌌 astro-ph.IM · astro-ph.CO

JAXtronomy: A JAX port of lenstronomy

Alan Huang , Simon Birrer , Daniel Gilman , Natalie B. Hogg , Anowar J. Shajib , Aymeric Galan , Nan Zhang This is my paper

Pith reviewed 2026-06-28 13:08 UTC · model grok-4.3

classification 🌌 astro-ph.IM astro-ph.CO
keywords gravitational lensingJAXlenstronomyautomatic differentiationjust-in-time compilationGPU computingastrophysics softwarecosmology
0
0 comments X

The pith

JAXtronomy reimplements lenstronomy in JAX while preserving the original API exactly.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces JAXtronomy as a complete reimplementation of the lenstronomy gravitational lensing package using the JAX library. Its main goal is to keep the application programming interface identical to the original so that existing code and workflows continue to work without modification. By using JAX, the new package gains just-in-time compilation for speed, automatic differentiation for gradient-based optimization, and the ability to run on graphics processing units or across multiple central processing unit cores. These changes matter because gravitational lensing analysis often involves computationally intensive modeling of light bending around massive objects to study cosmology and dark matter distributions.

Core claim

We introduce JAXtronomy, a re-implementation of the gravitational lensing software package lenstronomy using JAX. Our core design principle of JAXtronomy is to maintain an identical API to that of lenstronomy. The main JAX features utilized in JAXtronomy are just-in-time compilation, which can lead to significant reductions in execution time, and automatic differentiation, which allows for the implementation of gradient-based algorithms that were previously impossible. Additionally, JAX allows code to be run on GPUs or parallelized across CPU cores, further boosting the performance of JAXtronomy.

What carries the argument

JAXtronomy, the JAX reimplementation of lenstronomy that preserves the original API

If this is right

  • Execution times for lensing calculations can be reduced through just-in-time compilation.
  • Gradient-based algorithms become feasible due to automatic differentiation.
  • Computations can be accelerated by running on GPUs or parallelized over CPU cores.
  • Existing lenstronomy code can be used directly with JAXtronomy without changes.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Similar API-preserving ports could be applied to other astronomy simulation packages.
  • Integration with JAX ecosystems might enable machine learning approaches to lensing model fitting.
  • Users could test performance gains on specific lensing problems to quantify benefits.

Load-bearing premise

A full reimplementation in JAX can preserve exact numerical behavior and API compatibility with the original lenstronomy.

What would settle it

Running identical inputs through both packages and finding outputs that differ beyond floating-point precision, or finding API calls that fail to match in signature or behavior.

read the original abstract

Gravitational lensing is a phenomenon where light bends around massive objects, resulting in distorted images seen by an observer. Studying gravitationally lensed systems provides insights into cosmology and astrophysics, including constraints of the expansion rate of the Universe and the distribution of dark matter. Thus, we introduce JAXtronomy, a re-implementation of the gravitational lensing software package lenstronomy (Birrer, 2021; Birrer & Amara, 2018) using JAX (Bradbury et al., 2018). JAX is a Python library that uses an accelerated linear algebra (XLA) compiler to improve the performance of computing software. Our core design principle of JAXtronomy is to maintain an identical API to that of lenstronomy. The main JAX features utilized in JAXtronomy are just-in-time compilation, which can lead to significant reductions in execution time, and automatic differentiation, which allows for the implementation of gradient-based algorithms that were previously impossible. Additionally, JAX allows code to be run on GPUs or parallelized across CPU cores, further boosting the performance of JAXtronomy.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 1 minor

Summary. The paper introduces JAXtronomy, a re-implementation of the gravitational lensing package lenstronomy using the JAX library. Its core design principle is to maintain an identical API to lenstronomy while utilizing JAX features such as just-in-time compilation for performance gains, automatic differentiation for gradient-based methods, and support for GPUs or parallel CPU execution.

Significance. If the port achieves functional equivalence with identical numerical outputs and API compatibility, it would enable faster lensing computations and new optimization techniques in cosmology and astrophysics, representing a useful methods contribution in astro-ph.IM.

major comments (2)
  1. [Abstract] Abstract: The manuscript states the intent to deliver performance improvements via JIT compilation and autodiff but supplies no benchmarks, timing results, or numerical verification that JAXtronomy matches lenstronomy outputs, leaving the central claim of a functional reimplementation untested.
  2. [Abstract] Abstract: The design goal of preserving exact numerical behavior while reimplementing in JAX is asserted without any test cases, comparison tables, or validation sections demonstrating equivalence on standard lensing problems.
minor comments (1)
  1. The paper would benefit from explicit discussion of how JAX-specific features (e.g., tracing limitations) were handled to maintain API identity with the original lenstronomy codebase.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for their constructive comments. We agree that the current manuscript version does not provide the necessary benchmarks or validation tests to substantiate the performance and equivalence claims, and we will revise the paper to include these elements.

read point-by-point responses

  1. Referee: [Abstract] Abstract: The manuscript states the intent to deliver performance improvements via JIT compilation and autodiff but supplies no benchmarks, timing results, or numerical verification that JAXtronomy matches lenstronomy outputs, leaving the central claim of a functional reimplementation untested.

    Authors: We acknowledge that the submitted manuscript does not include benchmarks or numerical verification. This omission leaves the central claims unsupported in the current version. In the revised manuscript we will add a results section containing timing benchmarks (CPU, GPU, and JIT vs. non-JIT) and direct numerical comparisons demonstrating that JAXtronomy reproduces lenstronomy outputs to machine precision on standard test problems. revision: yes

  2. Referee: [Abstract] Abstract: The design goal of preserving exact numerical behavior while reimplementing in JAX is asserted without any test cases, comparison tables, or validation sections demonstrating equivalence on standard lensing problems.

    Authors: We agree that the manuscript currently asserts API and numerical equivalence without supporting test cases or tables. We will incorporate a validation subsection that presents side-by-side comparisons on standard lensing configurations (e.g., SIS, SIE, and NFW profiles with point and extended sources) together with difference maps and quantitative metrics confirming functional equivalence. revision: yes

Circularity Check

0 steps flagged

No significant circularity in software port description

full rationale

The manuscript is a description of a JAX reimplementation of lenstronomy that preserves the original API. It contains no derivations, equations, predictions, fitted parameters, or first-principles results. The central claim is the factual existence of delivered code with matching behavior, which cannot reduce to its own inputs by construction. No self-citation chains, ansatzes, or uniqueness theorems are invoked as load-bearing steps. This is the expected non-finding for a pure software translation paper.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

No free parameters, axioms, or invented entities are introduced; the contribution is a software engineering port with no mathematical modeling or new physical postulates.

pith-pipeline@v0.9.1-grok · 5756 in / 1006 out tokens · 22201 ms · 2026-06-28T13:08:45.065093+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

10 extracted references · 9 canonical work pages

  1. [1]

    Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., … Viola, F. (2020). The DeepMind JAX Ecosystem . http://github.com/google-deepmind Bingham, E., Chen, J. P., Janko...

    2020

  2. [2]

    https://doi.org/10.3847/1538-4357/ac1108 Huang et al. (2026). JAXtronomy: A JAX port of lenstronomy. Journal of Open Source Software , 11(120),

    work page doi:10.3847/1538-4357/ac1108 2026

  3. [4]

    https://doi.org/10.3847/1538- 4357/ab5b90 Gilman, D., Birrer, S., Nierenberg, A., Treu, T., Du, X., & Benson, A. (2019). Warm dark matter chills out: Constraints on the halo mass function and the free-streaming length of dark matter with eight quadruple-image strong gravitational lenses. Monthly Notices of the Royal Astronomical Society , 491(4), 6077–610...

    work page doi:10.3847/1538- 2019

  4. [5]

    , keywords =

    https://doi.org/10.3847/1538-4357/ac6de4 Harris, C. R., Millman, K. J., Walt, S. J. van der, Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., Kerkwijk, M. H. van, Brett, M., Haldane, A., Río, J. F. del, Wiebe, M., Peterson, P., … Oliphant, T. E. (2020). Array programming with NumPy....

    work page doi:10.3847/1538-4357/ac6de4 2020

  5. [7]

    J., Treu, T., Birrer, S., & Sonnenfeld, A

    https: //doi.org/10.3847/1538-4357/adf95c Shajib, A. J., Treu, T., Birrer, S., & Sonnenfeld, A. (2021). Dark matter haloes of massive elliptical galaxies at z ∼ 0.2 are well described by the Navarro-Frenk-White profile. Monthly Notices of the Royal Astronomical Society , 503(2), 2380–2405. https://doi.org/10.1093/ mnras/stab536 Shajib, A. J., Wong, K. C.,...

    work page doi:10.3847/1538-4357/adf95c 2021

  6. [8]

    Y., Shajib, A

    https://doi.org/10.21105/joss.07081 Tan, C. Y., Shajib, A. J., Birrer, S., Sonnenfeld, A., Treu, T., Wells, P., Williams, D. M., Buckley-Geer, E. J., Drlica-Wagner, A., & Frieman, J. (2024). Project Dinos I: A joint lensing-dynamics constraint on the deviation from the power law in the mass profile of massive ellipticals. Monthly Notices of the Royal Astr...

    work page doi:10.21105/joss.07081 2024

  7. [9]

    https://doi.org/10.3847/1538-4357/aca525 Wagner-Carena, S., Lee, J., Pennington, J., Aalbers, J., Birrer, S., & Wechsler, R. H. Huang et al. (2026). JAXtronomy: A JAX port of lenstronomy. Journal of Open Source Software , 11(120),

    work page doi:10.3847/1538-4357/aca525 2026

  8. [11]

    2024, ApJ, 975, 297, doi:10.3847/1538-4357/ad6e70

    https://doi.org/10.3847/1538-4357/ad6e70 Wang, H., Suyu, S. H., Galan, A., Halkola, A., Cappellari, M., Shajib, A. J., & Cernetic, M. (2025). GPU-accelerated gravitational lensing and dynamical (GLaD) modeling for cosmology and galaxies. Astronomy & Astrophysics , 701, A280. https://doi.org/10.1051/ 0004-6361/202554861 Wedig, B., Daylan, T., Birrer, S., C...

    work page doi:10.3847/1538-4357/ad6e70 2025

  9. [12]

    M., Treu, T., Birrer, S., Shajib, A

    https://doi.org/10.3847/1538-4357/adc24f Williams, D. M., Treu, T., Birrer, S., Shajib, A. J., Wong, K. C., Morishita, T., Schmidt, T., & Stiavelli, M. (2025). TDCOSMO - XX. WFI2033–4723, the first quadruply imaged quasar modeled with JWST imaging. Astronomy & Astrophysics , 703, A118. https: //doi.org/10.1051/0004-6361/202554359 Huang et al. (2026). JAXt...

    work page doi:10.3847/1538-4357/adc24f 2025

  10. [13]

    https://doi.org/10.21105/joss. 09685. 7

    work page doi:10.21105/joss