JAX Autodiff from a Linear Logic Perspective (Extended Version)

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Giusti, Giulia, Pagani, Michele
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866918174599938048
author Giusti, Giulia
Pagani, Michele
author_facet Giusti, Giulia
Pagani, Michele
contents Autodiff refers to the core of the automatic differentiation systems developed in projects like JAX and Dex. Autodiff has recently been formalised in a linear typed calculus by Radul et al in arXiv:2204.10923. Although this formalisation suffices to express the main program transformations of Autodiff, 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 Autodiff into a linear $λ$-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 arXiv:2204.10923). As a byproduct, we show that unzipping, one of the transformations used to implement backpropagation in Autodiff, is, in fact, optional.
format Preprint
id arxiv_https___arxiv_org_abs_2510_16883
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle JAX Autodiff from a Linear Logic Perspective (Extended Version)
Giusti, Giulia
Pagani, Michele
Programming Languages
Logic in Computer Science
Autodiff refers to the core of the automatic differentiation systems developed in projects like JAX and Dex. Autodiff has recently been formalised in a linear typed calculus by Radul et al in arXiv:2204.10923. Although this formalisation suffices to express the main program transformations of Autodiff, 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 Autodiff into a linear $λ$-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 arXiv:2204.10923). As a byproduct, we show that unzipping, one of the transformations used to implement backpropagation in Autodiff, is, in fact, optional.
title JAX Autodiff from a Linear Logic Perspective (Extended Version)
topic Programming Languages
Logic in Computer Science
url https://arxiv.org/abs/2510.16883