Autodiff survey

autodiff-survey.md

Modern Autodiff Architectures

Comparative report and research outlook — evidence and probes current through 22 July 2026

Executive summary

Automatic differentiation (AD) is no longer one architecture. Modern systems differentiate at several levels: PyTorch records eager tensor operations and can later stage a joint forward/backward graph; Zygote rewrites typed Julia SSA; JAX interprets a small primitive language and transposes a linearized Jaxpr; Enzyme transforms optimized LLVM IR. Clad works nearer the C++ AST, Swift's language-integrated experiment transforms SIL, and Rust's experimental std::autodiff is a language-facing route into Enzyme. The transformation level determines what semantics the differentiator can see, which effects it can handle, and which optimizations it can exploit.

Four conclusions survive both the literature review and the local probes:

1. JVP and VJP are the useful architectural atoms. “Forward mode” and “reverse mode” describe how they are evaluated; gradients, Jacobians and Hessian-vector products are compositions around them.
2. Mutation is an architectural choice, not a minor feature. PyTorch permits guarded eager mutation and can functionalize it for compilation. JAX exposes functional updates. Zygote's SSA transform still rejects common array mutation, while Mooncake and Enzyme explicitly model it. Alias analysis and effect semantics determine whether reverse code is sound.
3. Reverse-mode memory is a compiler policy. A tape, a pullback closure and residual arguments are three representations of information needed later. AOTAutograd partitioning, jax.remat, Enzyme's augmented forward pass and checkpoint schedules decide what to save, recompute or move.
4. No level dominates. High-level IR retains types, structure and useful diagnostics; low-level IR sees foreign and optimized code. The strongest research direction is multi-level AD with a shared rule protocol, not one universal tape.

For Wolfram specifically, the answer is yes, but not as a general public grad transform. The Wolfram Language has symbolic D, undocumented/internal neural-network backpropagation, and documented automatic gradient construction for typed functions embedded in CompiledLayer. More substantially, the Mathematica-based commercial AceGen system has for decades combined symbolic manipulation, forward/reverse AD, simultaneous expression optimization and C/Fortran generation, especially for finite elements. These are serious efforts with narrower interfaces and domains than JAX, Zygote or PyTorch's public differentiable-programming APIs.

1. A common model

The terminology here follows Baydin et al.'s broad JMLR survey of automatic differentiation. AD evaluates derivatives by applying the chain rule to a program's elementary operations; it is neither finite differencing nor merely computer-algebraic simplification.

For a differentiable program f : R^n -> R^m at primal input x:

• A JVP (Jacobian-vector product), J_f(x) v, pushes an input tangent v forward. Its cost is largely independent of m, so forward mode suits few inputs or directional derivatives.
• A VJP (vector-Jacobian product), u^T J_f(x), pulls an output cotangent u backward. Its cost is largely independent of n, so reverse mode gives a scalar-output gradient in roughly one backward sweep.
• A linearization returns f(x) and a reusable linear map for JVPs. A pullback returns f(x) and a reusable map from output cotangents to input cotangents. Transposition mechanically turns a linear JVP program into its reverse linear program when the IR and primitive transpose rules permit it.
• Residuals are primal intermediates that reverse code needs. They may live on an operator tape, in a pullback closure, or as explicit values passed from augmented forward code. “Saved tensors” are PyTorch's prominent example.
• Checkpointing/rematerialization discards selected residuals and recomputes them in reverse. It trades time for memory; the scheduling problem can be optimized globally, as Checkmate demonstrates.
• Operator overloading changes the run-time meaning of primitives (dual numbers or tracked tensors). Tracing runs a program with abstract values to capture an IR. Source transformation rewrites program representation into primal and derivative programs. Real systems combine them.

The phrase “source transformation” therefore needs a level: Python bytecode or FX graph, Julia typed SSA, Swift SIL, Clang AST, MLIR dialect, or optimized LLVM IR. Each retains a different subset of source intent.

2. Comparison at a glance

The table describes the public architecture, not every backend optimization. “Higher order” means a supported composition path exists, not that every primitive or custom rule is arbitrarily differentiable.

System	Capture / transform level	Execution and control flow	Mutation / aliasing	Higher order and rules
PyTorch eager	Dynamic tensor-operation graph of backward Nodes	Ordinary Python executes; only the path taken is recorded	In-place operations guarded by leaf rules, saved-tensor version counters and view semantics	Double backward where nodes/rules support it; custom autograd.Function
torch.func / AOTAutograd	Functional transforms plus FX/ATen joint forward-backward graph	Transforms compose; compilation specializes/guards Python behavior	Functionalization rewrites many mutations and views to pure equivalents	jvp, vjp, grad, jac*, hessian, vmap; transformable custom Functions
Zygote	Lowered/typed Julia SSA to primal + pullback	Native Julia control flow transformed after specialization	Ordinary scalar code works; array mutation is a central documented limitation	ChainRules rrule; higher-order support is incomplete and often mixed with ForwardDiff
JAX	Primitive tracing into Jaxpr; linearization + transposition	Eager transforms or staged XLA execution; dynamic staged control flow uses lax.cond/loops	Arrays immutable; functional update syntax and effect-aware primitives	First-class composable transforms; custom_jvp / custom_vjp
Enzyme	Optimized LLVM IR (and related compiler integrations)	Native compiled control flow, memory and calls where analyzable	Designed for mutation and aliasing using activity/type analysis and shadow memory	Forward, reverse and mixed modes; frontend-dependent custom rules
Mooncake	Julia IR transformation/rule system	Native Julia control flow within supported compiler subset	Explicitly designed to support mutation and alias-aware reverse rules	Reverse plus forward-over-reverse; extensible rule interface
ForwardDiff	Dual-number operator overloading	Native Julia generic code	Works when code accepts dual element types; mutation of dual-compatible storage is possible	Nested duals; custom behavior through generic methods/ChainRules bridges
ReverseDiff	Runtime instruction tape, optionally compiled	Dynamic tape or fixed recorded tape	Tracked arrays restrict some mutation	Gradient/Jacobian/Hessian APIs; tape compilation improves reuse but freezes recorded paths
TensorFlow	Eager GradientTape, optionally tf.function graphs	Eager path recording; graph control-flow ops when staged	Resource-variable mutation supported within TensorFlow semantics	Nested/persistent tapes, ForwardAccumulator, registered gradients
Clad	Clang C/C++ AST and generated C++	Native language control flow; templates and CUDA in supported subsets	Can transform explicit loads/stores with C++-level type visibility	Forward/reverse/Jacobian/Hessian; custom pushforwards and pullbacks
Swift differentiation	Typed SIL compiler transformations	Language-integrated @differentiable functions	Governed by Swift ownership/SIL support; still experimental	JVP/VJP registrations and generated derivatives
Rust std::autodiff	Rust attribute/frontend, Enzyme at LLVM	Native compiled code in nightly experiment	Enzyme-backed memory model; substantial restrictions remain	Forward/reverse activity annotations; experimental

