From e06a527b4cc5b0e688b84c22ad3ef42d14e10e18 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 30 Jul 2026 18:06:25 +0100 Subject: [PATCH] docs: add a JAX citations section, including Prodigy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PyAutoGalaxy's citations page had no JAX section at all, and files/citations.bib carried no jax, optax, nufftax, finufft or jax_zero_contour entries — so a user running the JAX path had nothing to cite, even though the NUFFTax and Jax-Zero-Contour sections already assumed those keys existed. Mirror PyAutoLens's JAX section and the entries it references, and add Prodigy (ICML 2024) for the learning-rate-free af.MultiStartProdigy search, cited alongside optax rather than instead of it since optax.contrib.prodigy is an implementation of a separately published method. Co-Authored-By: Claude Opus 5 --- docs/general/citations.md | 33 ++++++++++++++++++++++++ files/citations.bib | 53 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+) diff --git a/docs/general/citations.md b/docs/general/citations.md index 2616c7fe..b8c644b6 100644 --- a/docs/general/citations.md +++ b/docs/general/citations.md @@ -17,6 +17,39 @@ You should also specify the non-linear search(es) you use in your analysis (e.g. the main body of text, and delete as appropriate any packages your analysis did not use. The citations.bib file includes the citation key for all of these projects. +## JAX + +**PyAutoGalaxy** runs on a NumPy backend by default and an optional +[JAX](https://github.com/jax-ml/jax) backend for just-in-time compilation, automatic +differentiation, and GPU/TPU execution. If you run any analysis on the JAX path, please +cite JAX under the citation key `jax`. JAX-specific components that are also cited under +their own keys when used are `optax` (gradient-based optimizers, key `optax`), the +interferometer non-uniform FFT (`nufftax` and its FINUFFT algorithm, keys `nufftax` and +`finufft`, see below), and the critical-curve/caustic solver (`Jax-Zero-Contour`, key +`jax_zero_contour`, see below). + +If you use the learning-rate-free `af.MultiStartProdigy` search, also cite the Prodigy +method under the key `prodigy`. It is a published algorithm in its own right, run through +Optax's implementation (`optax.contrib.prodigy`), so cite it alongside `optax` rather than +in place of it: + +```bibtex +@inproceedings{prodigy, +author = {Mishchenko, Konstantin and Defazio, Aaron}, +title = {{Prodigy: An Expeditiously Adaptive Parameter-Free Learner}}, +booktitle = {Proceedings of the 41st International Conference on Machine Learning}, +series = {Proceedings of Machine Learning Research}, +volume = {235}, +pages = {35779--35804}, +publisher = {PMLR}, +url = {https://proceedings.mlr.press/v235/mishchenko24a.html}, +year = {2024} +} +``` + +The reference implementation is at and the preprint +at . + ## Jax-Zero-Contour If you use the zero-contour method for critical curve and caustic computation (the default in diff --git a/files/citations.bib b/files/citations.bib index 399fc146..7d035281 100644 --- a/files/citations.bib +++ b/files/citations.bib @@ -378,3 +378,56 @@ @article{Enzi2026 year = {2026}, url = {https://arxiv.org/abs/2606.30620}, } + +@software{jax, +author = {Bradbury, James and Frostig, Roy and Hawkins, Peter and Johnson, Matthew James and Leary, Chris and Maclaurin, Dougal and Necula, George and Paszke, Adam and VanderPlas, Jake and Wanderman-Milne, Skye and Zhang, Qiao}, +title = {{JAX}: composable transformations of {Python}+{NumPy} programs}, +url = {https://github.com/jax-ml/jax}, +year = {2018} +} + +@software{optax, +author = {{DeepMind} and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and others}, +title = {The {D}eep{M}ind {JAX} Ecosystem}, +url = {http://github.com/google-deepmind}, +year = {2020} +} + +@inproceedings{prodigy, +author = {Mishchenko, Konstantin and Defazio, Aaron}, +title = {{Prodigy: An Expeditiously Adaptive Parameter-Free Learner}}, +booktitle = {Proceedings of the 41st International Conference on Machine Learning}, +series = {Proceedings of Machine Learning Research}, +volume = {235}, +pages = {35779--35804}, +publisher = {PMLR}, +url = {https://proceedings.mlr.press/v235/mishchenko24a.html}, +year = {2024} +} + +@software{jax_zero_contour, +author = {Krawczyk, Coleman}, +title = {{CKrawczyk/Jax-Zero-Contour}: Version 2.0.0}, +publisher = {Zenodo}, +version = {v2.0.0}, +doi = {10.5281/zenodo.15730415}, +url = {https://doi.org/10.5281/zenodo.15730415}, +year = {2025} +} + +@software{nufftax, +author = {Gragas and Oudoumanessah, Geoffroy and Iollo, Jacopo}, +title = {nufftax: Pure {JAX} implementation of the Non-Uniform Fast Fourier Transform}, +url = {https://github.com/GragasLab/nufftax}, +year = {2026} +} + +@article{finufft, +author = {Barnett, Alexander H. and Magland, Jeremy F. and af Klinteberg, Ludvig}, +title = {A parallel non-uniform fast {F}ourier transform library based on an 'exponential of semicircle' kernel}, +journal = {SIAM J. Sci. Comput.}, +volume = {41}, +number = {5}, +pages = {C479--C504}, +year = {2019} +}