JAX Autodiff refers to the core of the automatic differentiation (AD) systems developed in projects like JAX and Dex. JAX Autodiff has recently been formalised in a linear typed calculus by Radul et al in POPL 2023. Although this formalisation suffices to express the main program transformations of AD, the calculus is very specific to this task, and it is not clear whether the type system yields a substructural logic that has interest on its own.
We propose an encoding of JAX Autodiff into a linear $\lambda$-calculus that enjoys a Curry-Howard correspondence with Girard's linear logic. We prove that the encoding is sound both qualitatively (the encoded terms are extensionally equivalent to the original ones) and quantitatively (the encoding preserves the original work cost as described in Radul et al). As a byproduct, we show that unzipping, one of the transformations used to implement backpropagation in JAX Autodiff, is, in fact, optional.
Thu 15 JanDisplayed time zone: Brussels, Copenhagen, Madrid, Paris change
14:00 - 15:40 | |||
14:00 25mTalk | A Logic for the Imprecision of Abstract Interpretations POPL Marco Campion Inria Paris - ENS - Université PSL, Mila Dalla Preda University of Verona, Roberto Giacobazzi University of Arizona, Caterina Urban Inria Paris - ENS - Université PSL DOI | ||
14:25 25mTalk | Big-Stop Semantics: Small-Step Semantics in a Big-Step Judgment POPL David M. Kahn Denison University, Jan Hoffmann Carnegie Mellon University, Runming Li Carnegie Mellon University DOI | ||
14:50 25mTalk | JAX Autodiff from a Linear Logic Perspective POPL DOI | ||
15:15 25mTalk | U-Turn: Enhancing Incorrectness Analysis by Reversing Direction POPL Flavio Ascari University of Konstanz, Roberto Bruni University of Pisa, Roberta Gori Diaprtimento di Informatica, Universita' di Pisa, Italy, Azalea Raad Imperial College London DOI | ||