Memory and systems trade-offs are easier to compare separately:

System	Residual / memory policy	Foreign code	Compilation and debugging character
PyTorch	Saved tensors eagerly; checkpointing; AOT partitioner can rematerialize	Opaque calls need a custom Function or decomposition	Excellent eager graph inspection and anomaly/version diagnostics; compiled stack adds guards and graph breaks
Zygote	Pullback closures capture residuals	ccall is generally not differentiable without a rule	Julia source/IR visibility helps inspection; failures can surface in generated pullbacks
JAX	Residuals made explicit by partial evaluation; remat changes saving policy	Needs a primitive plus lowering and derivative rules	make_jaxpr is precise; tracer errors explain staging violations, though downstream compiler IR is lower-level
Enzyme	Augmented forward decides caches/tape; shadow memory holds derivatives	Strongest opportunity when foreign code reaches compatible LLVM IR	Cross-language reach and optimization; source diagnostics and lost high-level intent remain challenges
Mooncake	Explicit forward/reverse data structures and rule-defined storage	Unsupported calls need rules	Julia-level errors and IR inspection; younger ecosystem
TensorFlow	Tape stores intermediates; persistent tapes retain resources	Custom op gradients required	Mature profiler/graph tooling; eager vs graph boundaries matter

3. PyTorch: a dynamic tape that learned to compile

Eager architecture

PyTorch autograd records a directed graph as tensor operations execute. A result with requires_grad carries a grad_fn pointing to a backward Node; calling backward or autograd.grad traverses those nodes and accumulates cotangents. The graph is reconstructed on every forward call, so ordinary Python branches and bounded loops differentiate the path actually taken. The official autograd mechanics note documents saved tensors and graph recreation; the autograd API covers graph inspection and hooks.

Backward formulas save selected primal tensors. Since an in-place write could silently invalidate one, each tensor has a version counter; a saved tensor's counter is checked when unpacked. Leaf tensors requiring gradients and views have additional mutation restrictions. torch.autograd.Function gives an opaque operation explicit forward, backward, context-saving and, where needed, transform hooks. Rules intended to compose with torch.func must obey the documented transformability constraints.

End-to-end walkthrough. For y = sum(sin(x)^2), eager execution creates nodes for sin, multiplication/power and reduction, saving what their backward formulas require. Seeding dy = 1, reduction broadcasts a cotangent, squaring produces 2 sin(x), and SinBackward multiplies by cos(x), yielding 2 sin(x) cos(x). Our graph walk found SinBackward0, multiplication/division, addition and AccumulateGrad nodes in the branch-and-loop workload. The VJP was numerically identical to torch.func.grad, and forward-over-reverse produced the expected Hessian-vector product.

Functional transforms and AOTAutograd

torch.func exposes composable grad, vjp, jvp, Jacobian/Hessian and vmap transforms (API). Its functionalization pass replaces many in-place updates and views with non-mutating operators while preserving input updates when required. In our view-mutation probe, functionalization exposed clone, slice_copy, add, and slice_scatter ATen operations and preserved [0.4, 1.2, 3.4] as the gradient.

This is distinct from eager autograd. The PyTorch 2 architecture paper describes AOTAutograd running eager autograd with fake tensors while a dispatcher captures a joint forward/backward graph. Decomposition and functionalization make it compiler friendly. A partitioner then splits forward from backward and can use a min-cut strategy to save expensive activations but rematerialize cheap ones. Thus the compiled system turns local save_for_backward choices into a graph-level memory decision.

Architectural limitation. Eager flexibility is path-specific, while staging must guard or graph-break around Python and unsupported effects. In-place semantics, views, hooks, custom Functions and transform composition form a large compatibility surface; the official torch.func limitations make clear that not every eager program is transformable. PyTorch's solution is a layered stack, not a single uniform IR.

4. Zygote: reverse transformation of Julia SSA

Zygote differentiates after Julia lowering and specialization. Its documented _pullback internals use generated-function machinery to inspect lowered SSA and emit a primal plus a reverse pullback. The forward code returns the primal result and a generated Pullback struct whose fields are closure-converted residuals/pullbacks. Calling it with an output cotangent walks the transformed SSA backward and returns cotangents for arguments and captured values.

Julia's specialization is a major advantage: generic methods have concrete types and dispatch targets before transformation. Zygote can differentiate ordinary scalar loops, branches, structs and closures without requiring users to express programs as a tensor graph. It delegates many domain-specific rules to ChainRulesCore: an rrule returns (primal, pullback), while an frule specifies a pushforward. Tangent types distinguish structural NoTangent from an additive ZeroTangent, and thunks can delay expensive cotangent work.

End-to-end walkthrough. In y = sum(sin.(x).^2), Julia specializes the broadcast/reduction calls. Zygote emits forward code returning y and residual pullbacks. A seed of 1 enters the reduction rule, broadcast rules pull it into elementwise cotangents, and the sine/power rules return 2 .* sin.(x) .* cos.(x). In our common probe this gave [0.1186703954, 0.4589044329, 1.2766774339], matching PyTorch, JAX, ForwardDiff and finite differences. A captured closure also returned both the argument gradient [4.5, 9.0] and the captured-scale cotangent 15.0 in the analogous Python probe, illustrating why pullbacks are closures rather than only arrays.

Architectural limitation. Mutation is not an incidental missing primitive. The Zygote limitations page calls array mutation the most important limitation and also records problems around foreign calls, exception handling and some higher-order derivatives. Our differentiable view update failed exactly as documented: “Mutating arrays is not supported.” Reverse transformation needs a coherent account of aliases, overwrites and old values; pure SSA values do not alone supply one. Users often rewrite code functionally, add an rrule, or compose ForwardDiff over Zygote for higher derivatives. This is useful evidence about the transform's boundary, not merely a bug list.

