Papers With

← All papers
First page of Handling JAX

Handling JAX

Oghenevwogaga Ebresafe, Daniel Zheng, Gordon Plotkin, Ningning Xie

cs.PL Nov 2, 2026 · v1
verification formalization type-theory
Read PDF arXiv abstract
TL;DR
The jaxpr effect-handler calculus and the proofs that its handler-eliminating translation preserves typing and semantics are mechanized in Lean 4.
Abstract
Languages that target hardware accelerators typically expose a restricted, first-order core without recursion or first-class functions, leaving programmers without a principled way to express computational effects. This paper presents a technique for compiling effect handlers to such a constrained target. We introduce a novel translation that statically eliminates effect handlers, and prove that it preserves both typing and semantics. Moreover, we address the challenges of integrating effect handlers with structured control flow, where continuations cannot be constructed statically. The calculus and all our results are fully mechanized in the Lean 4 proof assistant. Finally, we implement our design as a library for JAX, a Python domain-specific language for accelerator programming, and demonstrate its expressiveness through practical case studies.
Problem

Accelerator-oriented languages such as JAX's jaxprs expose a first-order core without recursion or first-class functions. This leaves programmers with no principled way to express computational effects like mutable state and exceptions. JAX libraries therefore rely on ad-hoc, mutually incompatible mechanisms.

Approach

The authors define a calculus formalizing core jaxprs extended with algebraic effect handlers. They give a novel translation that statically eliminates handlers into the handler-free subset. For unbounded loops, where continuations cannot be built statically, tail-resumptive handlers are handled in place and abortive handlers by a source-level rewrite. The calculus and the type- and semantics-preservation proofs are mechanized in Lean 4, and the design is implemented as the JAX library effecax.

Results

The translation is proven to preserve typing and semantics, and every paper result has a Lean counterpart. Case studies refactor Diffrax's ODE integration core and probabilistic programming patterns with effecax. These were evaluated for expressiveness, composition with JAX control flow, and performance on an A100 GPU.