Handling JAX

Introduction

“When walking about the countryside of Italy, the people will not hesitate to tell you that JAX has ‘una anima di pura programmazione funzionale’.” – JAX documentation¹

Languages that target hardware accelerators are deliberately restrictive. To generate efficient code for GPUs and TPUs, they typically expose a purely functional, first-order core: programs are built from array primitives and a fixed set of structured control-flow constructs such as conditionals and loops, without recursion or first-class functions. Examples include array languages such as Futhark [21], Dex [36], and Accelerate [7], and compiler intermediate representations such as XLA [17]. While these restrictions enable aggressive optimization, they leave programmers without a principled way to express computational effects such as mutable state and exceptions. Without language support, each such feature is typically implemented via a bespoke program transformation or an external library.

We study this problem in the setting of JAX [6], a domain-specific language embedded in Python for accelerator-oriented array computations, widely adopted in high-performance numerical computing and large-scale machine learning, including for Google’s large language models [8, 15, 18] and by Anthropic [2]. JAX traces Python programs into a first-order intermediate representation, jaxprs, with exactly the restrictions described above, and lacks direct support for effectful features. As a result, popular JAX libraries resort to different, often ad-hoc mechanisms to implement these features, causing inconsistencies and interoperability issues across the ecosystem (see § 2).

We address this challenge using algebraic effect handlers [41, 42], a flexible and composable mechanism for defining structured control-flow abstractions. While effect handlers have been incorporated into various languages [3, 29, 48], including recent work on low-level languages such as C [1, 28], C++ [16], and WebAssembly [38], all such implementations rely on features that first-order languages lack: higher-order functions, recursion, or a capturable runtime stack. The central question of this paper is therefore: how can effect handlers be compiled to a first-order target without higher-order functions or recursion? We resolve this by introducing a translation that statically eliminates handlers, carefully extending the technique across loop boundaries where continuations cannot be constructed statically.

Contributions. Following the high-level overview in §2, we present our contributions in the following sections:

• Formalization (§3): We formalize $\lambda_{\mathrm{JAXPR}}$, a calculus that extends a core subset of JAX’s intermediate representation, a first-order language with second-order control-flow primitives, with parameterized algebraic effect handlers. We establish preservation and progress with effects.

• Translation (§4): We introduce a novel continuation-based translation that statically compiles effect handlers away without needing higher-order functions or recursion, and prove that it preserves both typing and semantics.

• Iterative control flow (§5): We analyze the challenges of translating effect handlers across loop boundaries, where continuations cannot be constructed statically, and present a mechanism to support tail-resumptive and abortive handlers while preserving handler parameters and scoping.

• Mechanization: We fully mechanize our calculus, its type system and operational semantics, translation, and all soundness proofs in the Lean 4 proof assistant.

• Implementation (§6): We implement effecax, an effect handler library for JAX, outlining its programming interface and our implementation.

• Evaluation (§7): We evaluate effecax across practical case studies, demonstrating that it provides an expressive, flexible, and natural interface for JAX programs.

While this paper focuses on JAX, where effecax provides immediate utility to a broad user base, neither the problem nor our solution is specific to it: the compilation applies to any first-order target, such as XLA’s intermediate representation. More broadly, we hope this work offers a new perspective on compiling effect handlers to low-level target languages. Lastly, we discuss related work in §8.

Overview

Motivation

JAX adopts a functional programming paradigm: its compilation pipeline is designed for pure Python functions free of side effects. This model aligns naturally with modern machine learning workloads, which predominantly consist of static mathematical operations connected in computation graphs. Functional purity simplifies graph optimization and yields significant performance gains [14]. Moreover, purity enables JAX to reliably trace Python functions into its intermediate representation, jaxpr²: a first-order, statically typed language in A-normal form [13], with a fixed set of second-order control-flow primitives such as cond and while. JAX also provides powerful function transformations, such as jax.grad for automatic differentiation, jax.vmap for automatic vectorization, and jax.jit for just-in-time compilation of functions to hardware accelerators via XLA [17].

The effectiveness of this functional programming model is evidenced by JAX’s widespread adoption in the machine learning community [2, 8, 15, 18]. Its paradigm of composable transformations is so powerful that it inspired PyTorch’s functorch³, which has been fully integrated into core PyTorch⁴.

However, while the functional programming model is highly desirable for machine learning workloads, practical ML programming is inherently stateful and side-effect intensive. ML training pipelines must manage mutable state, handle exceptions, perform asynchronous computation, and occasionally use specialized effects such as coroutines in probabilistic programming. In short, computational effects abound. This creates a fundamental tension: JAX’s functional core provides no principled mechanism for needed side-effectful computations. To illustrate this, we examine two effects common in everyday ML programming: mutable state and exception handling.

Case 1: Mutable state. To use mutable references, developers can either leverage JAX’s experimental built-in Ref arrays⁵ or rely on external libraries. The Ref approach uses a custom program transformation that converts programs containing Refs into purely functional programs by automatically threading state parameters through the computation.

Alternatively, JAX-based libraries such as Flax NNX⁶, a popular neural network library, expose Modules that are mutable by default. To maintain compatibility with JAX’s functional requirements, NNX implements its own transformations, such as flax.nnx.jit. This transformation wraps jax.jit by “lifting” the mutable module: it splits the module into a static structure and a pure state, applies the transformation to those pure components, and then propagates any state updates back to the original mutable object.

Both approaches involve trade-offs. While Ref integrates tightly with JAX’s built-in transformations, it is constrained by JAX’s tracing rules: for instance, Refs cannot be returned from jit-compiled functions or from the bodies of higher-order primitives (e.g. while), and cannot be passed where arrays are expected. Conversely, flax.nnx.jit provides a more ergonomic mutable interface, but requires a duplicated ecosystem of custom transformations and incurs Python performance overhead in training loops from repeatedly lifting stateful objects⁷.

Case 2: Exceptions. A similar situation exists for exception handling. JAX provides an experimental built-in transformation called checkify⁸, which rewrites a program that may raise errors into a functionally pure version that returns error values. JAX libraries such as Equinox⁹, on the other hand, offer equinox.error_if to raise runtime errors imperatively during execution.

While checkify adheres to JAX’s functional paradigm and composes cleanly with other transformations, it introduces significant boilerplate by forcing developers to manually propagate and inspect error values throughout the computation. It also lacks support for early termination, which can be important for performance. Conversely, Equinox makes it easy to insert error checks into existing models without refactoring, but degrades performance by forcing synchronous execution barriers between the hardware accelerator and the host.

Current landscape and limitations. As a result, JAX programmers must constantly choose between specialized, ad-hoc JAX transformations or fragmented external libraries. Neither is ideal. Dedicated internal transformations encourage code duplication and lack modularity. For instance, the Ref transformation implements state mutation using a substantial amount of code, while the checkify transformation relies on an entirely separate implementation. Despite this, both transformations share similar structures and “parameter plumbing” patterns. This redundancy inflates JAX’s codebase and its surface for compiler bugs.

Alternatively, relying on custom libraries often results in performance degradation, poor tool integration, and ecosystem fragmentation. Interoperability remains a painful challenge. For example, disparate state management mechanisms across libraries require complex, fragile manual state conversions¹⁰. Moreover, combining multiple effects from different libraries may not work as expected.

Furthermore, if a developer requires a specific effect for which no custom JAX transformation or third-party library exists, they are forced to sacrifice expressive power or resort to brittle, ad-hoc workarounds. Given JAX’s central role in machine learning, the lack of a unified approach to computational effects remains a fundamental limitation.

Key Ideas

To resolve these limitations, we seek a side-effect framework that has:

• Extensibility: Developers should be able to define custom side effects as well as specify their corresponding handling logic. Multiple effects should compose seamlessly; for example, a program should be able to manage mutable state while simultaneously throwing exceptions.

• Integration: The mechanism should integrate seamlessly with the language’s existing facilities. Moreover, since inspecting the intermediate representation is standard practice for debugging and optimization, it should produce clean, idiomatic code in that representation.

• Backend preservation: Accelerator-oriented languages rely on a downstream compiler, XLA in the case of JAX, to target hardware. The mechanism should remain fully compatible with existing backends and require no changes to that compiler.

In this work, we propose to extend JAX with algebraic effect handlers [40, 43], a flexible framework for extensible and composable computational effects in functional settings. Concretely, we extend the jaxpr representation with effect-handler constructs. As an example, consider the function safe_div, which performs a throw operation on division by zero (types are omitted; the syntax is detailed in § 3):

We handle it with a handler whose parameter (args) is unused and thus takes the unit value $\star$, runs the handled computation (body), which calls safe_div on 42 and 0, returns its result unchanged (ret), and handles throw by resuming the computation with the fallback value 0 (impl):

We then compile these constructs away by a translation that produces standard, vanilla jaxprs, ensuring that existing transformations and backends apply unchanged. Realizing this plan, however, raises two key challenges, which we discuss next.

Challenge 1: a first-order target. Effect handler implementations typically capture operation continuations either at runtime, by copying or manipulating the stack [1, 16, 38, 48], or at compile time, via continuation passing style (CPS) translation that represents continuations as closures [23, 46, 47]. Jaxprs offer neither facility, as they have no runtime stack to capture, and no higher-order functions or recursion to represent continuations.