5. JAX: interpreters over a differentiable primitive language

JAX's key abstraction is not a tape but an extensible interpreter. A primitive has implementations, abstract evaluation, lowering and transformation rules. Transforms run Python with tracers whose operations bind primitives; the trace either computes tangents, batches values, or stages equations into a typed, functional Jaxpr. The compact Autodidax implementation is the clearest primary account of this architecture.

Forward AD attaches primal and tangent values to JVP tracers. Reverse AD is factored rather than independently hard-coded: JAX linearizes with JVP rules, uses partial evaluation to separate known primal work and explicit residuals, then transposes the resulting linear Jaxpr. That decomposition explains why grad, jvp, vjp, vmap and jit can compose. custom_jvp and custom_vjp override mathematical behavior at an abstraction boundary (custom-rule guide).

Staged Python must be abstractly executable. Static loops may unroll, while data-dependent control flow uses structured primitives such as lax.cond, scan, fori_loop or while_loop. Arrays are immutable; .at[...] expresses a functional update that compilers may lower in place. JAX's tracing guide explains concrete-value errors at the Python/staging boundary.

End-to-end walkthrough. Tracing sum(sin(x)^2) produces Jaxpr equations for sin, integer power and reduce_sum. JVP rules add cos(x) * x_dot and the linearized power rule. Partial evaluation retains needed primals; transposition starts from scalar cotangent 1, transposes the reduction and multiplies through the linear equations. Our gradient Jaxpr contained sin, cos, integer_pow, reduce_sum, broadcasts, multiplies and add_any; JVP, VJP and grad agreed, and nested transforms produced the same Hessian-vector product as PyTorch and Mooncake.

jax.remat/checkpoint changes what the partial evaluator is allowed to save, requesting recomputation during reverse; the rematerialization notebook explains the policy. Our printed Jaxpr contained the rematerialization primitive.

Architectural limitation. JAX transforms a restricted, effect-aware primitive language, not arbitrary Python. Our ordinary data-dependent Python if worked eagerly but under jit raised TracerBoolConversionError; direct array assignment raised an immutability error. Both have explicit JAX forms, but libraries must provide primitives, lowerings and derivative rules for opaque foreign work. Composition is exceptionally systematic inside the language and deliberately bounded outside it.

6. Enzyme: AD after optimization

Enzyme differentiates compiler IR, principally LLVM. Its NeurIPS 2020 paper describes a pipeline that performs type and activity analysis, differentiates instructions and memory effects, and allows the derivative to participate in ordinary compiler optimization. Since many languages lower to LLVM, one engine can serve C/C++, Julia, Rust, Swift experiments and mixed-language kernels when sufficient analyzable IR survives.

An active value can affect a requested derivative; inactive values need no shadow. Enzyme associates differentiable storage with shadow memory for tangents or adjoints. Reverse mode creates an augmented forward that executes the primal and caches values that cannot safely be reconstructed, then a reverse function that consumes that tape, visits memory effects backward and accumulates shadows. This can handle explicit mutation and aliases because loads, stores and calls are visible, though correctness depends on type/activity information and supported calling conventions. The official Enzyme usage documentation and Julia API expose forward, reverse and mixed modes with Const, Active and Duplicated annotations.

End-to-end walkthrough. For a loop that copies x, mutates a view and sums squares, LLVM IR exposes allocation/copy, pointer arithmetic, loads, stores and the reduction. Activity analysis marks the input/output path; augmented forward saves overwritten values or other needed residuals and allocates derivative shadows. Reverse replays the loop backward, updates the view's shadow and propagates it to the aliased base allocation. The local Enzyme.jl 0.13.190 probe returned [0.4, 1.2, 3.4], agreeing with analytical differentiation, while the same source pattern was rejected by Zygote and ReverseDiff.

Differentiating optimized IR creates unusual opportunities: constant propagation, alias analysis, inlining and dead-code elimination can simplify the primal before AD; optimization can then simplify the derivative. It also risks losing high-level facts—mathematical operator identity, source structure, shape intent or safe abstraction boundaries—that would make better rules and errors.

Architectural limitation. LLVM validity is not differentiability. Inline assembly, opaque libraries, integer/pointer tricks, incomplete type information, concurrency and effects may defeat analysis or require custom rules. Low-level errors are harder to explain at source. Cross-language reach is strongest when the relevant code is available as compatible IR; a precompiled opaque binary is still opaque. Enzyme shifts the frontier substantially, but does not erase FFI boundaries.

7. Compiler-native AD as a transformation ladder

C++ AST: Clad

Clad is a Clang plugin that derives C++ functions at the AST level. It retains overloads, types, templates and source constructs, can emit inspectable C++ for forward/reverse/Jacobian/Hessian modes, and supports CUDA subsets. Its custom derivative mechanism names pushforward and pullback rules. Compared with LLVM AD, it can produce better source-oriented diagnostics and use semantic operator identity; it sees less of fully inlined cross-language code and must track C++'s complex alias/effect semantics itself. This report did not build Clad; those observations and generated-code examples are source-derived.

Typed language IR: Julia and Swift

Zygote and Mooncake demonstrate two generations of Julia-IR AD. Zygote's primal/pullback rewrite excels on functional Julia but exposes mutation limits. Mooncake also transforms Julia IR, with a rule system designed around reverse-data storage and support for mutation. Our Mooncake 0.5.40 probe differentiated the mutating view and a forward-over-reverse Hessian-vector product successfully. Typed Julia IR retains multiple dispatch and structures, while Julia compiler integration can still lower the result to LLVM.

Swift's differentiation experiment integrates generated JVP/VJP functions with the language and transforms the compiler's typed SIL, whose ownership and calling-convention information is richer than LLVM. It remains experimental; the project's February 2026 status update reports Swift 6.3 work including throwing functions, WebAssembly and performance fixes rather than declaring the feature stable. Swift was not installed locally, so this evidence is source-derived.

LLVM and multi-level IR: Enzyme and MLIR

