Automatic Functional Differentiation in JAX - Lin, Min (International Conference on Learning Representations, 2024)