We present a novel compilation strategy for this setting; to our knowledge, this is the first work to study effect handlers under such restrictions. Conceptually, our compilation combines a CPS transformation, which compiles handlers into continuations, with partial evaluation to statically evaluate all handler-related constructs. As a result, neither handlers nor continuations appear in the output, leaving only primitive computations for the hardware. Fusing CPS with static reduction is a classic idea underlying one-pass CPS transformations [9, 10] and CPS translations of effect handlers without administrative redexes [23, 47]; we apply it here directly to the handler constructs themselves. On the example above, the translation produces the following jaxpr:

Here, the do throw in the then branch has been statically evaluated to its handling code, here just the value 0, while the primitive computations are left untouched for execution on the hardware.

Challenge 2: loop boundaries. This strategy, however, requires the continuation of every operation to be known statically. Consider an operation performed inside a while loop but handled outside it, as in Figure 1. The continuation of the operation comprises not only the rest of the current iteration but also all remaining iterations, that is, the loop itself. Translating it would then unroll the loop. As a result, the translation does not terminate. Unfortunately, this is a common use case: the state and exception effects of § 2.1 are typically performed inside training and solver loops whose handlers live outside. We support these most common cases by identifying two classes of handlers that can nevertheless be compiled across loop boundaries.

We focus on tail-resumptive handlers, which resume the computation exactly once at the end of the operation clause and are common in practice [52, 54]. For such common clauses, we avoid materializing continuations entirely by handling operations in-place and inlining the clause directly into the loop body. This approach introduces two subtleties that we systematically address: threading handler parameters through the loop via an extension we call the final continuation, and preventing naive inlining from altering the handler scope. Compared to prior work, our approach is novel in several aspects: it applies to jaxprs, supports parameterized handlers, and preserves semantics without restrictions [52, 54].

Additionally, we support abortive handlers, which never resume and are common for exceptions. We achieve this through a complementary source-level rewrite of the while expression, denoted $e \rightsquigarrow e'$, before translation. Importantly, our abortive handlers are more general than JAX’s checkify, which must complete all iterations before checking for errors; in contrast, ours catches the exception inside the loop and aborts at the next condition check.

These two classes do not cover every handler; for example, a clause that resumes more than once is not supported across a loop boundary. We discuss this limitation further in § 8.

The rest of the paper proceeds as follows. § 3 presents the calculus, § 4 the translation, and § 5 its extension across loop boundaries; all results are mechanized in Lean 4. § 6 describes the effecax library built on the translation, and § 7 evaluates it through case studies.

The Calculus

This section presents $\lambda_{\mathrm{JAXPR}}^{\varepsilon}$, a calculus formalizing a core subset of jaxprs extended with effect-handling constructs. At its heart, $\lambda_{\mathrm{JAXPR}}^{\varepsilon}$ is simply a first-order calculus with second-order primitives: functions cannot be passed as values, but a fixed set of primitives, such as conditionals and loops, take jaxprs as arguments. This design is not specific to JAX; rather, it captures the essential shape of accelerator-oriented intermediate representations. At the same time, the calculus remains faithful to actual jaxprs, allowing our translation to apply directly to the implementation. To our knowledge, this is the first formalization of a core jaxpr language.

Syntax

Figure 2 presents the syntax of $\lambda_{\mathrm{JAXPR}}$. The calculus is a statically-typed, first-order language in A-normal form (ANF) [13]. We focus on modeling a minimal set of primitives; additional primitives can be added in a straightforward manner. The syntax comprises three syntactic categories: types, terms (which include expressions and jaxprs), and primitives. We assume disjoint name spaces for variables ($x$), constants ($c$), basic primitive functions ($f$), and effect operations ($\mathit{op}$).

Types $\sigma$ are product types, written as sequences of array types for typing tensors (i.e. multidimensional arrays). Array types $\alpha$ have the form $d[s]$, where $d$ represents the underlying scalar data type (e.g. float, int), and $s$ is a shape tuple specifying the number and sizes of the tensor dimensions (e.g. $3 \times4$). For scalar values, we often omit the empty shape tuple for brevity, e.g. writing bool instead of $\mathrm{bool}[]$. We also often write $\star$ for the unit type (and its unit value).

Expressions $e$ consist of a sequence of let bindings $\mathop{\mathrm{let}} b^{*} = p\ a^{*}\ \mathop{\mathrm{in}}\ e$ that end with a return statement $\mathop{\mathrm{return}}\ a^{*}$. Binders $b$ are typed variables. Atoms $a$ are either variables $x$ or literals $l$ specifying tensors. Literals $l$ include constants $c$ and tensors $[l^{*}]$. A jaxpr $j$ is a lambda expression, written $\{\lambda b^{*}.e\}$.

Primitive operations $p$ include both basic functions $f$ and higher-order primitives $\hat{p}$ built from jaxprs. These higher-order primitives include function calls (call $[\ \mathrm{body}:j\ ]$), that call the jaxpr $j$ when applied within let bindings; conditionals (cond $[\ \mathrm{then}:j_{1},\ \mathrm{else}:j_{2}\ ]$), that take a boolean argument and execute the appropriate branch; while loops (while $[\ \mathrm{cond}:j_{1},\ \mathrm{body}:j_{2}\ ]$) with condition jaxpr $j_{1}$ and body jaxpr $j_{2}$; and scans (scan $[\ \mathrm{body}:j,\ \mathrm{carry}:a_{1},\ldots,a_{n}\ ]$) that provide a functional loop-fold, where $j$ iterates over input arrays while updating a carry state initialized by $a_{1},\ldots,a_{n}$.

To support algebraic effects, $\lambda_{\mathrm{JAXPR}}$ introduces primitives for performing and handling effectful computations. First, the primitive do $\mathit{op}$ performs an effectful operation $\mathit{op}$. Second, following Leijen [29], Schuster et al. [47], we represent resumptions using a resume primitive. Finally, the syntax includes a parameterized handler primitive: handle $[\ \mathrm{args}:a^{*},\ \mathrm{body}:e,\ \mathrm{ret}:j,\ \mathrm{impl}:\mathit{op}^{*}\mapsto j^{*}\ ]$. This consists of four components: the arguments $a^{*}$ provide the initial value for the handler parameters; the body expression $e$ provides the computation being handled; the return jaxpr $j$ converts the final handler parameters and the body’s result into the final output; and the implementation mapping $\mathit{op}^{*}\mapsto j^{*}$ gives the jaxpr clauses defining the operations.

Throughout this paper, we use an overline to denote sequences; for example, $\bar{a}$ represents the sequence $a_{1},\ldots,a_{n}$ ($n \ge0$), and $\bar{a}:\sigma$ indicates that each $a_{i}:\sigma_{i}$ has the corresponding type $\sigma_{i}$.

Typing

Figure 3 presents the typing rules for $\lambda_{\mathrm{JAXPR}}$. The typing context $\Gamma$ is a finite sequence of variable typings $x:\alpha$ and at most one resumption typing resume: $\sigma\to\tau$. We assume no variable appears twice within a context, and we work up to $\alpha$-equivalence as usual. We use in and out for the input and output types of an operation, and par for the type of a handler’s parameter. We write $\Gamma\vdash\bar{a}:\sigma$ for the component-wise derivation $\Gamma\vdash a_{i}:\sigma_{i}$ for all $i$ (assuming $|\bar{a}|=|\sigma|$).

There are four typing judgments; most rules are straightforward. The judgment $\Gamma\vdash\bar{a} : \alpha$ types atoms. In rule T-Const, primitive constants $c$ have predefined base types $d$ (e.g. $t$, $f : \mathsf{bool}$ and $\star: \star$). Rule T-Array forms a tensor of shape $(n,s)$ from $n$ tensors of shape $s$. The judgment $\Gamma\vdash e : \sigma$ types expressions. Rule T-Let checks that the arguments $\bar{a}$ of a primitive $p : \sigma\to\tau$ have type $\sigma$, and extends the context with $\bar{x} : \tau$ to type the body $e$. Rule T-Ret types the returned atoms. The judgment $\Gamma\vdash j : \sigma\to\tau$ types jaxprs as standard lambda abstractions.

The judgment $\Gamma\vdash p : \sigma\to\tau$ types primitives. Basic primitives $f$ (rule T-Basic) are assumed to have given signatures (for example, tensor concatenation $(::) : d[s],d[ns] \to d[(n+1)s]$). Rule T-Call promotes a jaxpr to a primitive of the same type. Rule T-Cond requires both branches to have the same type $\sigma\to\tau$ and yields type $(\mathsf{bool},\sigma) \to\tau$. Rule T-While follows the standard typing of functional loops, with condition $j_1 : \sigma\to\mathsf{bool}$ and state update $j_2 : \sigma\to\sigma$. Rule T-Scan takes an initial carry of type $\sigma$ and a body $j : (\sigma,d_1[s_1]) \to(\sigma,d_2[s_2])$, and yields type $(d_1[n,s_1]) \to(\sigma,d_2[n,s_2])$.