LLVM maximizes exposure to inlined, optimized, cross-language machine-oriented code. MLIR supplies a framework of progressively lowered dialects, making it an attractive place to preserve tensor, control-flow or domain meaning longer and later reach LLVM. It is important not to overstate the current platform: MLIR is a multi-level compiler framework, not itself one standard, universal AD dialect. Enzyme-related integrations can operate in this ecosystem, but maturity varies by dialect and frontend.

Rust frontend backed by Enzyme

Nightly Rust exposes an experimental std::autodiff attribute interface. The unstable compiler documentation states that rustc uses Enzyme and requires an Enzyme-enabled toolchain component. Rust supplies user-facing function/activity annotations; lower-level LLVM AD generates the body. This is a clean example of a language frontend backed by a shared lower-level differentiator. The installed Rust 1.96 toolchain lacked the Enzyme component, so no plugin build was attempted.

The actual trade-off

Higher transformation level tends to preserve	Lower transformation level tends to expose
Mathematical operator identity and shapes	Inlined implementation and memory operations
Source locations and actionable diagnostics	Cross-language code lowered to common IR
Ownership, algebraic data types, structured control flow	Alias analysis and target-specific optimization
Intentional custom derivative boundaries	Dead code, constants and optimization opportunities
Separate-compilation/API structure	The code actually executed after lowering

A compiler-native design should therefore pass derivative meaning down and analysis facts up. Separate compilation needs derivative ABI/rule metadata; foreign calls need summaries or IR; and diagnostics need a provenance chain back to source. Merely choosing “high” or “low” leaves value on the table.

8. Context systems

TensorFlow. tf.GradientTape records operations executed while watched values are active, retains intermediates needed by registered gradients and then releases them unless persistent. Nested tapes give higher derivatives; tf.autodiff.ForwardAccumulator provides JVPs. The official autodiff guide and advanced guide show the eager model and Jacobian/Hessian patterns. tf.function adds staged graphs and structured control flow. Its architecture is closest to PyTorch eager plus graph capture, with particularly mature resource-variable and deployment semantics.

ForwardDiff. ForwardDiff uses tagged dual numbers whose coefficients carry one or several directional derivatives. Chunking chooses how many input columns are propagated per run; nested duals give higher derivatives. Its implementation guide makes the model explicit. It is simple, robust on generic Julia numerical code and often the right complement to reverse engines, but scales poorly for very many inputs and fails where code insists on concrete non-dual element types or opaque calls.

ReverseDiff. ReverseDiff records operations in an instruction tape and can compile a tape for repeated execution. The API documentation warns that compiled tapes capture the execution path at recording input. Our probe recorded a positive branch, then replayed at a negative input: the dynamic gradient was [-1,-1], while the compiled tape returned [-2,-4] from the stale branch. This intentionally demonstrates that tape compilation is specialization, not a free optimization. Its tracked array also rejected the view mutation.

Tracker. Flux's older Tracker represents classic operator-overloaded reverse mode: tracked values build a graph of operations and pullback closures as the program runs. The current Flux guide describes it as Flux's original, more traditional operator-overloading approach. It remains historically useful for understanding Julia's move from runtime tapes toward compiler transformation and shared ChainRules, but it is no longer the principal general architecture considered here.

9. Wolfram ecosystem: serious, narrower attempts

Symbolic differentiation is adjacent, not equivalent

Wolfram Language's D and Derivative are powerful symbolic operators. They manipulate mathematical expressions and know large libraries of identities. That is valuable for closed-form simplification, exact derivatives and rule generation, but it is not by itself program-level AD over mutation, closures, foreign calls and arbitrary compiled control flow.

Neural-network layers plainly perform backpropagation internally, but public evidence should not be stretched into claims about proprietary implementation. The relevant documented public bridge is CompiledLayer: it accepts a typed compiled function and attempts to construct its automatic gradient for use as a neural-network layer, subject to port/shape restrictions; the user can supply a separate gradient function if automatic construction fails. That is genuine AD of a typed function, but scoped to the layer contract, not a general public grad[f] for arbitrary Wolfram programs.

The public Wolfram Compiler overview and FunctionCompile documentation expose typed compilation and IR/tooling, but as of the report date document no general-purpose differentiation transform comparable to jax.grad, Zygote's gradient, or Enzyme's autodiff. Absence here means “not found in the documented public interface,” not a claim about internal code.

AceGen: the substantial precedent

AceGen is a commercial Mathematica application for automatic code generation, especially numerical and finite-element kernels. Its manual describes a hybrid pipeline: symbolic representation, automatic differentiation in forward and reverse forms, simultaneous optimization of expressions and derivatives, and generation of C/Fortran and other target code. This is not a toy package; it has been used over a long period with the companion AceFEM and continues to appear in recent finite-element research—for example a 2025 work on automatically generated finite-element code.

AceGen's significance is architectural. A computer-algebra environment can use domain identities and global expression optimization before and during AD, generating compact derivative kernels rather than interpreting a generic tape. Its limitations relative to modern differentiable programming are interface and scope: it is a separate proprietary product and code-generation workflow, strongly associated with mechanics/FEM, rather than a ubiquitous transform that composes with every Wolfram Language program. Nonetheless it is the clearest answer to whether a serious Wolfram-ecosystem attempt exists.

Bottom line for Wolfram

There are three distinct layers, and conflating them causes confusion:

1. symbolic D/Derivative: broad mathematics, not general program AD;
2. CompiledLayer: documented, real automatic gradients for a constrained typed neural-layer function;
3. AceGen: mature symbolic/AD/optimization/code-generation technology with a specialized commercial workflow.

So the ecosystem has important AD work and a particularly interesting symbolic hybrid, but no documented public general transform at the Wolfram Compiler level as of 22 July 2026.

10. Reproducible probes and observations

The checked-in experiments/workloads.md specifies scalar/array branches and loops, mutation/views, closures, structured transforms, custom reverse rules, rematerialization and IR inspection. Python and Julia lock files pin all packages; experiments/README.md gives the commands. Full captured values are in experiments/results/python.txt and experiments/results/julia.txt.

