Design Space 02: Differentiation Through Effects

Series: AI-Native Programming Language Research — Design Space Memos Document 2 of 6 (planned: shapes, AD-through-effects, sparsity, distribution, verification, hardware abstraction)

1. Problem statement

Reverse-mode automatic differentiation over pure mathematical functions is well-understood and well-implemented. The semantics are clean (chain rule), the algorithms are textbook (Pearlmutter & Siskind, 2008; Elliott, 2018), and production systems (PyTorch autograd, JAX, Zygote.jl) handle the pure-functional case correctly.

Real machine-learning programs are not pure mathematical functions. They mutate buffers, sample from distributions, branch on data, iterate to convergence, raise exceptions on numerical error, do I/O, and call into black-box code. Each effect breaks AD's clean compositional story in a different way, and the existing systems handle each effect as a separate special case. Mutation is handled by tape-based capture. Randomness is handled by manually picking a gradient estimator. Implicit definitions are handled by a separate implicit_diff library. Control flow is handled by graph-capture heuristics.

This is a textbook situation where a unifying framework would replace ten special cases with one. Algebraic effects with handlers is the strongest candidate for that framework, and the relationship between effects and AD is deeper than is generally recognized — reverse-mode AD is itself a CPS transform, and CPS is the implementation strategy for effect handlers. They are, in a real sense, the same machinery.

This memo: surveys the effects that break naive AD, compares how existing systems patch each one, sketches the effect-handler unification, and identifies what's open.

2. Background: where pure-AD ends

Quick review to fix vocabulary, then we leave it behind.

