# Bringing PDEs to JAX with forward and reverse modes automatic differentiation

arXiv: https://arxiv.org/abs/2309.07137 · HTML: https://mlsanity.com/paper/2309.07137/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 4 of its 11 references that are in the map.
- First posted on arXiv 31 Aug 2023.

## Builds on (most cited first, 4 of 4)

- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [Learning to Control PDEs with Differentiable Physics](https://mlsanity.com/paper/2001.07457/) (2020)
- [Neural reparameterization improves structural optimization](https://mlsanity.com/paper/1909.04240/) (2019)
- [Neural network augmented inverse problems for PDEs](https://mlsanity.com/paper/1712.09685/) (2017)

## Abstract (arXiv)

Partial differential equations (PDEs) are used to describe a variety of physical phenomena. Often these equations do not have analytical solutions and numerical approximations are used instead. One of the common methods to solve PDEs is the finite element method. Computing derivative information of the solution with respect to the input parameters is important in many tasks in scientific computing. We extend JAX automatic differentiation library with an interface to Firedrake finite element library. High-level symbolic representation of PDEs allows bypassing differentiating through low-level possibly many iterations of the underlying nonlinear solvers. Differentiating through Firedrake solvers is done using tangent-linear and adjoint equations. This enables the efficient composition of finite element solvers with arbitrary differentiable programs. The code is available at github.com/IvanYashchuk/jax-firedrake.