Probe	Observed result
Branch + bounded loop	PyTorch, JAX eager, Zygote and ForwardDiff agreed with central finite differences: gradient 2.88930719036
JVP/VJP/gradient	PyTorch, JAX and Julia engines agreed on primal 0.26130650521 and gradient [0.11867039543, 0.45890443289, 1.27667743390]
Hessian-vector product	PyTorch, JAX and Mooncake agreed on [1.17344515063, -1.09511367013, -0.77858170803]
Mutation through a view	PyTorch eager/functionalized, Enzyme and Mooncake returned [0.4,1.2,3.4]; Zygote and ReverseDiff rejected it; JAX required .at functional update
Data-dependent staged Python branch	JAX jit raised TracerBoolConversionError; the structured lax equivalent is the supported representation
Opaque custom reverse rule	PyTorch autograd.Function and JAX custom_vjp implemented a stable softplus rule, returning primal 50 and gradient 1 at input 50
Rematerialization	PyTorch checkpoint executed the counted forward region twice; JAX's printed Jaxpr contained the rematerialization primitive
Compiled tape branch change	ReverseDiff replay returned a stale-branch gradient, reproducing its documented recording restriction
IR inspection	Eager PyTorch nodes, functionalized ATen graph, Jaxpr primitives and optimized primal LLVM IR were captured

These are behavior/correctness studies, deliberately not timing claims. The restricted workspace sandbox denied Julia precompile-worker pipe creation; --compiled-modules=no made the original focused probes runnable but imposed a large cold-start cost. A controlled outside-sandbox rerun later confirmed normal package-image construction and is reported separately in julia-latency-followup.md. Swift, Clad, Rust-Enzyme and Wolfram examples are explicitly source-derived because those toolchains/kernel components were unavailable. wolframscript was installed, but no local licensed Wolfram kernel was available.

11. Prioritized research outlook

1. Multi-level AD with retained semantics. Attach derivative summaries and mathematical identities to high-level operations, lower them with provenance, and allow lower-level activity/alias analysis to feed facts back upward. Julia-to-LLVM and Rust-to-Enzyme show pieces; MLIR offers infrastructure, but a portable derivative ABI and rule metadata remain open.

2. Effect- and alias-aware differentiation. Purity simplifies transposition, yet scientific programs mutate arrays, use RNGs, throw exceptions and perform I/O. Functionalization is valuable but cannot be the only answer. Ownership, effect systems and shadow-memory analyses should expose which mutations are reversible, which values must be logged, and why unsupported effects fail.

3. Reusable JVP/VJP rule protocols. JAX primitive rules, ChainRules frule/rrule, PyTorch custom Functions and Enzyme custom rules express related mathematics with incompatible packaging and tangent types. A language-neutral protocol needs structured tangents, zero/lazy cotangents, effects, batching and lowering—not merely a callback returning an array.

4. Memory policy as an explicit compiler problem. Residual selection, checkpointing, recomputation, host/device offload and compression should be jointly optimized under memory/time constraints. AOTAutograd partitioning, jax.remat, Enzyme tapes and Checkmate are steps toward policy separated from local derivative definitions.

5. Cross-language and heterogeneous differentiation. The derivative path should cross library, CPU/GPU and language boundaries without pretending an opaque binary is transparent. Options include link-time IR, derivative companion symbols, verified summaries and custom-rule contracts. Diagnostics must identify the exact boundary and required annotation.

6. Higher-order, sparse, implicit and discontinuous programs. Nested AD is not sufficient when rules intentionally stop derivatives, perturbations collide, Jacobians are sparse, fixed points need implicit differentiation, or branches create distributional/discontinuous behavior. Systems need explicit semantics and specialized algorithms, not silent zeros.

7. Verification and derivative diagnostics. Custom rules are trusted code today. Frameworks should automatically test adjoint identities, compare JVPs and VJPs, use finite/complex-step checks where applicable, track nondifferentiable points and explain failures at source. Formal work on the denotational correctness of reverse AD suggests how transformation correctness can become more than testing.

8. Symbolic/AD hybrids. AceGen is evidence that symbolic simplification, common-subexpression optimization and AD can be co-designed. Modern systems should preserve recognized linear algebra, sparsity, invariants and exact identities long enough to generate smaller, more stable derivative programs, while falling back to program-level AD for general control and effects.

Conclusion

PyTorch is the pragmatic layered architecture: flexible eager recording plus a functional/compiler path. Zygote is the clearest demonstration of elegant language-IR pullback generation and of the cost of not fully modeling mutation. JAX offers the most systematic transform algebra through primitives, tracers, Jaxpr, linearization and transposition. Enzyme reaches deepest into real compiled programs and memory, gaining mutation and cross-language scope while surrendering some high-level meaning. Mooncake, Clad, Swift and Rust show the design space is still moving toward compiler-native, multi-level AD.

The most credible future is not one winner. It is a stack in which high-level mathematical rules, typed effect-aware transformation, low-level activity and alias analysis, and explicit memory scheduling cooperate—and in which an error at any level can still tell the programmer what happened in their source.

julia-latency-followup.md

Julia compilation latency after the package-image bump

A follow-up to the autodiff survey — evidence and measurements current through 22 July 2026

Short answer

The concern that Julia's build/compile experience has plateaued over the last two to three years is substantially justified, with one important qualification: Julia 1.9 and 1.10 delivered large improvements in caching and loading already compiled library code, and 1.11 made small-process startup cheaper. They did not make arbitrary new application specializations compile an order of magnitude faster. For compiler-heavy libraries such as Enzyme and Mooncake, the remaining first-call work is still measured in ten-second units on this machine.

The local results make the split unusually clear:

• ForwardDiff's cold package construction was 19 seconds, but its warm import is 0.13 seconds. A gradient of a new user function still costs 0.87 seconds in every fresh process.
• Zygote goes from 58 seconds cold to 0.89 seconds warm, but its new-function gradient costs 1.77 seconds.
• Mooncake goes from 141 seconds cold to 0.48 seconds warm, but preparing a derivative for the user function costs 9.76 seconds and allocates 789 MB while compiling.
• Enzyme goes from 277 seconds cold to 1.33 seconds warm, but its first derivative costs 11.45 seconds and allocates 848 MB while compiling.
• Repeating the call in the same process takes microseconds (Mooncake cache preparation: about 40 ms). Persistent sessions still amortize latency very effectively; short-lived tools and tests do not.

These are single diagnostic observations on Windows/Julia 1.12.6, not a cross-version benchmark. Exact outputs and scripts are in experiments/julia_latency/.

