Search arXiv⌕ Search
Home/Archive/arXiv

arXiv · 2611.03918

Handling JAX

Oghenevwogaga Ebresafe·Daniel Zheng·Gordon Plotkin·Ningning Xie
View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

cs.PL

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Oghenevwogaga Ebresafe, Daniel Zheng, Gordon Plotkin, Ningning Xie. 2026-11-02. Handling JAX. https://arxiv.org/abs/2611.03918

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Proving the Wrong Program: Miscompilations as a Threat to zkVM Guarantees

A zero-knowledge virtual machine (zkVM) lets an untrusted prover convince a verifier that a program, typically Rust compiled to RISC-V, ran correctly. The proof covers the compiled binary, so a miscompilation can make a correct zkVM produce a valid proof of a computation that differs from the developer's source. We present the first empirical study of miscompilation as a threat to zkVM guarantees, on the production zkVMs SP1 and Jolt. We assemble 88 known rustc and LLVM miscompilations with Rust reproducers, using an LLM agent to recover Rust triggers for 60 of 124 LLVM bugs reported only as LLVM IR. Many of the bugs carry over to zkVMs: 48 of the 88 make a guest commit a wrong output on SP1 and 46 on Jolt, including the same 15 rustc bugs on both. Most bugs that do not reproduce never reach the faulty code on the zkVM's target, e.g. because they need an instruction-set extension. Existing guest programs are rarely affected by accident. Across 13 guests of 10 popular SP1 projects, the buggy and the fixed compiler produce byte-identical binaries in 447 of the 552 guest-bug pairs that build, and the other 105 have the same code or agree on every input we ran. An attacker who contributes code to a guest, however, can trigger these bugs on purpose. For about $40 per project, an LLM agent found patches that look like ordinary optimisations, preserve the behaviour of the guest's source, and expose a known bug in 12 of 366 project-bug pairs, across 11 of 20 projects. In four cases the guest commits a wrong output, including an Ethereum client that rejects a valid block, and in the other eight it fails on a valid input. Reviewing or formally verifying the zkVM or the guest's source cannot catch such a patch, because both are correct. zkVM projects should treat their compiler as part of their trusted base and track the known bugs of their toolchain.

cs.PL↗
arXiv2026

Properly defining Erlang lists

We define list types, distinguishing proper lists, which end in the empty list, from improper lists, which may end in any value. Both are encoded as recursive types of cons cells, implemented in etylizer. This work found one OTP compiler bug. We propose a one-line sound improvement to Dialyzer's list handling that finds six instances of dead code in the OTP code base. We compare BEAM VM type checkers and catalogue how other languages design lists and cons cells.

cs.PL↗
arXiv2026

Linking as Evaluation: Correct Separate Compilation via First-Class Environments

Separate compilation is central to modular software: components are compiled independently and linked into an executable. Separate compilation is correct when a linked program behaves like the program formed by merging its sources. This property is less routine to prove than type soundness because it requires a formal account of linking. We present $λ_{\mathsf{SCE}}$, a simply typed calculus in which linking, modules, and environments are all first-class. Programmers control module linking through a link operator, which a type-directed elaboration translates into ordinary terms of a core calculus of first-class environments. The core calculus's existing evaluation rules then perform the linking: linking as evaluation. Linking independently compiled modules yields the same core term as elaborating their source link. Therefore, compilation and linking commute by construction. As a result, separate compilation correctness follows directly from the semantics preservation for both static and dynamic linking. All meta-theory is mechanized in Lean 4, and we implement a compiler to WebAssembly in OCaml.

cs.PL↗