For effects, each operation is assumed to have a given signature $op : in \to out$. Rule T-Do uses this signature, and requires an argument of type $in$ to produce a result of type $out$. Rule T-Handler types handler expressions: the initial parameters $\bar{a}$ have type $par$, the handled computation $e$ has type $\tau$, and the return clause $j$ maps the final parameters and the result to the overall output type $\tau'$. Each operation clause $j_i$ is typed under a context extended with resume : $par,out_i \to\tau'$, so that resuming takes updated parameters and the operation result and yields the handler’s output. The handler has type $\star\to\tau'$. Rule T-Resume simply retrieves the type of resume from the context.

Operational Semantics

Figure 4 presents the operational semantics of $\lambda^\varepsilon_{\mathrm{JAXPR}}$, which serves as the reference behavior against which we prove that our translation preserves semantics. The top of Figure 4 defines auxiliary notations used in the operational semantics. The application $\{\lambda\bar{x} : \sigma.e\}\bar{a}$ expands to $e[\bar{a}/\bar{x}]$, and $j@\bar{a}_1@\bar{a}_2$ abbreviates $j@\langle\bar{a}_1,\bar{a}_2\rangle$. A context $C$ is either a hole $\square$ or a sequence of let-bindings ending in a hole. Every well-formed expression can be uniquely written as $e=C[\textbf{return }\bar{a}]$. The operation $e_1 \triangleright e_2$

then takes $e_1 = C[\textbf{return }\bar{a}]$, substitutes $\bar{a}$ for the binders $\bar{b}$ in $e_2$, and plugs the result into the hole of $C$. Plugging is capture-avoiding (our mechanization uses a locally nameless representation, and our implementation performs $\alpha$-renaming when needed). This plugging is similar to re-normalization, which is a well-known technique used to keep ANF terms valid during reduction [26].

In rule S-FUNC, we assume each basic primitive $f:\sigma\to\tau$ comes with a total, deterministic evaluation relation $f(\bar{a})\to\bar{a}'$ (e.g. $(l::[l])\to[ll]$), and the results are substituted for the binders $\bar{b}$.

Rule S-CALL evaluates a call. However, reducing $\text{let }\bar{b}=\textbf{call}[\textbf{body}:j]\ \bar{a}\ \textbf{in}\ e$ directly to $\text{let }\bar{b}=j@\bar{a}\ \textbf{in}\ e$ would violate ANF. So instead, rule S-CALL reduces to $j@\bar{a}\triangleright\bar{b}.e$, which runs the jump application and binds its results in the rest of $e$.

The control-flow rules are standard. Rule S-COND reduces a conditional to a call of the branch $j_c$ selected by the boolean $c$. Rule S-WHILE evaluates the condition $j_1$ and, if true, runs one iteration of $j_2$ before calling the while primitive again; otherwise it returns the current state $\bar{a}$. Rules S-SCAN1 and S-SCAN2 evaluate a scan by induction on the input array, running the body on the first element and recursing on the rest.

The final four rules are for handlers. Rule S-HCONG evaluates the handled computation, and rule S-HRET applies the return clause to the current parameters and the result once the computation finishes. For operations, rule S-HDO1 forwards an operation to the handler does not handle, while rule S-HDO2 calls the corresponding clause $j$ with the handler parameters and operation arguments.

_footer: Vol. 1, No. 1, Article . Publication date: November 2026.

Within the clause, resume is replaced by a call whose body $j_k$ captures the continuation following the operation and reinstalls the handler. We inspect $j$ only to determine the parameter type $\tau$ needed to construct $j_k$. Appendix § 10 illustrates these rules with two example programs that use handlers for exceptions and for state.

Metatheory

We prove that $\lambda^{\varepsilon}_{\mathrm{JAXPR}}$ enjoys the following desirable properties:

The third case of progress arises because $\lambda^{\varepsilon}_{\mathrm{JAXPR}}$ does not track effects in types, so a well-typed program may perform an unhandled operation; the theorem nevertheless guarantees that this is the only way it can get stuck. An effect type system [22, 29] would rule this out, but is orthogonal to our focus, and JAX itself does not statically track effects.

Effect-Handling Translation

We compile effect handlers away via a novel translation from $\lambda^{\varepsilon}_{\mathrm{JAXPR}}$ to its handler-free subset $\lambda_{\mathrm{JAXPR}}$, while leaving primitive computations to the hardware. This section presents the core translation without control flow; control flow is addressed in § 5.

The Translation

Figure 5 presents the translation rules. The function $\mathcal{E}\llbracket\mathbb{H} \mid\mathbb{P} \mid\mathbb{K} \vdash e \rrbracket$ translates an expression $e$ with respect to a handler context $\mathbb{H}$ mapping operations to clauses, their parameters $\mathbb{P}$, and a list of continuations $\mathbb{K}$, which are jaxprs in our setting. The list $\mathbb{K}$ contains both return clauses $j_r$ of handlers and normal continuations $k$, which we assume are kept distinguishable. (The special form $\langle k\rangle$ is discussed in § 5.1 and can be ignored for now.) The function maintains the invariant that the length of $\mathbb{H}$ matches both the number of return clauses in $\mathbb{K}$ and the length of $\mathbb{P}$. We assume that the output type $\tau_o$ of the elaboration is given, which will always be provided (e.g. in Theorem 4.3).

The first three rules translate $\mathbf{return}\ \bar{a}$ according to the continuation list. If the list is empty (E-Ret1), the translation is complete. Otherwise, if the innermost continuation is a normal continuation $k$ (E-Ret2), we apply it to $\bar{a}$ and translate the result with $k$ removed. If it is a return clause $j_r$ (E-Ret3), we apply $j_r$ to the handler parameters $\bar{a}_h$ and $\bar{a}$, and continue after popping the innermost handler $h$, its parameters, and its return clause from their lists.

Primitive functions are preserved, since they are to be evaluated on the hardware; thus, rule E-Func keeps the function and translates the body. For call $\mathop{\mathbf{call}}$ and cond, when the handler and continuation lists are empty (E-Call1 and E-Cond1), we simply translate the subexpressions recursively, using the jaxpr translation given at the bottom of the figure. Otherwise (E-Call2 and E-Cond2), the subexpressions may perform operations, so we capture the current continuation $\{\lambda\bar{b}.e\}$, prepend it to the continuation list, translate the subexpressions, apply the result to $\bar{a}$, and return the result bound to fresh variables $\bar{z}$.

Translating a do expression (rule E-Hop) is slightly more involved. First, we look up the operation clause $op \mapsto j$ in the handler $h$, together with its corresponding parameters $\bar{a}_h$ and return clause $j_r$. This lookup relies on the structural property that the length of the handler list $\mathbb{H}'$ equals both the number of return clauses in $\mathbb{K}'$ and the length of the parameter list $\mathbb{P}'$. We then translate $j$ applied to $\bar{a}_h$ and $\bar{a}$ under the lists preceding $h$. The complication is that $j$ may use resume, which

where $j'_1=\mathcal{J}\llbracket\mathcal{H}\mid\mathcal{P}\mid\mathcal{K}\cdot\{\lambda\bar{b}.e\}\vdash j_1\rrbracket$, $j'_2=\mathcal{J}\llbracket\mathcal{H}\mid\mathcal{P}\mid\mathcal{K}\cdot\{\lambda\bar{b}.e\}\vdash j_2\rrbracket$

where $op\mapsto j\in h\land op\notin\mathcal{H}\land|\mathcal{H}|=|\mathcal{K}|=|\mathcal{P}|\land h:p\mathbin{\mathrm{ar}}$

replaced by its continuation. Yet we cannot translate that continuation immediately: resume may be invoked inside $j$ under various contexts, so the context for translating the resumed expression $e$ is unknown at this point. We therefore substitute resume with a call primitive whose body receives the updated parameters $\bar{p}$ and the operation results $\bar{x}$, and returns the special expression $\mathcal{D}\llbracket h\cdot\mathcal{H}'\mid\bar{p}\cdot\mathcal{P}'\mid j_r\cdot\mathcal{K}'\vdash e\rrbracket$, which captures the lists up to $h$ along with the updated parameters. Rule E-HResume then resolves a $\mathcal{D}\llbracket\mathcal{H}'\mid\mathcal{P}'\mid\mathcal{K}'\vdash e\rrbracket$ construct by concatenating the captured lists with the current ones and translating $e$ in the combined context. Therefore, if resume is used several times inside $j$, $e$ is translated under each context as intended.

Finally, E-Handle1 and E-Handle2 translate handle expressions. When the handler list is empty (E-Handle1), we translate $e_b$ with the current handler and bind its result to the translation of $e$ under empty lists. Otherwise (E-Handle2), we push the handler, its parameters, the continuation $\{\lambda\bar{b}.e\}$, and the return clause onto the lists to translate the handled expression.

Elaboration Soundness and Semantics Preservation

First, it is worth noting that our translation is a selective CPS transformation [34]: CPS is performed only when the stack is non-empty.

Lemma 4.1 (Pure programs). If $e$ contains no handler-related constructs, then $\mathcal{E}\llbracket\bullet\mid\bullet\mid\bullet\vdash e\rrbracket=e$.