Three costs hiding behind “build time”

It helps to name the bills separately:

1. Installation and artifact acquisition resolves packages and downloads binary artifacts. Enzyme's LLVM artifacts are large, but download time was not included in the recorded cold runs because the depot already contained them.
2. Package precompilation infers package code and, since Julia 1.9, stores native code in package images. This is a disk-persistent, usually one-time cost for a particular Julia/package/preferences/CPU configuration.
3. Application specialization compiles methods for the concrete function and types encountered in the user's program. Unless that workload was included in a package image or sysimage, it repeats in every new Julia process.

The original autodiff run accidentally mixed these. The restricted execution sandbox denied Windows precompile-worker pipes. --compiled-modules=no bypassed the worker but also disabled serialized inference and native package-image reuse, forcing six compiler-intensive packages through one process. A controlled run of the identical normal command outside the sandbox succeeded. The five-minute episode was therefore real wall time but not a valid normal-package measurement.

What actually improved since the COVID-era work

Julia 1.6–1.8: inference and invalidation groundwork

The 2020 invalidation work explained why newly added methods could discard large trees of inferred code and why native caching would be ineffective until those trees became stable (Julia invalidation analysis). Julia 1.8 improved preservation of type-inference results. This was essential groundwork, but code generation still recurred in new processes.

Julia 1.9: the visible bump

Julia 1.9 added native code to package caches. The official 1.9 retrospective reported large TTFX reductions—often one to two orders of magnitude on curated workloads. It also stated the trade-off plainly: precompilation became roughly 10–50% more expensive and cache files became larger. PrecompileTools lets package authors choose representative workloads whose native code is stored.

This is the bump users remember. It moved work from “first use in every session” to “construct a package image once.” It did not remove the work, and it only caches signatures exercised or requested during package precompilation.

Julia 1.10: loading and parallel package construction

Julia 1.10 attacked warm loading and cache construction. Its release report shows the artificial 650-package OmniPackage load falling from 48.0 to 19.1 seconds, alongside fewer invalidations, better large-method-table scaling, parallel precompilation-on-import and pidfile coordination. LLVM work for large package images became parallel—but the same report notes that native-image parallelism is disabled on Windows because of COFF limitations. That matters for the very large Mooncake and Enzyme images measured here.

Julia 1.11–1.12: startup, distribution groundwork and tooling

Julia 1.11 excised standard libraries from the base system image and reported a small-script startup improvement from 113 ms to 92 ms in its release measurement. It also made package caches relocatable, groundwork for serving prebuilt caches rather than constructing all of them locally.

Julia 1.12 did not advertise another broad JIT-latency step. It added exactly the instrumentation needed to see the problem: --trace-compile-timing, @trace_compile and @trace_dispatch (1.12 highlights). Optional BOLT-built Julia/LLVM binaries improved compiler-heavy benchmarks around 10% (more in combination with PGO/LTO), but the documented BOLT path is Linux-only and is a build configuration, not a universal algorithmic reduction. Experimental --trim/JuliaC targets deployable static applications with restricted dynamic dispatch; it does not make an ordinary interactive session's novel methods free.

So the trajectory since 1.10 has emphasized cache portability, startup, deployment and observability. Those are meaningful improvements, but they match the subjective impression that application-specific compilation throughput has not had another dramatic jump.

What is moving in the 1.13 development cycle

The near-term work is relevant but not yet a reason to declare the cliff solved. The Julia project's January 2026 development summary reports a new API intended to let non-native compilers such as GPUCompiler cache owned CodeInstances during precompilation, addressing cases previously dropped by serialization. That could directly help compiler libraries with a cost shape like Enzyme's. It also reports broader AOT thread use (Documenter improved from 16 to 13 seconds in one test) and LLVM -time-trace integration. The February summary describes work on detachable background package precompilation. These are development items, not guarantees in stable Julia 1.12; backgrounding a build improves workflow but not total compute, while reusable non-native CodeInstances could remove real repeated work if the ecosystem can exploit them.

What the local traces say

Package	Cold image/load	Warm load	First user call	Same-process repeat	Direct image
ForwardDiff 1.4.1	19.31 s	0.129 s	0.870 s	11 µs	3.2 MiB
Zygote 0.7.11	57.59 s	0.886 s	1.767 s	32 µs	27.7 MiB
ReverseDiff 1.17.0	31.35 s	2.820 s	1.840 s	135 µs	19.9 MiB
Mooncake 0.5.40	141.01 s	0.475 s	9.758 s	39.5 ms	72.6 MiB
Enzyme 0.13.190	276.98 s	1.334 s	11.454 s	3 µs	46.3 MiB

The direct image size excludes dependency images. Cold times are order-dependent because later packages reused dependencies built earlier; they describe the user experience in this pinned depot, not intrinsic package scores.

Most importantly, invalidation was not the culprit in these first-call runs. Reported recompilation time was zero or negligible. The work was fresh specialization:

• ForwardDiff emitted 31 timed compile entries; its gradient and configuration signatures were tied to Main.objective and this concrete vector type.
• Zygote emitted 80. _pullback(Context, typeof(Main.objective), Vector{Float64}) took about 1.19 seconds, followed by a large specialized pullback type.
• Mooncake emitted 633, spread over SSA-to-ID mapping, compiler inference, generated AD statements and the final derived rule.
• Enzyme emitted 327. The outer Enzyme.gradient(..., typeof(Main.objective), Vector{Float64}) entry covered 12.31 seconds; a nested typeinf_local call using EnzymeInterpreter covered 4.99 seconds. Timed entries are nested and cannot be summed.

This is the fundamental cache boundary: a library package cannot precompile code for a function that does not exist until the user's Main module is loaded. Compiler-based AD magnifies it because specialization must compile not just the numeric kernel but an interpreter/IR transform and generated derivative code.

ReverseDiff shows a different cost shape. @time_imports attributed about 2.42 seconds directly to loading ReverseDiff even with a valid 19.9 MiB package image, and the process allocated roughly 214 MB during import. Package images remove code generation but still must be mapped/deserialized, validated and integrated with the current method world. More cached code can therefore exchange TTFX for cold construction, disk, memory and warm-load work.

