Handling JAX

Oghenevwogaga Ebresafe, Daniel Zheng, Gordon Plotkin, Ningning Xie

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.