# Automatic Functional Differentiation in JAX

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 36 references that are in the map.
- First posted on arXiv 30 Nov 2023.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Lagrangian Neural Networks](https://mlsanity.com/paper/2003.04630/) (2020)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [The simple essence of automatic differentiation](https://mlsanity.com/paper/1804.00746/) (2018)
- [Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator](https://mlsanity.com/paper/1803.10228/) (2018)
- [Correctness of Automatic Differentiation via Diffeologies and Categorical Gluing](https://mlsanity.com/paper/2001.02209/) (2020)
- [You Only Linearize Once: Tangents Transpose to Gradients](https://mlsanity.com/paper/2204.10923/) (2022)
- [Efficient Differentiable Programming in a Functional Array-Processing Language](https://mlsanity.com/paper/1806.02136/) (2018)
- [λS: Computable Semantics for Differentiable Programming with Higher-Order Functions and Datatypes](https://mlsanity.com/paper/2007.08017/) (2020)
- [CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators](https://mlsanity.com/paper/2305.06347/) (2023)
- [A Language for Evaluating Derivatives of Functionals Using Automatic Differentiation](https://mlsanity.com/paper/2210.06095/) (2022)

## Built on by (most cited first, 2 of 2)

- [Physics-informed Neural Networks for Functional Differential Equations: Cylindrical Approximation and Its Convergence Guarantees](https://mlsanity.com/paper/2410.18153/) (2024)
- [Archimedean Copula Inference via Taylor-Mode AD](https://mlsanity.com/paper/2605.23134/) (2026)

## Abstract (arXiv)

We extend JAX with the capability to automatically differentiate higher-order functions (functionals and operators). By representing functions as a generalization of arrays, we seamlessly use JAX's existing primitive system to implement higher-order functions. We present a set of primitive operators that serve as foundational building blocks for constructing several key types of functionals. For every introduced primitive operator, we derive and implement both linearization and transposition rules, aligning with JAX's internal protocols for forward and reverse mode automatic differentiation. This enhancement allows for functional differentiation in the same syntax traditionally use for functions. The resulting functional gradients are themselves functions ready to be invoked in python. We showcase this tool's efficacy and simplicity through applications where functional derivatives are indispensable. The source code of this work is released at https://github.com/sail-sg/autofd .