microJAX: A Differentiable Framework for Microlensing Modeling with GPU-Accelerated Image-Centered Ray Shooting

Abstract

We introduce microJAX, the first fully differentiable implementation of the image-centered ray-shooting (ICRS) algorithm for gravitational microlensing. Built on JAX and its XLA just-in-time compiler, microJAX exploits GPU parallelism while providing exact gradients through automatic differentiation. The current release supports binary- and triple-lens geometries, including limb-darkened extended-source effects, and delivers magnifications that remain differentiable for all model parameters. Benchmarks show that microJAX matches the accuracy of established packages and attains up to a factor of $\sim$5-6 speed-up in the small-source, limb-darkened regime on an NVIDIA A100 GPU. Since the model is fully differentiable, it integrates seamlessly with probabilistic programming frameworks, enabling scalable Hamiltonian Monte Carlo and variational inference workflows. Although the present work focuses on standard microlensing magnification models, the modular architecture is designed to support upcoming implementations of microlensing higher-order effects, while remaining compatible with external likelihood frameworks that incorporate advanced noise models. microJAX thus provides a robust foundation for precise and large-scale surveys anticipated in the coming decade, including the Nancy Grace Roman Space Telescope, where scalable, physically self-consistent inference will be essential for maximizing scientific return.