footer: , Vol. 1, No. 1, Article . Publication date: November 2026.

We then prove a group of mutually recursive structural properties of the elaboration:

Lemma 4.2 (K-family lemmas). (K-push) $\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash e_1\mathbin{\triangleright} b.e_2\rrbracket=\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\cdot\{\lambda b.e_2\}\vdash e_1\rrbracket$. (K-compose) $\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\cdot\{\lambda x.e_1\mathbin{\triangleright}b.e_2\}\cdot\mathbb{K}'\vdash e\rrbracket=\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\cdot\{\lambda b.e_2\}\cdot\{\lambda x.e_1\}\cdot\mathbb{K}'\vdash e\rrbracket$. (K-delay) $\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash\mathcal{D}\llbracket\mathbb{H}\cdot\{\lambda b.e_2\}\vdash e_1\rrbracket\rrbracket=\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash\mathcal{D}\llbracket\mathbb{H}'\vdash e_1\mathbin{\triangleright}b.e_2\rrbracket\rrbracket$.

We now turn to elaboration soundness. First, we define well-formedness, denoted $\Gamma\vdash\mathbb{H}\bowtie\mathbb{P}\bowtie\mathbb{K}:\tau_1\to\tau_2$ as a judgment indicating that the lists take an expression of type $\tau_1$ and produce one with type $\tau_2$. Here, $\Gamma\vdash h:par\Rightarrow\tau$ abbreviates the operation-clause premises of rule T-HANDLER for a handler that takes parameters $par$ and has result type $\tau$. The last premise of WF-CONT ensures continuations always end with a return clause.

(WF-EMPTY)

(WF-CONT)

(WF-HANDLER)

We now establish that our translation preserves typing; a similar statement holds for jaxprs.

Theorem 4.3 (Elaboration Soundness). Given $\Gamma\vdash\mathbb{H}\bowtie\mathbb{P}\bowtie\mathbb{K}:\tau_1\to\tau_{n+1}$, if $\Gamma\vdash e:\tau_1$, then $\Gamma\vdash\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash e\rrbracket:\tau_{n+1}$.

Next, we prove semantics preservation. Semantic preservation proofs usually proceed by constructing a simulation diagram [30]; we first prove that such a diagram exists and commutes:

Lemma 4.4 (Simulation step). If $e\longrightarrow e'$, then $\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash e\rrbracket\longrightarrow^{*}\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash e'\rrbracket$

Using the simulation step, we then prove semantics preservation:

Theorem 4.5 (Semantics preservation). If $e\longrightarrow^{*}_{\mathrm{return}}v$, then $\mathcal{E}\llbracket\bullet\mid\bullet\mid\bullet\vdash e\rrbracket\longrightarrow^{*}_{\mathrm{return}}v$.

Handling Control-Flow Primitives

The translation of the previous section does not account for control-flow primitives such as while and scan. Consider translating $\mathcal{E}\llbracket\mathbb{H}\mid\mathbb{P}\mid\mathbb{K}\vdash\mathbf{let}\ b=\mathbf{while}\ [\mathit{cond}:j_1,\mathit{body}:j_2]\ a\ \mathbf{in}\ e\rrbracket$. Translating the loop body $j_2$ requires extending the continuation list with the current continuation, which contains a copy of the while expression itself. Since continuations are unrolled during translation, we eventually return to translating the same while expression, and the translation fails to terminate! The root cause is that statically eliminating handlers requires statically determining the continuation of each operation. Bounded loops such as scan can be unrolled, but for unbounded loops no valid jaxpr continuation can be constructed. We are unaware of prior work addressing this challenge. Following §2.2, we handle tail-resumptive handlers in-place (§5.1) and abortive handlers by a source-level rewrite (§5.3), and prove the tail-resumptive translation sound (§5.4).

Handling In-Place

The left of Figure 6 defines tail-resumptive expressions ($\lvert e\lvert$), jaxprs ($\lvert j\lvert$), and primitives ($\lvert p\lvert$): an expression is tail-resumptive if it ends in resume, or if it is a let binding whose primitive is tail-resumptive, or whose primitive contains no resume and whose body is tail-resumptive; a jaxpr is tail-resumptive when its body is, and call and cond are when their jaxprs are.

From a tail-resumptive expression, the judgment $\mathcal{R}\llbracket\mathbb{H}\mid\mathbb{P}\mid\omega\vdash e\rrbracket$ extracts the pre-resume computation, which we inline instead of materializing a continuation. It tracks the handler context $\mathbb{H}$, since the expression may perform further tail-resumptive operations, and threads the parameters $\mathbb{P}$, which such operations may update. The type $\omega=\tau,par,out$ annotates generated let-expressions: $\tau$ is the type of $\mathbb{P}$, $par$ the parameter type of the handler evaluated in-place, and $out$ the operation’s

output type. We omit the straightforward rules for jaxprs and primitives. Rule R-Resume returns the current parameters $\mathsf{P}$ alongside the arguments to resume, i.e. the updated parameters and the operation’s output. Rule R-Let1 recurses into the primitive. Rule R-Let2 handles let bindings where body $e$ is tail-resumptive, and primitive $p$ does not mention resume. Here, the body can be inlined recursively, and $p\bar{a}$ is translated under $\mathsf{H}$, allowing for tail-resumptive operations of $\mathsf{H}$. Therefore, $\mathcal{R}\llbracket\cdot\rrbracket$ and $\mathcal{E}\llbracket\cdot\rrbracket$ are mutually recursive. Notably, $p$ may install new handlers whose operations need not be tail-resumptive, as local handlers capture continuations without crossing the loop boundary.

The final continuation. A final subtlety of R-Let2 is that because operations of $\mathsf{H}$ can update the parameters, the translation of $p\bar{a}$ must return the final parameters. To address this, we introduce a new continuation form $\langle k\rangle$, the final continuation, which marks the end of the computation and receives both the current parameters and the result (as in E-Ret4). Accordingly, R-Let2 translates the binding under the final continuation $\langle\{\lambda\bar{y}:\tau,\sigma.\operatorname{return}\bar{y}\}\rangle$, so that the updated parameters are returned, bound to $\bar{z}$, alongside the result $\bar{x}$, and then translates $e$ using $\bar{z}$ as the new parameters.

Handling in-place. Rule E-Hop2 in Figure 7 handles an operation in-place. It looks up the clause $op\mapsto j$, with $j\mid l$, and the parameters $\bar{a}_h$ in the handler $h$, without inspecting the continuations, and forms a call whose body is the inlined resumption of $j$, computed under the handlers $\mathsf{H}$ preceding $h$, the parameters $\mathsf{P}$, and their types. The call threads the updated parameters $\bar{y}$ and returns the output $\bar{x}$; the body $e$ is then translated under the original context with the updated parameters.

Notably, this extended translation relaxes the invariant that the number of return clauses in $\mathbb{K}$ always matches the number of handlers in $\mathsf{H}$ (e.g. in E-Ret4). This is desired: when translating do, if the corresponding handler has no return clause in $\mathbb{K}$, rule E-Hop cannot apply, leaving rule E-Hop2 to handle only tail-resumptive operations.

Translation of while. In rule E-While, the loop’s return type is extended with the parameters; the translated loop takes the current parameters $\mathsf{P}$ and yields the updated ones $\bar{z}$. The condition $j_1 ignores the parameters and translates $e_1$ under empty contexts¹¹; the body $j_2$ translates $e_2$ with the input parameters, using the final continuation to return the updated ones.

An Example

Resumption inlining must preserve handler scoping, which is subtle when a handler’s clause itself performs operations. Consider the following program, where the outermost and innermost handlers handle $op_1$ with clauses $j_1$ and $j_3$, the middle one handles $op_2$, and its $op_2$ clause performs $op_1$:

When the innermost body performs $op_2$, naively taking the pre-resume part of the $op_2$ clause, i.e. the computation before resume, would yield the translation

This changes the program’s semantics: $op_1$ is now handled by the innermost handler ($j_3$), although it was originally performed within the scope of the outermost handler ($j_1$). To prevent this, Xie et al. [52], Xie and Leijen [54] either restrict resumptions to be applied in the very scope of their original handler context, or add a runtime construct that dynamically adapts handler contexts.

Our system needs neither, as inlining naturally preserves handler scopes: because inlining recursively processes let-bindings, it is performed under the handler context that preceded the one being inlined, here $j_1$ (§ 5.1). Instead of the naive pre-resume part, the inlining emits a call to the following translation, so that $op_1$ is correctly handled by $j_1$:

Abortive Handlers

Handling in-place covers tail-resumptive handlers. Abortive handlers never resume, so there is no resumption to inline; instead, we support them by a source-level rewrite of the loop that leaves the translation untouched. Concretely, we support operations $op_i$ (for $i \in1 \ldots n$) with signatures $in_i \to out_i$ whose abortive handlers reside outside the loop by the following rewrite rule; for space reasons, we give the full definitions of the rewritten components in Appendix B.§ 2.

