microJAX: Differentiable Microlensing Modeling

• microJAX is a differentiable microlensing framework that leverages JAX and GPU acceleration for precise, image-centered ray shooting finite-source calculations.
• It re-engineers standard microlensing magnification models with custom JVP rules and parallel root solving (Ehrlich-Aberth method) to enable exact gradient computation.
• The framework supports binary and triple-lens geometries and integrates with probabilistic inference workflows for scalable, gradient-based Bayesian analysis.

microJAX is a differentiable microlensing modeling framework built on JAX and its XLA just-in-time compiler, designed around GPU-accelerated image-centered ray shooting (ICRS) for high-precision finite-source calculations. It is presented as the first fully differentiable implementation of ICRS for gravitational microlensing, with current support for binary- and triple-lens geometries, limb-darkened extended-source effects, and magnifications that remain differentiable for all model parameters. Its central significance lies in combining physically detailed microlensing calculations with automatic differentiation and accelerator-oriented execution, so that forward modeling, gradient-based optimization, Hamiltonian Monte Carlo, and variational inference can be carried out within a single computational framework (Miyazaki et al., 3 Oct 2025).

1. Scope and scientific motivation

microJAX addresses the regime in which microlensing calculations become numerically demanding: extended sources, multiple lenses, caustic crossings, and surface-brightness structure. Point-source magnification is comparatively straightforward, but finite-source magnification near caustics requires numerically accurate area integration over lensed images, and this is precisely the regime in which standard approximations can fail. The framework is motivated by the anticipated data volume and model complexity of upcoming Roman Space Telescope microlensing surveys, for which scalable modeling of binary and triple lenses, limb darkening, and higher-order effects is expected to be operationally important (Miyazaki et al., 3 Oct 2025).

The framework’s purpose is therefore not limited to fast light-curve synthesis. It is designed to make high-precision microlensing calculations compatible with modern autodiff and probabilistic inference stacks. The paper explicitly positions this as enabling physically self-consistent inference at scale, while remaining compatible with external likelihood frameworks that incorporate advanced noise models. A common misconception is that differentiable microlensing models are necessarily restricted to simplified or approximate magnification formulas; microJAX instead targets standard microlensing magnification models, but re-engineers the numerical machinery so that exact gradients through automatic differentiation can be exposed at the model level (Miyazaki et al., 3 Oct 2025).

2. Microlensing formulation

For an $N$-lens point-mass system, microJAX uses the complex lens equation

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$

with $\sum_i \epsilon_i = 1$. The corresponding polynomial form has degree

$N_{\rm deg}=5(N-1),$

so binary lenses produce degree-5 polynomials and triple lenses degree-10 polynomials. Up to $5(N-1)$ image roots may arise, although not all polynomial roots are physical and must be checked against the lens equation (Miyazaki et al., 3 Oct 2025).

For a point source, magnification is written as

$$A(\bm{w}) = \sum_{j=1}^{N_{\rm image}} \frac{1}{|\det J(\bm{z}_j)|},$$

with Jacobian determinant

$$\det J(\bm{z}_j)=1-\left|\sum_{i=1}^{N}\frac{\epsilon_i}{(\overline{\bm{z}_j-\overline{\bm{a}_i}})^2}\right|^2.$$

Critical curves satisfy $\det J=0$, and their source-plane images are caustics. These are the loci around which finite-source treatment becomes essential (Miyazaki et al., 3 Oct 2025).

For finite sources, microJAX computes a surface-brightness-weighted image-plane integral. In ICRS, grid points $\bm{z}_j^{(k)}$ in the image plane are mapped back to the source plane as

$\bm{w}_j^{(k)} = f_{\rm map}(\bm{z}_j^{(k)}),$

and source membership is encoded by

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$0

The finite-source magnification is then approximated by

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$1

with

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$2

This formulation is the mathematical basis for the framework’s treatment of arbitrary surface-brightness profiles and for its robustness near caustics (Miyazaki et al., 3 Oct 2025).

3. Image-centered ray shooting and root solving

The defining algorithmic choice in microJAX is the use of image-centered ray shooting rather than contour integration. Conventional ICRS implementations often rely on branching, adaptive region growth, and other control-flow patterns that are poorly aligned with GPU execution and JAX’s static compilation model. microJAX restructures the method into a static, vectorized pipeline: it samples the source limb uniformly with $$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$3 points,

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$4

