Differentiation of High-Level Language Semantics

Michael Innes

TL;DR: Describes how the Zygote AD is able to differentiate language features, like closures, mutation and concurrency.

Abstract: Though analytic differentiation (AD) is a program transformation, AD tools have typically supported only very limited program representations, consisting of primitive mathematical operations and basic structured control flow. Zygote, an AD for the Julia language, instead operates on Julia code. This presents an interesting challenge for the AD implementor: the program representation now contains not just mathematical operations, but arbitrary control flow, user-defined functions, recursion, data structures, mutation, metaprogramming, foreign function calls, specialised hardware, and even concurrency and parallelism primitives. This paper explains how Zygote handles these high-level features safely and efficiently, making an unusually large set of Julia programs differentiable.