The idea is that the loop records an abortive operation and exits, and the operation is re-performed after the loop. The loop state carries an integer flag $i$ (initially 0) identifying the recorded operation, if any, together with its arguments $\bar{v}_i$ (initially zeros). The body $j'_2$ wraps the original body in a handler for all $op_i$ that discards the resumption and returns the loop state, $i$, and the operation arguments; the condition $j'_1$ additionally checks $i = 0$; and after the loop, a cond on $i$ (where branches abbreviates nested conditions on an integer) returns the loop results if $i = 0$ and otherwise re-performs $op_i$ with the stored arguments. Both this rewrite and the one for effectful conditions (Appendix B.§ 1) are formally proven to preserve typing and semantics; the proofs proceed by a simulation between the source loop and the rewritten one.

Elaboration Soundness and Semantics Preservation

To show that the new translation rules preserve elaboration soundness, we first establish the soundness of resumption inlining:

We then extend our well-formedness relation $\Gamma\vdash\mathsf{H} \bowtie\mathsf{P} \bowtie\mathbb{K} : \tau_1 \to\tau_2$ with a fourth case for the final continuation $\langle k\rangle$: The final continuation may only occur as the leftmost element of $\mathbb{K}$: it is the last continuation to run, as it ends the computation:

That is, $k$ receives the current parameters $\overline{\mathit{par}}_i$ alongside the result of type $\tau_1$, and produces the final result of type $\tau_2$. Using this extension and Lemma 5.1, elaboration soundness (Theorem 4.3) extends to the new translation rules.

Next, we turn to semantics preservation. The simulation diagram in § 4.2 says the following: if $e \longrightarrow e'$, then the translated program should take one or more steps, $\llbracket e \rrbracket\longrightarrow^* \llbracket e' \rrbracket$. For CPS transformations, however, this diagram notoriously often doesn’t commute. In particular, administrative redexes introduced by CPS change the syntactic shape of the target program, so target reductions may not correspond to source reductions. We avoided this issue in § 4.2 because there the translation generates no administrative redexes, and when it changes the shape of a program by pushing a continuation inward, the K-family lemmas (Lemma 4.2) restore syntactic equality. Resume inlining, however, is different: rather than following the source semantics of handlers, which resumes a captured continuation, it executes the handler’s body in place. We therefore relate the inlined program to the CPS translation, which is already proven to match the source semantics. Because the two differ by administrative redexes, the simulation diagram of § 4.2 no longer holds. Plotkin [39] handles this issue with a colon translation, which contracts administrative redexes at the head of a term as it translates. More recent work [35] instead uses a logical relation to show that, although the syntactic shape of the translated program changes, the programs remain semantically equivalent. Our simulation proof combines both: a logical relation captures semantic equivalence, and, in the spirit of Plotkin’s colon translation, our simulation relation accounts for the administrative redexes that separate the two programs.

First, we adapt this logical relation to our setting:

The expression relation says that two expressions are related when every final value produced by the first is the same value produced by the second. A similar relation holds for jaxprs.

Simulation between translations. Since the tail-resumptive translation differs syntactically from the core translation, we start by relating the two by a simulation $e_1 \preceq_j e_2$ (Figure 8) between a program $e_1$ on the core side and a program $e_2$ on the tail-resumptive side, modulo a pending continuation $j$.

!Simulation relation $e \simeq^{j} e$; Sim-Resume: $j_1@\bar{a} \simeq^{j} \operatorname{return}\bar{a}$; Sim-Admin: $\operatorname{let}\bar{x}:\sigma=\operatorname{call}[\operatorname{body}:j_1]\bar{a}\operatorname{in}\operatorname{return}\bar{x}\simeq^{j}e$; Sim-Inline: $e\simeq^{j}j_1@\bar{a}\triangleright\bar{x}.e_2$, $e\simeq^{j}\operatorname{let}\bar{x}:\sigma=\operatorname{call}[\operatorname{body}:j_1]\bar{a}\operatorname{in}e_2$; Sim-Prefix: $e_2\simeq^{j}e_4$, $e_1\triangleright\bar{x}.e_2\simeq^{j}e_1\triangleright\bar{x}.e_4$; Sim-Compose: $e_2\simeq^{j}e_4$, $e_1\simeq^{j}e_3\triangleright\bar{x}.e_4$, $\{\lambda\bar{x}:\operatorname{out}.e_2\}$, $e_1\simeq^{j}e_3\triangleright\bar{x}.e_4$ bbox=0.09,0.11,0.91,0.26

The pending continuation $j$ is what the core side still runs once the tail-resumptive side returns, so Sim-Resume relates $j_1@\bar{a}$ to $\operatorname{return}\bar{a}$. The two sides may be unfolded to different degrees: Sim-Admin lets the core side wrap a continuation in an administrative call, and Sim-Inline lets the tail-resumptive side keep a call in place of its inlined body. Sim-Prefix is congruence under a common prefix, and Sim-Compose handles sequencing: when the tail-resumptive side continues with $e_4$, the core side instead runs its continuation $e_2$, so the prefixes are related with $e_2$ pending. Like Plotkin’s colon translation, this relation bridges the syntactic gap that administrative redexes create. Plotkin’s translation removes the redexes at the head of a term; our relation keeps them and describes how the terms differ.

The relation is preserved by reduction:

When the tail-resumptive side returns, the core side runs to the pending continuation applied to the same atoms: if $e_1 \simeq^{j} \operatorname{return}\bar{a}$, then $e_1 \longrightarrow^{*} j@\bar{a}$. At the identity continuation $id=\{\lambda\bar{x}.\operatorname{return}\bar{x}\}$ this is $\operatorname{return}\bar{a}$ itself. Therefore:

Semantics Preservation. Because resumption inlining is an instance of the simulation at the identity continuation, by Lemma 5.5 we obtain:

Using Lemma 5.6, we update the simulation step lemma (Lemma 4.4) with the logical relation: each source step is matched up to $\preceq$.

Since $\preceq$ is transitive (Lemma 5.2), this extends from single steps to whole runs: a closed program and its translation return the same values.

Implementation

We implement effex as a JAX library, leveraging our effect-handling extension to jaxprs.

Programming with effecax. JAX offers a NumPy-like interface [19] together with composable transformations. effecax extends this interface with effects and handlers, as Figure 9a illustrates. Effects are declared with their type signatures (via abstract_eval) and invoked like ordinary functions, as get and set in double_mul. A handler is installed with handle, which takes the handled computation, an initial parameter, a default return clause, and the operation cases. Each case receives the handler parameter, the operation arguments, and the continuation k, which resumes the computation with the updated parameter and the operation’s result.

Programs are traced into our extended jaxprs by JAX’s tracing infrastructure [14]. For example, tracing double_mul_handled yields the jaxpr in Figure 9b, which contains the effect primitives delimited_handler, resume, and effect. Like all jaxpr primitives, they are statically typed; each carries a PyTree field describing the structure of its shape-dependent arguments.

Compiling effect handlers. We implement the translations described in §4 and §5 that compile our extended jaxprs to standard JAX jaxprs. Our implementation implements $\alpha$-conversion and capture-avoiding substitution to systematically rename variables and ensure freshness. In-place handling also serves as an automated optimization for non-iterative code. For instance, because the handlers in Figure 9b are tail-resumptive, they are inlined during the translation to produce a more compact program shown in Figure 10.

Case Studies

We explore the usage of effecax across two real-world domains in the JAX ecosystem: numerical differential equation solving (§7.1) and probabilistic programming (§7.2). We select these domains because they represent common JAX use cases, both benefit from incorporating effects, and both feature mature, established JAX libraries that serve as baselines for comparison. A complete overview of these domains is beyond our scope; we refer interested readers to Kidger [27] and van de Meent et al. [50] for further background. Using these case studies, we evaluate effecax through the following research questions:

• RQ1: Can effecax express effect patterns that are awkward to write in today’s JAX?

• RQ2: Does effecax compose with JAX’s built-in primitives, such as control flow?

• RQ3: Does expressing effect patterns with effecax come at a performance cost?

Throughout this section, all runtimes are measured on an NVIDIA A100 (SXM4, 40 GB) GPU; compilation (Figure 12 and Table 2) is host-side work and runs on the machine’s CPU, a 12-vCPU Intel Xeon (Cascade Lake, 2.20 GHz) host with 85 GB RAM.

	EX (s)	OX (s)	NP (s)
Proportion	1.37 (0.05 + 1.32)	1.82 (0.12 + 1.69)	1.34 (0.04 + 1.30)
Stoch. volatility	0.48 (0.09 + 0.39)	0.90 (0.19 + 0.70)	0.59 (0.10 + 0.48)
Bayesian NN	0.48 (0.13 + 0.36)	0.78 (0.19 + 0.59)	0.46 (0.09 + 0.37)
Semi-sup. HMM	1.44 (0.20 + 1.25)	1.45 (0.15 + 1.30)	1.98 (0.72 + 1.27)
UCBadmit GLMM	0.88 (0.15 + 0.73)	0.87 (0.07 + 0.80)	1.37 (0.25 + 1.12)
HSGP	1.64 (0.31 + 1.33)	1.80 (0.26 + 1.54)	1.81 (0.31 + 1.49)