maps these points to the image plane, and then uses polar-coordinate binning to identify compact annular sectors that enclose each image. The paper states that this binning scales as $$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$5, in contrast to more adaptive approaches that can scale like $$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$6 and are harder to JIT-compile (Miyazaki et al., 3 Oct 2025).

Each annular sector is integrated on a fixed $$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$7 grid using vectorized operations. Architecturally, the code decomposes the workflow into root solving, image-region construction, inverse-ray integration, and light-curve batching. It uses vmap for data-parallel operations and lax.scan for memory-efficient sequential accumulation across sectors or epochs. The dominant ray-shooting workload has memory scaling

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$8

which is the paper’s stated estimate for the main integration cost (Miyazaki et al., 3 Oct 2025).

For root finding, microJAX adopts the Ehrlich-Aberth method rather than companion-matrix eigensolvers or sequential Newton-like solvers. The update is

$$\bm{w} = \bm{z} - \sum_{i=1}^{N}\frac{\epsilon_i}{\overline{\bm{z}-\overline{\bm{a}_i}}},$$9

Because all root estimates are updated in parallel with fixed iteration count, the method maps naturally to JAX’s static graph and to GPU execution. The paper reports that this enables solution of about $\sum_i \epsilon_i = 1$0 10th-degree complex polynomials in under one second on an NVIDIA A100 GPU (Miyazaki et al., 3 Oct 2025).

4. Differentiability architecture

Differentiability is not treated as an incidental by-product of the JAX implementation; it is a primary design constraint. While many operations can be differentiated directly by JAX, microlensing calculations contain intrinsically discontinuous or piecewise elements, including source-boundary membership tests and interpolation-based edge corrections. microJAX addresses this by defining custom JVP rules with @custom_jvp, so that the model remains differentiable even where naïve reverse-mode or forward-mode propagation through the raw numerical procedure would be ill-posed or unstable (Miyazaki et al., 3 Oct 2025).

For the polynomial roots, the framework uses an implicit-function gradient,

$\sum_i \epsilon_i = 1$1

where

$\sum_i \epsilon_i = 1$2

This permits gradient propagation through the Ehrlich-Aberth solver without differentiating through every iteration. For ray-inclusion tests, the Heaviside indicator $\sum_i \epsilon_i = 1$3 is replaced in the JVP by a smooth sigmoid-like approximation with a tunable steepness factor. The framework also defines custom smooth gradients for the piecewise correction function used in limb-darkened boundary segments (Miyazaki et al., 3 Oct 2025).

A common misunderstanding is that “fully differentiable” here means that the primal computation is made everywhere smooth. That is not what the implementation does. The primal path retains hard geometric boundaries where needed for accurate image-plane integration, while the differentiation path introduces custom smooth gradient rules for the nonsmooth components. The paper also notes a practical limitation: lax.scan currently limits reverse-mode AD in some parts, so the present implementation mainly supports forward-mode differentiation for those components (Miyazaki et al., 3 Oct 2025).

5. Supported models, numerical accuracy, and performance

The current release supports binary- and triple-lens geometries. In the triple-lens examples, the parameterization is $\sum_i \epsilon_i = 1$4, where $\sum_i \epsilon_i = 1$5 is the binary mass ratio, $\sum_i \epsilon_i = 1$6 the binary separation, $\sum_i \epsilon_i = 1$7 the third-lens mass ratio relative to the primary, and $\sum_i \epsilon_i = 1$8 the third-lens position in polar coordinates relative to the binary midpoint. The framework also supports limb-darkened extended sources with the linear limb-darkening law

$\sum_i \epsilon_i = 1$9

where $N_{\rm deg}=5(N-1),$0 and $N_{\rm deg}=5(N-1),$1 is chosen so that the total source flux is normalized (Miyazaki et al., 3 Oct 2025).