This scaling pressure is recognized by ecosystem developers. The open Julia issue on monorepo subpackages uses OrdinaryDiffEq as a case study: after package binaries pushed first solves below a second, comprehensive precompilation approached an hour before tuning, and a rare solver contributed 1.5 seconds merely through lowering. The practical response was to split one repository into dozens of packages. That is dependency architecture compensating for package image granularity, not a faster compiler.

Why the plateau is technically hard

Julia deliberately combines open-world method extension, aggressive concrete-type specialization and interactive redefinition. These are central strengths. They also make reusable native code conditional on the method world, preferences, CPU target and concrete caller types.

Several tensions follow:

• Coverage versus construction cost. More precompile workloads reduce future latency but make installation/updates slower and package images larger.
• Genericity versus signature explosion. Numerical libraries work over broad type combinations. Caching a representative slice cannot cover every function, element type, dimension, AD nesting and backend.
• Interactivity versus stable AOT assumptions. New methods and redefinitions require world-age and invalidation checks. A conventional closed-world binary can omit them; Julia's REPL cannot generally do so.
• Compiler libraries compile compilers. Enzyme, Mooncake, GPUCompiler and similar tools specialize inference and IR passes on user programs. Their latency is partly productive work that a generic library image cannot know in advance.
• Cache granularity versus ecosystem usability. A package is a relatively coarse unit. Splitting optional functionality improves loading and image selectivity but increases release and dependency-management complexity.

What helps today

For exploratory work, a long-lived Julia process plus Revise remains the highest leverage answer. The second-call measurements—microseconds after 1–11 seconds of compilation—show why. Julia 1.12's binding redefinition work may reduce some restart pressure further.

For repeatable project workloads, create a small project-specific “Startup” package and place representative calls inside PrecompileTools.@compile_workload. The current PrecompileTools documentation explicitly supports this. It can persist the exact derivative paths missing from library-owned package images. The trade-off is rebuilding that cache when the workload, dependencies, preferences or relevant methods change.

For a stable application, a PackageCompiler sysimage with a precompile execution script can include application specializations and eliminate both package load and first-call costs. Its documentation warns that embedded packages take precedence over project versions, so sysimages are best treated as versioned build artifacts, not an invisible global default. JuliaC/--trim is worth watching for closed-world command-line deployment, but its current dynamic-dispatch restrictions make it a different product model.

During package development, disabling selected PrecompileTools workloads through Preferences can shorten rebuild cycles; turn them back on for user-facing or CI images. Reduce hard dependencies, move optional integrations into extensions, and use smaller packages when package-image granularity is the limiting factor.

For diagnosis on Julia 1.12:

using InteractiveUtils
@time_imports using MyPackage
@trace_compile my_first_workload()
@trace_dispatch my_first_workload()

Use SnoopCompile when invalidations or inference are suspected. In these AD measurements the timed trace was more decisive: it showed new user-function specialization, not recompilation.

Finally, never use --compiled-modules=no as a performance workaround. It is valuable for debugging broken caches/process restrictions, but it deliberately throws away the main latency improvements delivered since Julia 1.8/1.9.

Assessment

Julia is much better than its pre-1.9 self at reusing work that package authors anticipated. Warm imports of ForwardDiff, Zygote, Mooncake and Enzyme in this study are seconds or fractions of seconds rather than their 19–277 second package builds. That is a genuine success.

But the user's complaint survives intact at the application boundary. Novel specializations—especially transformations of arbitrary user programs—still compile at a rate that makes short-lived Julia processes unattractive. Recent releases have improved the placement, parallelism, portability and diagnosis of that work more than the raw throughput of doing it. The result is a system that feels excellent once warm and can still feel surprisingly 2019-like whenever a new function/type combination crosses a compiler-heavy path.

The next convincing leap would need at least one of: substantially faster inference/codegen, portable served package/application images, finer-grained cache reuse across caller boundaries, or a practical closed-world deployment mode that retains enough of Julia's generic programming model. Julia 1.12 has pieces of all four directions, but none yet erases the first-call cliff measured here.

torchscript-chronicle.md

The Rise and Fall of TorchScript

A chronicle, 2016–2026

---

Prologue: The Two-Framework Problem

PyTorch's founding virtue was define-by-run. The model is the Python that executes it; there is no graph to build, no session to run, and pdb works. Researchers defected en masse from the define-and-run frameworks, and the reason was never subtle — you could print a tensor in the middle of your forward pass.

The founding virtue was also the founding liability. As the PyTorch team put it in The Road to 1.0: PyTorch "pretty much never knows the operation you'll run next." No lookahead means no whole-program optimization, no serializable artifact, and no way to run the model without a Python interpreter in the loop.

Facebook, PyTorch's largest stakeholder, solved this by having two frameworks. Research happened in PyTorch; production happened in Caffe2, which shipped to over a billion phones and had the server inference, the TensorRT integration, and the ARM/Intel tuning. The path from one to the other was a manual rewrite — slow, error-prone, and performed by the people least interested in doing it.

The first bridge was ONNX (September 2017), riding a tracer PyTorch had carried since 0.3. Trace the model, emit a graph, hand it to Caffe2. It worked for feedforward vision models and broke on anything with a data-dependent branch.

In May 2018, the two projects announced they were merging.

---

Act I: The Promise (December 2018 – 2019)

PyTorch 1.0 shipped in December 2018, and its headline feature was torch.jit. The pitch was carefully non-threatening: your code continues to work as-is. Production-readiness was an opt-in annotation. You wrote normal PyTorch, added a decorator, and got back an artifact that ran in a C++ runtime with no Python anywhere.

Two doors led into the graph:

• torch.jit.trace — run the model on example inputs, record the dispatched ATen ops. Zero code changes. Silently discards control flow and bakes in shapes.
• torch.jit.script — compile the Python source directly. Handles loops and branches. Requires that your source lie inside TorchScript: a statically typed subset of Python.

TorchScript proper was the ambitious door, and for a while the ambition kept growing. It began as an export format and became a compiler. By 2019 it had a typed IR with control flow, a profiling executor that specialized on observed shapes, alias analysis, a kernel fuser, and — as ezyang's PyTorch internals records — its own symbolic autodiff, with derivative formulas written in TorchScript source and stored in symbolic_script.cpp.