Optimizers and Solvers

Given a vector field $f$ and an initial value $y_0$, a solver numerically approximates the trajectory $y(t)$ where $\mathrm{d}y/\mathrm{d}t = f(t,y)$. The core is an integration loop: starting from $y_0$, it repeatedly calls a one-step method to advance the solution over time. In real-world solvers, the loop is a complex imperative program. In JAX, Diffrax¹² implements this loop using ad-hoc mechanisms described in § 2.1. In particular, Diffrax’s integration core¹³ carries a single record of 20 variables serving orthogonal purposes. The structure of this core loop is sketched in Figure 11a.

We contribute a refactor of Diffrax’s core integration module using effecax.

Expressiveness. With effecax, each of these concerns becomes an independent handler and the loop body is written in direct style. The 20-field record is replaced by one state handler per concern; taking a step and adapting the step size are effects (try_step, adapt) interpreted by handlers; and early termination is an exception (fail) rather than a flag threaded through the state and re-checked in the loop condition. The resulting program is written modularly in direct-style and results in a $\approx22\%$ shrinkage in code size (measured by lines of code). The same loop written in effecax is shown in Figure 11b. With effecax, parameterized handlers allow us to update the state in-place, without having to worry about threading it correctly. Moreover, the loop doesn’t take in parameters itself, and the loop’s conditional doesn’t directly check the results flag, rather a single fail() operation immediately aborts the loop. With effecax, the logic can be seen directly without any obfuscation behind state threading (RQ1). Note also that the effectful loop body sits inside the standard while_loop and cond (RQ2).

Performance. We compare the performance of our refactor against Diffrax. From Diffrax’s default benchmark suite, we select all programs that exercise the integration loop. Each sweeps one parameter, plotted on the x-axis:

• (a) a reaction-diffusion PDE on an $N \times N$ grid, sweeping the grid size $N$ (increasing grid size, hence per-step work);

• (b) a small neural ODE, whose vector field is a small MLP (batch size 1024, i.e. each solver step runs the MLP on 1024 states), sweeping the integration length: a longer integration takes more solver steps, hence more work per step;

• (c) a jump-step SDE (batch size 1024, i.e. each solver step advances 1024 independent SDE paths), likewise sweeping the integration length;

• (d) the classic Lotka-Volterra system (a tiny two-variable ODE) as a work-precision study [44] (increasing the error tolerance and measuring how long the solve takes).

Figure 12 demonstrates the effecax refactor is competitive on the three substantial workload benchmarks (a)-(c) (RQ3). The right panel of Figure 12 shows the work-precision plot of the Lotka-Volterra benchmark (d). As the error grows ($x$-axis), the runtime falls ($y$-axis) since a looser tolerance does less work. The two implementations roughly match across the entire sweep. Lotka-Volterra is a two-variable system whose steps do minimal numerical work, so the runtime is dominated by the integration-loop machinery itself which is precisely where a per-step overhead of the effect-handler compilation would show up if there were one (RQ3).

Finally, we compare compilation time (Figure 12, bottom right) from Diffrax’s compile-time benchmark. The benchmark deliberately stresses the compiler: it is written in such a way that produces large jaxprs. For Diffrax, this is the cold-start cost of the first call: tracing its (larger) integrator program to a jaxpr and letting XLA compile it. For effecax, it additionally includes an ahead-of-time stage that compiles the handlers away (tracing occurs before this step, XLA compilation after). Despite this extra stage, the effecax refactor compiles faster on the forward case and matches Diffrax on the forward with gradient case: the effect handler compilation pass is cheap and Diffrax spends longer tracing since it has a larger program.

Probabilistic Programming

Probabilistic models describe a joint distribution $p(x,z)$ over a space of latent variables $z$ and observations $x$. Bayesian inference computes the posterior $p(z \mid x) = p(x,z)/p(x)$. The marginal likelihood $p(x) = \int p(x,z)\,\mathrm{d}z$ quickly becomes intractable in higher dimensions of latent variables, so inference algorithms approximate the posterior instead. Probabilistic programming languages (PPLs) [50] express models with sample (drawing from a distribution) and observe (conditioning on data) operations, and run inference to learn model parameters from observed data.

We contribute a new PPL library in effecax via inference patterns [32, 33]: an inference pattern provides inference operations, whose semantics is provided by handlers, thus achieving modularity and extensibility.

Expressiveness over NumPyro and Oryx. We compare with two representative JAX-based PPL systems: NumPyro [37] and Oryx [51]. Both implement effects based on ad-hoc mechanisms similar to those described in § 2.1, which impose certain restrictions. NumPyro implements effects using a global state mechanism (“messengers”)¹⁴ outside of JAX, so effects break within JAX control flow. NumPyro therefore introduces and requires users to use its specialized versions of cond and scan, and does not support while. Oryx implements effects using a dedicated JAX transformation for each effect. In particular, its harvest transformation¹⁵ is limited to state-like effects and currently does not support control flow primitives.

Figure 13 shows a Hidden Markov Model (HMM), whose hidden states evolve by a Markov process and are naturally expressed as a loop, in NumPyro and in effecax (full definitions in the artifact). NumPyro must use its specialized scan; effecax uses the standard JAX scan (RQ2).

These restrictions can also cause a loss of expressiveness in NumPyro and Oryx. Figure 14 shows two practical models. (1) A Hierarchical HMM (left) is naturally written with nested scans. However, NumPyro rejects nested scan with an error, and Oryx’s harvest is incompatible with JAX control flow entirely, rejecting any latent variable inside a lax.scan. (2) In a Regime-Switching SSM (right), NumPyro’s cond requires all sample sites to share one distribution class and rejects the model at trace time, while Oryx silently drops the latent variable inside lax.cond from the joint density, giving wrong results without warning. effecax expresses both models directly and correctly (RQ1, RQ2). Users can often restructure such models around these restrictions, but doing so sacrifices compositionality and hand-encodes the error-prone interaction between control flow and effects.

Performance. We compare effecax against NumPyro and Oryx across both runtime cost and compile-time cost using Hamiltonian Monte Carlo (HMC) inference. HMC is an ideal benchmark because its core loop repeatedly exercises the target log-density gradient, which is exactly the computation generated by each framework’s side-effect abstraction, and because it is the standard default sampler in NumPyro. We measure compile time for the log-density function, and report the total runtime spent running the log-density program over a full inference run, across 100 000 evaluations. We run these benchmarks over a wide range of real-world probabilistic models using varying large dataset sizes¹⁶: (a) proportion test logistic regression, (b) stochastic-volatility state space, (c) Bayesian neural network, (d) scan-marginalized semi-supervised HMM, (e) hierarchical GLMM based on UCBadmit dataset, (f) Hilbert-space GP approximation.

Table 1 reports each framework’s total time over that run, alongside a hand-written baseline; in both tables, the shaded cell marks the fastest entry per row. The result shows that recovering the log-density computation via effect handlers, harvest, or messengers adds minimal runtime overhead, matching the baseline within a few percent in most cases (RQ3). Table 2 shows the compilation time. It measures the cost of running each library’s side-effect elimination strategy to produce the log-density computation and then XLA-compiling it. Entries are formatted as $z$ ($x + y$) seconds, separating the framework’s effect-elimination transform time ($x$) and XLA compilation time ($y$) from their total ($z$). On half of the models, effecax compiles fastest of the three frameworks, and is within a few percent of the leader on the remaining ones (RQ3).

Model	Size	HW (s)	EX (s)	OX (s)	NP (s)
Proportion	3M	10.25	17.78	11.56	17.73
Stoch. volatility	2.5k	1.75	1.36	2.72	1.99
Bayesian NN	50k	8.97	11.52	9.11	8.99
Semi-sup. HMM	5k	9365	9526	9423	9434
UCBadmit GLMM	49k	6.66	6.67	6.65	7.29
HSGP	7.3k	5.24	4.20	5.23	4.93

Discussion and Related Work

Compiling functional languages to accelerators. GPUs offer a restricted programming model and lack support for function pointers and dynamic recursion. As a result, compiling higher-order functions typically requires defunctionalization [45] to produce first-order functions, a transformation that replaces each application with a call to an apply function that branches across all possible function forms. However, as Hovgaard et al. [24] observe, naive defunctionalization introduces the issue of branch divergence on GPUs, where extensive branching leaves many threads idle while a subset executes instructions, leading to suboptimal performance. To address this, Hovgaard et al. [24] introduced a specialized defunctionalization transformation for the Futhark language [20, 21]. They restrict the use of functions, specifically disallowing loops and conditionals from returning functions and forbidding arrays from containing functions, a design that brings the language remarkably close to jaxprs. Under these constraints, Futhark’s defunctionalization transformation specializes each application to the particular form of function that may occur at runtime. As Hovgaard et al. [24] note, “the result is essentially equivalent to inlining completely the apply function in a program produced by Reynolds defunctionalisation.” Our translation plays the same role for effect handlers: we conjecture that it yields the program one would obtain by CPS-transforming the handlers, defunctionalizing the resulting higher-order functions, and fully inlining the apply function. It does so directly, in a single pass over the first-order program, without materializing the higher-order intermediate or inducing branch divergence.

