jaxdae: A JAX-native Differentiable Solver for Differential-Algebraic Equations in Coupled Multi-physics

2026-07-25

AI Summary

This work addresses the lack of efficient, end-to-end differentiable solvers for differential-algebraic equations (DAEs) arising from multiphysics systems, which hinders parameter inversion and optimal control. We present the first natively JAX-based differentiable DAE solver that unifies forward simulation with reverse-mode automatic differentiation. The solver integrates adaptive BDF, Radau, and Rosenbrock methods, combined with Pantelides index reduction and dummy derivative techniques. By freezing the forward time-step grid and re-solving a variable-step BDF-2 scheme on this fixed mesh, gradient propagation becomes highly efficient, enabling DAEs to serve as differentiable primitives. A single call to jax.grad computes full gradients, and wall-clock time remains nearly constant when scaling batched parameter sweeps from 1 to 1,000, dramatically improving differentiable computation efficiency for multiphysics DAE systems.

Abstract

Many engineered models begin as partial differential equations. Spatial discretization converts them into ordinary differential equations coupled to algebraic constraints---conservation closures, constitutive laws, network topology---whose joint evolution is a differential-algebraic equation (DAE). Parameter inversion, uncertainty quantification, Bayesian inference, and optimal control all require gradients of this solve. The two software traditions that should supply them have not met: industrial acausal modeling tools simulate DAEs forward but stop at reverse-mode differentiation, while differentiable-physics frameworks in JAX handle explicit ODEs and PDEs but leave the algebraic-constraint layer untouched. Here we show that the forward DAE solve and its reverse-mode sensitivity can be unified in one JAX-native suite. jaxdae pairs adaptive BDF, Radau, and Rosenbrock integration with Pantelides index reduction and dummy derivatives, and makes the adaptive BDF path differentiable by freezing the accepted step grid and re-solving a variable-step BDF-2 on it for the backward pass. The full pipeline differentiates under one $\texttt{jax.grad}$ call, XLA fuses a batched parameter sweep into one program whose wall time stays nearly flat from batch~1 to~1000, and the DAE becomes a differentiable primitive for inference, control, and design.

Problem

• Differential-Algebraic Equations
• Differentiable Solver
• Reverse-mode Differentiation
• Multi-physics Modeling
• JAX

Innovation

• differentiable solver
• differential-algebraic equations
• JAX
• automatic differentiation
• multi-physics simulation