It also became load-bearing for nearly everything else. PyTorch Mobile (1.3, 2019) ran on the TorchScript Lite Interpreter. ONNX export went through the JIT. Quantization, TorchServe, and the LibTorch C++ API all assumed a scripted module. TorchScript was no longer a feature; it was the floor.

---

Act II: The Cracks (2019 – 2021)

The trouble was that TorchScript is not Python. It is a different language wearing Python's syntax, and the differences surface exactly when your model is interesting:

• static types, with Tensor inferred for unannotated arguments
• no *args/**kwargs, no inheritance, no arbitrary Python objects
• homogeneous containers only — no Dict[str, Union[Tensor, int]]
• Optional narrowing rules that reject code Python accepts
• no closures over non-scriptable state, no third-party library calls

And the failure mode was all-or-nothing. One unsupported construct anywhere in the module tree and torch.jit.script raised, often with an error naming a line you did not write, inside a library you did not control. The PyTorch 2 paper is blunt about this in its own retrospective: scripting "suffers from the all-or-nothing limitation," and users were "forced to rewrite" their models to satisfy it.

So "scriptability" became a property libraries had to maintain. torchvision annotated its entire surface and kept CI to defend it. Most of the ecosystem did not, and could not — HuggingFace Transformers and detectron2 shipped with tracing as the practical option, which meant shipping graphs that were silently wrong the moment a sequence length changed.

Meanwhile the compiler ambitions were fragmenting from within:

• Three fusers in four years. The original fuser, then TensorExpr/NNC, then NVFuser. Each was competent; none could reach across a whole model, because differentiable-subgraph formation kept shattering on unsupported ops, mutation, aliasing, and control flow.
• A second derivative database. symbolic_script.cpp covered on the order of 100 ops. derivatives.yaml, the eager engine's source of truth, covered 687, against ~2,600 in native_functions.yaml. Two hand-maintained tables for the same mathematics, guaranteed to drift, and nobody adding an operator wanted to write its derivative twice.
• A rival IR. torch.fx arrived in 2020–21 offering Python-to-Python symbolic tracing for graph transformation — a tacit admission that TorchScript IR was the wrong substrate for the transformations people actually wanted.
• A rival capture path. Lazy Tensors, deferring at the C++ dispatcher to feed XLA.

By 2021 PyTorch had four graph-capture mechanisms — jit.trace, jit.script, fx.symbolic_trace, LazyTensor — each with different coverage, and every one of them all-or-nothing.

---

Act III: The Insight (2021 – 2022)

TorchDynamo's contribution was not a better tracer. It was giving up on total capture.

Dynamo hooks CPython's frame evaluation API (PEP 523) and rewrites bytecode before execution. When it meets Python it cannot trace, it does not fail — it compiles what it has, drops into the interpreter for the offending code, and resumes with a new graph. Correctness is protected by guards that trigger recompilation when assumptions break. The result: partial capture that never forces a rewrite. Graph breaks cost performance, and performance is a thing users can trade away. Rewriting the model is not.

Three companion decisions finished the job:

Component	What it replaced
AOTAutograd	The second derivative table. It traces the forward, runs the real eager autograd engine over it, and captures the joint forward-backward graph — then min-cut partitions it to trade saved activations against recompute. Coverage becomes total by construction.
PrimTorch	The impossible backend contract. ~2,600 ATen ops decomposed to a few hundred primitives, so writing a backend stops being a career.
TorchInductor	The fuser churn. Generate Triton for GPU and C++ for CPU instead of hand-maintaining a fusion engine.

Announced December 2022, released March 2023. The paper's own measured claim was that Dynamo captures graphs more robustly than prior approaches at minimal overhead — with the prior approaches in question being PyTorch's own.

---

Act IV: The Slow Retirement (2023 – 2026)

There was no execution. TorchScript simply stopped being developed, and by mid-2023 users were opening issues asking for a deprecation calendar that did not yet exist.

What happened instead was decomposition, and it is the most instructive part of the story. TorchScript had been asked to be five things simultaneously. Its successors are five separate things:

TorchScript's job	Successor
Graph capture for optimization	TorchDynamo
Autodiff over captured graphs	AOTAutograd
Serializable, portable deployment IR	torch.export → ExportedProgram
On-device / mobile runtime	ExecuTorch (beta 2024, 1.0 in 2025)
Python-free C++ serving	AOTInductor

Each successor is narrower, and narrower turned out to be the whole point. torch.export can demand a sound, normalized, fully-shaped graph precisely because it is not also trying to be a JIT. Dynamo can tolerate graph breaks precisely because it is not also trying to be a serialization format.

The deprecation notices landed on the docs during the 2.9 cycle in late 2025. Formal deprecation came in PyTorch 2.10, with torch.export named as the replacement for the trace and script APIs, and ExecuTorch for the embedded runtime.

---

Epilogue: What It Got Right

It is easy to read this as a failure. It is more accurate to read it as a correct diagnosis with a wrong prescription.

TorchScript was right that eager execution cannot be deployed or globally optimized without a graph. Right that the graph needs a real IR with types, control flow, and alias information — the FX-based stack inherited that requirement rather than escaping it. Right that the artifact must be serializable and Python-free. Every one of those premises survived into PyTorch 2.

The wrong bet was on who does the accommodating. TorchScript asked users to write in a restricted dialect so the compiler could understand them — which contradicted the exact property that made people choose PyTorch. Dynamo asked the compiler to understand Python as actually written, and to degrade gracefully where it could not. Same goal, opposite direction of effort, and the burden moved off the user.

The second wrong bet was duplication. The symbolic autodiff table is the crisp example: a parallel implementation of knowledge that already existed elsewhere in the codebase, which could therefore only ever be a subset, and could only ever drift. AOTAutograd's advantage is not a cleverer algorithm; it is reusing the one implementation that was already correct. That principle generalizes well past compilers.

TorchScript is not gone so much as dissolved. The dispatcher it was built against, the ATen operator schemas, the alias annotations, the notion of a typed tensor IR — all of it is still load-bearing underneath torch.compile. What was retired was the demand that users learn a second language to reach it.

---

Sources: PyTorch blog (The Road to 1.0, 2.11 release notes), Caffe2 merge announcement (May 2018), Ansel et al., "PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation," ASPLOS '24; PyTorch source tree (symbolic_script.cpp, derivatives.yaml, native_functions.yaml), counts taken from main, July 2026.