The array language Dex [36] similarly uses a simplification pass that leaves only first-order, top-level functions that are fully applied at their call sites, provided that the overall program is of a non-functional type. Both Futhark and Dex normalize higher-order programs, whereas our transformation operates on first-order jaxprs and normalizes only effect handlers. Dex supports Accum, an effect for associative accumulation in parallel computations, and later generalizes this concept to parallel effect handlers [53], though this feature remains yet unimplemented. Each parallel handler must provide a traverse clause that defines the semantics of operations executed in parallel. Nevertheless, this strategy applies only to specific classes of effects, such as accumulation. Also, their parallel exception handler must finish parallel execution before checking for exceptions.

Effect handlers. Existing work on effect handler compilation [23, 47, 54] commonly employs a higher-order CPS transformation or targets pure lambda calculi. In contrast, our work targets a restricted, first-order language without first-class functions or recursion. While prior work has addressed the mechanization of CPS compilation (e.g. [35]), to our knowledge, no prior work has mechanized a semantics-preserving CPS transformation for effect handlers. Schuster et al. [46] explicitly note that they did not mechanize their CPS semantics-preserving result as their “mechanized translation introduces administrative reductions”. Prior work [5, 31, 32, 33, 37] has explored specialized effect handlers for PPL, but these implementations are largely specialized to this specific domain, while effecax provides a natural interface for integrating with JAX.

Limitations. We currently restrict handlers across loop boundaries to be tail-resumptive or abortive, so a clause that e.g. resumes multiple times is rejected. Capturing continuations in a loop requires materializing continuations as data and turning the loop body into a state machine, following Ikebuchi et al. [25] for coroutines, though extending this approach to multi-shot continuations remains non-trivial. However, dispatching over states in every iteration would reintroduce the branch divergence our translation avoids. A further obstacle is space. Non-tail resumptive handlers require space proportional to the number of times a loop runs. Unfortunately, this space requirement cannot be determined at compile time and XLA requires static bounds on memory usage for jaxprs. A similar restriction arises in reverse-mode automatic differentiation of while loops, which requires storing intermediate states of statically unknown memory bounds.¹⁸ We leave addressing this limitation to future work.

Beyond JAX. Our work is motivated by the practical constraints of JAX, in the same way that effect handler implementations for e.g. OCaml, WebAssembly, and C/C++ are shaped by the constraints of their targets. The core ideas, however, apply to any intermediate representation that operates under similar hardware-imposed restrictions. In addition to Futhark and Dex mentioned above, PyTorch has a similar story: its functorch component requires functional code, which PyTorch obtains through a custom functionalization pass¹⁹ that removes state mutation, much like JAX’s Ref transformation. This pass has had a number of bugs since its introduction,²⁰ and we believe effect handlers offer a more principled treatment of side effects in this setting. Lastly, by establishing a framework for effect handlers in JAX, we hope to open new avenues for applying existing reasoning and verification techniques of effect handler programs [4, 11, 12, 49] to this domain.

References

Example Programs

This appendix walks through two small programs, one using a handler for exceptions and one for state, showing for each how it evaluates under the operational semantics of § 3 and how it is translated by the rules of § 4. Throughout, $j_r$ denotes the handler’s return clause and $h$ its operation clauses.

Exceptions

Given an exception operation $\operatorname{throw} : \star\to\operatorname{int}$, safe_div takes two integers and divides them if the divisor is non-zero, and throws otherwise:

safe_div = { λm : int, n : int. let b : bool = eq n 0 in
           let r : int = cond [then: { λx : int, y : int. let z : int = do throw ⋆ in return z },
                                else: { λx : int, y : int. let z : int = div x y in return z }] b m n
           in return r }

We apply it to 42 and 0 under a handler whose throw clause $j$ resumes with the fallback value 0; the handler takes no parameter, and its return clause $j_r$ drops the parameter:

let v : int = handle [ args: ⋆,
                       body: let w : int = call [body: safe_div] 42 0 in return w,
                       ret: j_r = { λx : int. x in return y },
                       impl: throw ↦ j = { λx : ⋆, y : ⋆. let u : int = resume ⋆ 0 in return u }] ⋆
             in return v

Evaluation. Evaluation proceeds inside the handler body (S-HCONG). The call is inlined with $m = 42$ and $n = 0$ (S-CALL), eq evaluates to 1 (S-FUNC), and the then branch is selected and inlined (S-COND), so that, up to trivial returns, the body becomes let z : int = do throw ⋆ in return z. The operation is now performed directly under a handler that handles it, so S-HDO2 applies: it runs the clause $j$ with the handler parameter and the operation argument, replacing resume by a call to the captured continuation $j_k$, which reinstalls the handler around the rest of the body:

Inlining $j$ (S-Call) yields let $u$ : int = call [ body: $j_k$ ] $\star$ 0 in return $u$, and inlining $j_k$ reinstalls the handler with $z = 0$. Its body has returned, so S-HRet applies the return clause, and S-Call finishes:

The program’s result is 0.

Translation. We translate the program under empty lists. Rule E-Handle1 translates the handler body under the handler, with $h = \text{throw} \mapsto j$, continuation list $j_r$, and parameter $\star$, and binds the result to $v$ for the translation of return $v$ under empty lists, which is return $v itself (E-Ret1). The interesting part is therefore

By E-Call2, we capture the continuation $k_1=\{\lambda w:\text{ int}.\text{ return }w\}$ and translate the jaxpr safe_div under $j_r\cdot k_1$, producing let $z_1$ : int = call [ body: $\mathcal{J}\llbracket h\mid\star\mid j_r\cdot k_1\vdash\text{ safe\_div}\rrbracket$ ] 42 0 in return $z_1$. Rule E-JaxpR recurses into the body of safe_div; E-Func keeps the primitive eq; and E-Cond2 captures $k_2=\{\lambda r:\text{ int}.\text{ return }r\}$ and translates both branches under $j_r\cdot k_1\cdot k_2$. The else branch performs no operation: E-Func keeps div, and return $z$ is translated by popping $k_2$ and $k_1$ (E-Ret2, twice) and then $j_r$ (E-Ret3), each application reducing to return $z$, so the branch is unchanged. The then branch is where the handler is compiled away:

By E-Hop, we look up the clause $j$ for throw in $h$, replace resume in it by a call $p=\text{call [ body: }\{\lambda p:\star,z:\text{ int}.\mathcal{D}\llbracket h\mid p\mid j_r\cdot k_1\cdot k_2\vdash\text{ return }z\rrbracket\}\text{ ]}$ whose body delays the translation of the rest of the program, and translate the clause applied to the handler parameter and the operation argument under the lists preceding $h$, which are empty:

The steps are E-Call1 (the lists are empty, so the call is kept and its body translated), E-HResume (the delayed translation resumes under the captured lists), and finally the return rules (E-Ret2 twice and E-Ret3), which unroll $k_2$, $k_1$, and $j_r$ to return $z$. The operation has been eliminated: the call to $p$ simply yields the fallback value. The complete translation is

which is a vanilla jaxpr evaluating to 0, as the source program does.

State

Consider state operations $\operatorname{get} : \star\to\operatorname{int}$ and $\operatorname{set} : \operatorname{int} \to\star$. The handler below carries the state as its parameter, initially 1. Its get clause resumes with the current state both as the new parameter and as the operation result; its set clause resumes with the operation argument as the new parameter and $\star$ as the result; and its return clause returns the final state alongside the body’s result. The body reads the state, doubles it, stores it back, and returns 42:

Evaluation. The body immediately performs get, so S-Hdo2 runs the clause $j_g$ with the parameter 1 and the argument $\star$, with $\mathsf{resume}$ replaced by a call to the continuation

where $e_1$ is the rest of the body. Inlining $j_g$ and then $j_{k_1}$ (S-Call, twice) reinstalls the handler with parameter 1 and $x=1$. Inside it, add evaluates to 2 (S-Func, under S-Hcong), and the body performs set 2. Again S-Hdo2 runs $j_s$ with parameter 1 and argument 2; its $\mathsf{resume}\ 2\ \star$ reinstalls the handler with parameter 2 around the remaining body $\mathsf{return}\ 42$. That body has returned, so S-HRet applies $j_r$ to 2 and 42, and the program evaluates to 2, 42: the final state and the body’s result.