Forward mode: evaluate f on dual numbers (x, ẋ), propagate (f(x), f'(x)·ẋ). Cost is n_inputs × O(f). Best when inputs are few.

Reverse mode: evaluate f, recording a tape of operations; walk the tape backwards multiplying adjoints. Cost is n_outputs × O(f). Best when outputs are few — the case for loss functions.

Source-to-source AD (Tapenade, Zygote.jl, Enzyme, Tangent): transform the program at compile time into a program that computes derivatives. Preserves optimization opportunities; needs to handle the full source language, including its effects.

Operator-overloading AD (PyTorch autograd, ForwardDiff.jl): build the tape at runtime via overloaded operations. Handles dynamic control flow trivially but pays a runtime cost and limits compiler optimization.

Tracing AD (JAX): run the program with abstract values to recover a pure-functional graph, then differentiate the graph. Inherits all the limitations of tracing (memo 1, §3.2).

These approaches differ in where effects become a problem: source-to-source AD must transform effectful constructs, operator overloading must capture them on the tape, tracing must abstract them. The underlying mathematical issue is the same.

The key theoretical result, due to Pearlmutter & Siskind (2008) and refined by Wang et al. (2019), is that reverse-mode AD is structurally equivalent to a CPS transform: the backward pass is the continuation, and the adjoint is what the continuation receives. Wang's "Shift/Reset the Penultimate Backpropagator" makes this explicit by implementing reverse-mode AD with delimited continuations. This is the bridge to effect handlers, which I'll come back to in §6.

3. The effects that break AD

Six distinct effects, each breaking AD in its own way:

Each row is, in effect, a different research community. Let me take them one at a time.

3.1 Mutation and aliasing

The mathematical problem: if x is mutated in place, then x at the time of differentiation is not the same value as x at the time the derivative was computed. Tape-based AD handles this by capturing values eagerly — the tape stores x_before, not a reference to x. Source-to-source AD on SSA form sidesteps the issue (each variable is assigned once). But on a real program with mutable state, you need explicit aliasing analysis.

Enzyme (Moses & Churavy, NeurIPS 2020) is the strongest current answer. It performs AD on LLVM IR after optimization, where alias analysis is already done by the surrounding compiler infrastructure. This lets it handle mutation, in-place updates, and even C/C++/Rust code. The cost is that Enzyme's correctness depends on LLVM's alias analysis being sound — which it mostly is, but the tail cases are real.

The language-design insight: linear types or ownership make this clean. If a tensor has a unique owner (Rust's &mut T, linear types in Idris/Granule), then in-place mutation is safe under AD because aliasing is impossible by construction. Mojo's value semantics + ownership story is partially this. A from-scratch AI-native language should probably make tensors linear-by-default and require explicit copying when sharing is needed.

This intersects with memo 1 (shapes): a Tensor[shape, layout, owned] type captures both the shape contract and the aliasing contract, and AD checks both.

3.2 Randomness

The deepest of the six. Given a function f(x) = E_{z ~ p(z; θ)}[g(x, z)], the gradient ∂f/∂θ is not in general E[∂g/∂θ] — you have to differentiate through the measure, not just the integrand.

There are at least three families of estimators, and choosing among them requires non-local information about the program:

1. Pathwise / reparameterization (Kingma & Welling, 2014; VAE paper): rewrite z = h(ε; θ) where ε is parameter-free. Differentiate normally. Low variance, but only works when the distribution can be reparameterized (Gaussian, yes; categorical, no).

2. Score function / REINFORCE (Williams, 1992): ∂E[g]/∂θ = E[g · ∂log p(z; θ)/∂θ]. Always works, but high variance. Standard in RL.

3. Measure-valued / pathwise extensions / Gumbel-Softmax (Jang et al., 2017; Maddison et al., 2017): relax discrete distributions to continuous ones, then reparameterize.

Plus various variance reduction techniques (control variates, baselines, Rao-Blackwellization).

Schulman et al. (2015), "Gradient Estimation Using Stochastic Computation Graphs," is the closest thing to a unifying theory: model the program as a graph with stochastic and deterministic nodes, and the correct estimator falls out of the graph structure. Storchastic (van Krieken et al., 2021) is a more recent framework. Pyro, Gen.jl, and Stan all implement subsets of this story for probabilistic programming.

The language-design problem: the choice of estimator should be a compiler concern, not a user concern, given enough information about the distribution. But the existing systems make the user pick (pyro.sample(..., infer={'enumerate': 'parallel'}) etc.). A native language with effect-typed sampling could in principle let the user write let z ~ Normal(μ, σ) in g(x, z) and have the compiler choose reparameterization automatically. The handler for the Random effect is the gradient estimator.

This is, in my read, the single most compelling case for effect-typed AD: it converts an ad-hoc choice into a structured language feature.

3.3 Data-dependent control flow

Branches and loops where the condition depends on a tensor value (not a shape). The classic example: if x > 0 then x else -x — the absolute value, which is differentiable except at 0 but has different gradients on each side.

For static branches, reverse-mode AD records which branch was taken on the tape and replays it. For loops, the tape grows with each iteration, which is why checkpointing (Griewank, 1992; Chen et al., 2016) became standard — recompute activations during the backward pass instead of storing them all.

The harder case is recursion or unbounded iteration where you don't want to differentiate by unrolling. This brings us to:

3.4 Implicit definitions and fixed points

Many ML programs define a value implicitly: x* = argmin_x L(x, θ), or x* such that F(x*, θ) = 0. Naive reverse-mode AD unrolls the solver and differentiates through the iterations. This is correct but expensive: memory grows with iteration count, and if the solver is run to convergence, the gradients via the iterations are nearly zero anyway.

The right answer is the implicit function theorem: if F(x*, θ) = 0, then ∂x*/∂θ = -(∂F/∂x*)^{-1} · ∂F/∂θ. You solve one linear system instead of unrolling thousands of iterations.

Christianson (1994), "Reverse accumulation and attractive fixed points," is the foundational paper. Bai, Kolter & Koltun (2019) popularized this in deep learning with Deep Equilibrium Models (DEQ). JAXopt, TensorFlow Probability, and Theseus (PyTorch) all expose implicit differentiation as a library.

The language-design issue: implicit differentiation requires the user to declare that something is a fixed point — the compiler can't infer it from a while loop. An implicit or fixpoint effect, with a handler that knows to use IFT instead of unrolling, fits naturally. This is also where verification becomes relevant: the IFT requires the Jacobian ∂F/∂x* to be invertible, which is a precondition the language could in principle check.

3.5 Exceptions and partial functions

What is ∂(1/x)/∂x at x = 0? Mathematically undefined. In a real program, 1/x raises a division-by-zero exception. Some AD systems propagate NaN, some raise, some define a custom gradient that returns 0 at the singularity.

The deeper question is what gradient semantics you want for partial functions. Three choices:

1. Strict: gradient is undefined where the primal is undefined. Errors propagate.
2. Ignore: define gradient to be 0 at singularities. Hides bugs but allows training to proceed.
3. Custom: user provides a subgradient or a regularized gradient. Used for ReLU, hard thresholding, etc.

This is exactly what an effect handler is good for. An Exception effect with a handler that maps thrown exceptions to NaN, zero, or a custom subgradient encodes the policy in one place rather than threading it through every operator.

3.6 External calls and black boxes

When the primal is a call into a library you can't differentiate (a physics simulator, a CUDA kernel without source, an external service), AD needs a user-provided gradient. PyTorch's torch.autograd.Function, JAX's custom_vjp, and Julia's ChainRules.jl all expose this as an API.

This is uncontroversial — there's no way to differentiate a black box automatically — but the integration into the rest of AD is the issue. A custom-gradient registration is, again, a handler for the External effect.

4. Comparison across systems

How current systems treat each effect:

Patterns to notice:

• PyTorch's tape-based approach handles mutation and control flow naturally but requires manual estimators and a separate library for implicit diff.
• JAX is the most "pure" — it disallows mutation and forces functional patterns, which makes everything cleaner but pushes complexity to the user.
• Zygote.jl is the closest to source-to-source on a high-level language, but the Cassette-based overlay has known correctness gaps.
• Enzyme is the strongest on low-level effects (mutation, exceptions, FFI) because it operates on LLVM IR; it's weak on high-level effects (probabilistic semantics) because it doesn't know about them.
• PPLs treat randomness as the central effect and let the inference algorithm pick estimators, but inherit the host language's handling of everything else.

Nobody has all six effects handled compositionally. Each system optimizes for a different subset, and crossing system boundaries (using Enzyme for mutation + Pyro for sampling) requires manual glue.

5. Effect systems as the unifying framework

The thesis: algebraic effects with handlers are the right framework for AD over effectful programs.

Three reasons:

5.1 Compositionality

Algebraic effects compose: a program with Random + State + Exception is handled by composing the handlers, in any order that preserves semantics. Each handler can be specified independently. This matches what we want for AD: a Differentiate handler that composes with a Random handler (which provides the estimator) and a State handler (which provides aliasing-aware capture).

The contrast with monad transformers is informative. Haskell's mtl lets you compose StateT and RandT and ExceptT, but the order matters and the lifting boilerplate is real. Algebraic effects fix this with row-polymorphic handlers (Koka, Eff, OCaml 5).

5.2 The CPS connection

Wang et al. (2019) showed that reverse-mode AD is reducible to delimited continuations: the backward pass is a continuation, and adjoint accumulation is exactly the kind of metaprogramming shift/reset were designed for.

Effect handlers compile to CPS. AD compiles to CPS. They're the same compilation strategy. A language that has effect handlers as a primitive is, accidentally, a language that has the right machinery for AD.

OCaml 5's effect handlers (Sivaramakrishnan et al., PLDI 2021) are the most mature production implementation. Eio and Miou exploit them for concurrency. Nobody has yet exploited them for AD, but the path is clear: declare effect Diff : ..., write a handler that does adjoint accumulation, and you get reverse-mode AD as a library, not as a separate compiler pass.

5.3 The estimator-as-handler insight

The most concrete payoff. For a program:

let model x =
  let z ~ Normal(0, 1) in
  let y = neural_net(z, x) in
  y

The Normal sampling is an effect. Different handlers correspond to different gradient estimators:

• with reparameterize handle ... → pathwise gradient.
• with score_function handle ... → REINFORCE.
• with gumbel_softmax handle ... → relaxation-based.

The user writes one program; the inference algorithm (or the compiler, with sufficient analysis) picks the handler. This is the cleanest framing of the estimator-selection problem I know of, and it's missing from every production system.

5.4 What the effect-typed AD design looks like

Sketch:

effect Diff[a]:
  primal : a
  adjoint : a -> Unit  // accumulate adjoint

effect Random[a]:
  sample : Distribution a -> a

effect State[a]:
  get : Unit -> a
  set : a -> Unit

effect Implicit[a]:
  fixpoint : (a -> a) -> a  // declare an implicit definition

handler reverseMode : Diff handler  // reverse-mode AD
handler reparam : Random handler    // pathwise estimator  
handler score : Random handler      // REINFORCE
handler ift : Implicit handler      // implicit function theorem

A program is then differentiated by composing handlers. The compiler is free to pick handlers when it has enough information; the user can override.

This is not a fully worked-out design — type inference for higher-order effect-polymorphic AD is genuinely hard — but it is, I think, the right shape.

6. Open research problems

1. A worked-out effect-typed AD calculus. Brunel, Mazza & Pagani (POPL 2020) gave a typed reverse-mode AD calculus over linear lambda calculus. Extending this to algebraic effects, with a soundness proof, is a real research project. The closest published work I know is fragmented — pieces in Wang's work, in differential lambda calculus, in linear logic. Nobody has the whole picture.

2. Verified AD implementations. This is your wheelhouse. Coq formalizations of forward-mode AD exist; reverse-mode is harder and less complete. A CompCert-style verified AD for a non-trivial language fragment is open. Given your Frama-C work, this is a credible research direction.

3. Compositional gradient estimators. Schulman's stochastic computation graphs unified the manual estimators in 2015; nobody has packaged this into a language feature where the compiler picks the estimator from the graph structure. The pieces exist; the integration doesn't.

4. Aliasing-aware AD with linear types. Linear or affine types should make mutation-aware AD trivial — but the language design (which fragment of linearity to expose, how to handle views) is not settled. Granule, Idris 2, and Linear Haskell have parts of the answer; none integrates with AD.

5. Differentiating through control flow on dynamic shapes. This is the intersection with memo 1: when shape is data-dependent (MoE routing, masking, beam search), AD has to differentiate through the routing decision. The current approaches (straight-through estimators, Gumbel-Softmax) are heuristic. A principled answer requires an effect-typed shape system.

6. Higher-order AD with effects. Hessian-vector products, second-order optimization, meta-learning — all need higher-order derivatives. The composition of higher-order AD with effects is not well-understood. Pearlmutter's R{} operator and forward-on-reverse tricks work for pure functions; the effectful case is largely open.

7. Distributed AD as a effect. Communication during the backward pass (gradient all-reduce, ZeRO partitioning) is currently bolted on as a framework feature. As a Distributed effect with a handler, it could compose with the other effects. This crosses into memo 4 (distribution).

8. Termination and AD. Differentiating through a non-terminating program is undefined. Total languages (Idris, Agda, Lean) sidestep this; partial languages need an effect for divergence. How does that compose with AD? Open.

7. Recommendation for an AI-native language design

Concrete positions, in priority order:

• Algebraic effect handlers as the core abstraction, not monads or implicit framework state. OCaml 5 / Koka style.
• Differentiate is one effect among many, not a pervasive language feature. AD becomes a library that uses the effect machinery.
• Random's handler is the gradient estimator. Default to reparameterization where possible, score function elsewhere; let the user override with a handler annotation.
• Linear or affine types for tensors, so mutation-aware AD is sound by construction. Mojo's ownership model is roughly right.
• Implicit / fixpoint as a declared effect, with the IFT handler as default. No more separate implicit_diff libraries.
• Custom gradients via handlers, not via a separate Function API. with my_grad handle expensive_op(x) reads more naturally than registering a class.
• Compile-time effect inference, so users don't have to annotate everything. The annotations should mostly be at the boundaries.

The thesis: the current AD ecosystem is a collection of special-case patches because the host languages don't have effect handlers. Once you have effect handlers, AD becomes the simplest case of a composable program transformation, and the ten special cases collapse into one framework.

8. References

Foundational AD theory

• Pearlmutter, B. & Siskind, J. (2008). Reverse-mode AD in a functional framework: Lambda the Ultimate Backpropagator. TOPLAS. — The foundational paper for AD over higher-order functional programs.
• Elliott, C. (2018). The simple essence of automatic differentiation. ICFP 2018. — Categorical formulation; clean and worth reading.
• Griewank, A. & Walther, A. (2008). Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed. SIAM. — Standard reference for the field.

AD as CPS / continuations

• Wang, F. et al. (2019). Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator. ICFP 2019. — The CPS-AD bridge; essential for effect-handler thinking.
• Brunel, A., Mazza, D. & Pagani, M. (2020). Backpropagation in the Simply Typed Lambda-Calculus with Linear Negation. POPL 2020. — Type-theoretic foundations.
• Krawiec, F. et al. (2022). Provably Correct, Asymptotically Efficient, Higher-Order Reverse-Mode Automatic Differentiation. POPL 2022.

Modern AD systems

• Moses, W. & Churavy, V. (2020). Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients. NeurIPS 2020. (Enzyme.)
• Innes, M. (2018). Don't Unroll Adjoint: Differentiating SSA-Form Programs. arXiv:1810.07951. (Zygote / Diffractor design.)
• Bradbury, J. et al. JAX: composable transformations of Python+NumPy programs.
• Maclaurin, D. (2016). Modeling, Inference, and Optimization with Composable Differentiable Procedures. PhD thesis, Harvard. (Autograd origins; very readable.)
• Paszke, A. et al. (2017). Automatic differentiation in PyTorch. NIPS Workshop.

Stochastic AD

• Williams, R. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning. (REINFORCE.)
• Kingma, D. & Welling, M. (2014). Auto-Encoding Variational Bayes. ICLR 2014. (Reparameterization trick.)
• Schulman, J. et al. (2015). Gradient Estimation Using Stochastic Computation Graphs. NeurIPS 2015. — Unifies manual estimators.
• Jang, E., Gu, S. & Poole, B. (2017). Categorical Reparameterization with Gumbel-Softmax. ICLR 2017.
• Maddison, C., Mnih, A. & Teh, Y. (2017). The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. ICLR 2017.
• van Krieken, E. et al. (2021). Storchastic: A Framework for General Stochastic Automatic Differentiation. NeurIPS 2021.

Implicit differentiation

• Christianson, B. (1994). Reverse accumulation and attractive fixed points. Optimization Methods and Software. — Foundational.
• Bell, B. & Burke, J. (2008). Algorithmic differentiation of implicit functions and optimal values. — The IFT-AD connection.
• Bai, S., Kolter, Z. & Koltun, V. (2019). Deep Equilibrium Models. NeurIPS 2019. — Brought implicit differentiation to mainstream deep learning.
• Blondel, M. et al. (2022). Efficient and Modular Implicit Differentiation. NeurIPS 2022. (JAXopt.)

Algebraic effects and handlers

• Plotkin, G. & Pretnar, M. (2009). Handlers of Algebraic Effects. ESOP 2009.
• Bauer, A. & Pretnar, M. (2015). Programming with Algebraic Effects and Handlers. JLAMP.
• Leijen, D. (2017). Type Directed Compilation of Row-Typed Algebraic Effects. POPL 2017. (Koka.)
• Sivaramakrishnan, K. C. et al. (2021). Retrofitting Effect Handlers onto OCaml. PLDI 2021. — OCaml 5's implementation.
• Pretnar, M. (2015). An Introduction to Algebraic Effects and Handlers. MFPS tutorial. — Most readable starting point.

Probabilistic programming (estimator selection in practice)

• Bingham, E. et al. (2019). Pyro: Deep Universal Probabilistic Programming. JMLR.
• Cusumano-Towner, M. et al. (2019). Gen: A General-Purpose Probabilistic Programming System with Programmable Inference. PLDI 2019.
• Carpenter, B. et al. (2017). Stan: A Probabilistic Programming Language. Journal of Statistical Software.

Adjacent — relevant but not central

• Maclaurin, D., Duvenaud, D. & Adams, R. (2015). Gradient-based Hyperparameter Optimization through Reversible Learning. ICML 2015. — Differentiating through training.
• Chen, T. et al. (2016). Training Deep Nets with Sublinear Memory Cost. — Checkpointing for control flow.
• Griewank, A. (1992). Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. — Original checkpointing paper.

9. Connections to the rest of the series

• Memo 1 (shapes): data-dependent shapes interact with control-flow AD. MoE routing is the canonical hard case. An effect-typed shape system and an effect-typed AD system want to share the same effect framework.
• Memo 3 (sparsity): sparse gradients (Top-K, sign-SGD) and sparse activations (MoE again) are both estimator-choice problems disguised as data structure problems. Connects directly to §3.2.
• Memo 4 (distribution): gradient all-reduce and ZeRO are Distributed effects in the AD pass. The compositional handler story extends naturally.
• Memo 5 (verification): verified AD is the most credible "real research" angle in this series, and it builds on §5 + §6.

10. Next memo

Design Space 03: Compositional Sparsity and Structured Tensors. Sparse and structured tensors (block-sparse, low-rank, MoE routing, KV-cache compression) are currently library features bolted onto dense frameworks. The question: can sparsity patterns be language-level types, with composition rules and AD that respects them? This connects to type-level shapes (memo 1) and stochastic estimators (memo 2).