Accuracy is benchmarked against VBBinaryLensing for a binary lens with $N_{\rm deg}=5(N-1),$2 and $N_{\rm deg}=5(N-1),$3, using source radii from $N_{\rm deg}=5(N-1),$4 down to $N_{\rm deg}=5(N-1),$5, azimuthal resolutions $N_{\rm deg}=5(N-1),$6, and fixed $N_{\rm deg}=5(N-1),$7. The reported result is that relative magnification errors are below $N_{\rm deg}=5(N-1),$8 for sufficiently resolved cases, while for the smallest source size tested, $N_{\rm deg}=5(N-1),$9 is needed to keep the error below $5(N-1)$0. The paper identifies angular resolution as the critical factor in tiny-source, near-caustic events; radial quadrature converges more gently, whereas azimuthal boundary resolution controls whether discontinuities or spikes appear in the light curve (Miyazaki et al., 3 Oct 2025).

Performance is explicitly described as regime-dependent. In the small-source, limb-darkened regime on an NVIDIA A100 GPU, microJAX attains a speed-up of about $5(N-1)$1–$5(N-1)$2 relative to VBBinaryLensing on an AMD EPYC CPU. For large uniform sources, microJAX can be about $5(N-1)$3–$5(N-1)$4 slower than VBBinaryLensing. At the same time, its runtime is much less sensitive to source profile than contour integration: the paper states that runtime differs by less than a factor of 2 between uniform and limb-darkened sources, whereas VBBinaryLensing can be about $5(N-1)$5 faster for uniform than for limb-darkened profiles. The timings exclude compilation overhead, and the paper explicitly notes that direct CPU-versus-GPU comparison is hardware-dependent (Miyazaki et al., 3 Oct 2025).

6. Probabilistic inference workflows

Because the magnification model is differentiable with respect to physical parameters, microJAX can be embedded directly in probabilistic programming frameworks. The paper specifically mentions NumPyro and demonstrates Hamiltonian Monte Carlo with the No-U-Turn Sampler. This is a central departure from traditional microlensing toolchains, in which finite-source magnification is often treated as an expensive black-box likelihood component rather than as a differentiable primitive suitable for gradient-based sampling and optimization (Miyazaki et al., 3 Oct 2025).

The reported demonstrations include a synthetic Roman-like injection-recovery problem and a real-data application. In the synthetic case, the framework fits a 72-day light curve sampled every 12 minutes and recovers the injected parameters well, with convergence diagnostics $5(N-1)$6 and effective sample sizes above 800. In the real-data example, it analyzes OGLE-2014-BLG-0124 with a binary-lens model including parallax and recovers a plausible posterior with $5(N-1)$7 and ESS above 5000. The paper emphasizes that HMC captures non-Gaussian degeneracies much better than Fisher-matrix ellipses or local Gaussian approximations. This positions microJAX not merely as a forward solver, but as a framework for full Bayesian inference in high-dimensional microlensing problems (Miyazaki et al., 3 Oct 2025).

7. Broader JAX context, limitations, and extensions

microJAX belongs to a broader class of JAX-native scientific software that combines JIT compilation, accelerator execution, and differentiability in domain-specific numerical codes. A plausible implication is that it occupies, for microlensing, a role analogous to what jaxFMM does for adaptive Fast Multipole Method evaluation (Kraft et al., 19 Nov 2025), JANC does for differentiable reacting-flow simulation and JAX-based AMR (Wen et al., 18 Apr 2025), JAX-SCM does for modular atmospheric single-column modeling (Pierzyna, 23 May 2026), immrax does for interval analysis and mixed monotone reachability as function transforms (Harapanahalli et al., 2024), and gyaradax does for a minimal JAX/CUDA gyrokinetic solver (Galletti et al., 7 Apr 2026). In each case, the common design pattern is not merely “code written in Python,” but numerical infrastructure recast around JAX transformations, GPU execution, and differentiable programming.

Within that context, microJAX is specialized rather than universal. The present work focuses on standard microlensing magnification models, and the current release is limited to up to three lenses. The modular architecture is intended to support higher-order effects in future versions, and the paper explicitly identifies orbital motion, xallarap, more than three lenses, and astrometric microlensing as future extensions. It also points to realistic noise models and integration with differentiable Gaussian processes or state-space models as natural next steps. These statements suggest a framework intended to remain compatible with external likelihood machinery while extending the differentiable core beyond the current magnification-only scope (Miyazaki et al., 3 Oct 2025).