Translation. As before, E-Handle1 translates the body under $h=\operatorname{get}\mapsto j_g,\operatorname{set}\mapsto j_s$, the continuation list $j_r$, and the parameter 1, and binds the result to $v$. The body starts with get, so E-Hop applies: $\mathsf{resume}$ in $j_g$ is replaced by $p_1=\mathsf{call}\ [\mathsf{body}:\{\lambda p:\operatorname{int},x:\operatorname{int}.\ \mathcal{D}\llbracket h\mid p\mid j_r\vdash e_1\rrbracket\}]$, and $j_g[p_1/\mathsf{resume}]$ applied to 1 and $\star$ is translated under empty lists. By E-Call1 the call is kept and its body translated, and by E-HResume that body becomes $\mathcal{E}\llbracket h\mid p\mid j_r\vdash e_1\rrbracket$: the rest of the body, now translated with the updated parameter $p$ in place of 1. Rule E-Func keeps add, and the set operation is handled by E-Hop in the same way, with $p_2=\mathsf{call}\ [\mathsf{body}:\{\lambda p':\operatorname{int},z:\star.\ \mathcal{D}\llbracket h\mid p'\mid j_r\vdash\mathsf{return}\ 42\rrbracket\}]$ and the clause $j_s[p_2/\mathsf{resume}]$ applied to $p$ and $y$. Its delayed body $\mathcal{E}\llbracket h\mid p'\mid j_r\vdash\mathsf{return}\ 42\rrbracket$ is resolved by E-Ret3: the return clause is applied to the current parameter $p'$ and the result, giving $\mathsf{return}\ p'\ 42$. Assembling the pieces, the complete translation is

Both handlers and both operations are gone. The state is threaded explicitly: the outer call receives the initial state 1 twice, as the parameter and as the result of get; the inner call receives the doubled value 2 as the new parameter, which the inlined return clause pairs with 42. The jaxpr evaluates to 2, 42, as the source program does.

B Source-Level Rewrites of While Loops

Effectful Conditions

The translation of while in Figure 7 forbids operations of outer handlers in the loop condition. This restriction can be lifted by moving the condition into the loop body, with an additional boolean carry $r$, through the source-level rewrite

with the rewritten condition $j'_1$ and body $j'_2$ defined as follows:

The new condition $j'_1$ returns the boolean carry $r$. The new body $j'_2$ evaluates the original condition $j_1$ on the state $\bar{y}$ and branches on the result: if true, it runs the original body $j_2$ and returns $t$ together with the new state; otherwise it returns $f$ with the current state, which makes $j'_1$ terminate the loop at the next iteration.

Abortive Handlers

Abortive handlers never resume, so an operation performed inside a loop and handled abortively outside it needs no continuation: the loop can simply stop and the operation be re-performed after it. § 5.3 supports this by a source-level rewrite, applied before translation, for operations $op_i$ (for $i \in1 \ldots n$) with signatures $in_i \to out_i$ whose abortive handlers reside outside the loop:

The loop state is extended with an integer flag $i$, initially 0, that records which operation aborted the loop, if any, together with slots for its arguments, initially zeros. After the loop, a cond on the flag either returns the loop results or re-performs the recorded operation with the stored arguments, where branches abbreviates nested conditions on an integer. The rewritten components are defined as follows:

The extended state type $\sigma'$ adds to the original loop state $\sigma$ an integer flag $i$ and the argument types $\overline{in_i}$ of every operation. The condition $j'_1$ evaluates the original condition $j_1$ on the original state $\bar{z}$ and conjoins it with the check $i = 0$. The body $j'_2$ runs the original body $j_2$ under a handler with a unit parameter. Its return clause passes the updated state through with $i = 0$ and zeroed argument slots. Its clause for $op_i$ discards the resumption and returns the current state $\bar{z}$, the identifier $i$, and the operation arguments $\bar{v}$ placed in the $i$-th slot, with all other slots zeroed. Finally, the branch jaxprs $\bar{j}_i$ are selected on the flag after the loop: $j_0$ returns the loop results, while $j_i$ (for $i \ge1$) re-performs $op_i$ with the stored arguments $\bar{v}_i$.

Elaboration Soundness and Semantics Preservation

We write $\rightsquigarrow_1$ for the effectful conditions rewrite of §B.1 and $\rightsquigarrow_2$ for the abortive handlers rewrite of §B.2. We prove that both rewrites preserve typing.

The rewrites also preserve semantics: whenever the original loop returns, the rewritten loop returns the same atoms. For the rewrite in §B.1, we first start developing a syntactic relation (on the source program) that captures how the rewrite changes the syntax of the program ($\simeq_{\mathit{cond}}$). We then show that this relation is closed under step:

Using Lemma B.3, we then show semantics preservation:

For the semantics preservation of §B.2, we also develop a syntactic relation on the source program that shows how the rewrite changes the program shape ($\simeq_{\mathit{abort}}$). We then prove step simulation:

Using Lemma B.5, we can prove semantics preservation. Since the abortive rewrite is a peephole rewrite, thus such a contextual program could end in an unhandled operation. This operation could be an abortive do or not. The theorem therefore says that if they the step results in a do then the continuation of the do should also be in the relation, if it is not abortive. If it is abortive, it doesn’t matter because the continuation is aborted.

The abortive rewrite applies to a loop nested anywhere in a program, for instance inside the body of another loop or of a handler. We therefore state its semantics preservation for the abortive rewrite congruence $e_1 \rightsquigarrow^{*}_2 e_2$ of Figure 15, which applies $\rightsquigarrow_2$ at any number of positions in $e_1$. The judgment is parameterized by the set $A$ of abortive operations, whose handlers never resume, and rule Rw-Loop only rewrites loops whose $\mathit{op}_i$ are in $A$. Rule Rw-Loop states the rewrite of a loop $\mathop{\mathit{let}}\ \bar{x} : \sigma= \mathit{while}\ [\mathit{cond}: j_1,\mathit{body}:j_2]\ \bar{a}\ \mathit{in}\ \mathit{return}\ \bar{x}$ without its continuation: the rewritten loop ends in the dispatch $e_r$ on the flag, which is then composed with the original continuation $e$. The remaining rules are congruences. Rule Rw-Call rewrites a called jaxpr; the rules for the jaxprs of cond, while and scan are similar. rule Rw-Handle rewrites the body of an abortive handler $h$, i.e. one whose clauses never resume. rule Rw-Handle also has other variant which can rewrite the return and operation clauses of such handlers.

Using the lemmas Lemma B.5 and theorem B.6, we then prove semantics preservation for the contextual equivalence:

Mechanization

Table 3 lists the Lean counterpart of each result in the paper.

Result	Paper Name	Lean Name	File	Line
Lemma 3.1	Determinacy	determinacy	MiniJaxprH	2976
Theorem 3.2	Preservation	preservation	MiniJaxprH	8419
Theorem 3.2	Progress with effects	progress	MiniJaxprH	8130
Definition	$\mathcal{E}\llbracket\cdot\rrbracket$	translateExpr	Translation	213
Definition	$\mathcal{J}\llbracket\cdot\rrbracket$	translateJaxpr	Translation	456
Lemma 4.1	Pure programs	pure_programs	Kfamily	13
Lemma 4.2	K-push	K_push	Kfamily	3418
Lemma 4.2	K-compose	K_compose	Kfamily	3465
Lemma 4.2	K-delay	K_delay	Kfamily	3518
Theorem 4.3	Elaboration soundness	elaboration_soundness	ElaborationSoundness	1070
Lemma 4.4	Simulation step	simulation_step	SemanticsPreservation	3200
Theorem 4.5	Semantics preservation	semantics_preservation	SemanticsPreservation	4711
Definition	$\mathcal{R}\llbracket\cdot\rrbracket$	RExpr	Translation	463
Lemma 5.1	$\mathcal{R}\llbracket\cdot\rrbracket$ soundness	resumption_inlining_typed	ElaborationSoundness	1365
Definition	Logical relation	Refines	Simulation	121
Lemma 5.2	Reflexivity	Refines_refl	Simulation	124
Lemma 5.2	Transitivity	Refines_trans	Simulation	127
Lemma 5.3	Compatibility (Bind)	Refines_bind	Simulation	305
Lemma 5.3	Compatibility (Let)	Refines_let	Simulation	2455
Definition	Simulation relation	Rsim1	Simulation	409
Definition	Simulation relation (top level)	Rsim0	Simulation	456
Lemma 5.4	Step matching	sim_step	Simulation	1595
Lemma 5.5	Adequacy of $\simeq$	simulation_preserves_semantics	Simulation	2372
Lemma 5.6	Resumption inlining preserves semantics	resumption_inlining_preserves_semantics	SemanticsPreservation	1951
Theorem 5.7	Simulation step	simulation_step	SemanticsPreservation	3200
Theorem 5.8	Semantics preservation	semantics_preservation	SemanticsPreservation	4711
Theorem B.1	Effectful conditions: typing	translateWhile_preserves_typing	EffectfulConditional	436
Definition	$\simeq_{\mathrm{cond}}$	WSim	EffectfulConditional	784
Lemma B.3	Effectful conditions: step	WSim_step	EffectfulConditional	1242
Theorem B.4	Effectful conditions: semantics	translateWhile_semantics_preservation	EffectfulConditional	1756
Theorem B.2	Abortive handlers: typing	translateWhile_preserves_typing	Abortive	1619
Definition	$\simeq_{\mathrm{abort}}$	ASimCong	Abortive	5998
Lemma B.5	Abortive handlers: step	ASimCong_step	Abortive	6904
Theorem B.6	Abortive handlers: semantics	translateWhile_preserves_terminals	Abortive	5858
Definition	\❀^{*}	RewriteExprAnywhere	Abortive	508
Theorem B.7	Rewrite congruence: semantics	RewriteExprAnywhere_star	Abortive	7522