Zygote.jl

Zygote performs Automatic Differentiation (AD) by transforming functions into pullbacks. A pullback (represented by the function J in conceptual examples) takes a function and its inputs, returning the result of the forward pass and a 'back' function. The 'back' function (the adjoint) takes the gradient of the output and returns the gradient of the inputs.

For composite functions, Zygote chains these pullbacks together: the backward pass of a composite function $f(g(x))$ is the composition of the pullbacks of $f$ and $g$, propagating gradients from the output back to the input.

The gradient function is a high-level wrapper around Zygote.pullback. A pullback is a fundamental operation that returns two values: the result of the original function and a pullback function (often called back).

Mathematically, if $y = f(x)$, the pullback $\mathcal{B}_y$ implements the vector-Jacobian product (VJP). Given a gradient $\bar{y}$ (the derivative of a loss with respect to the output), the pullback computes the gradient with respect to the input: $\bar{x} = \mathcal{B}_y(\bar{y})$.

In Zygote, pullbacks implement the adjoint of the Jacobian, performing a left-multiplication ($v'J$).

While both represent a program's execution, they differ in fidelity:

• IR (Intermediate Representation): A lower-level version of source code (e.g., using goto instead of loops) that preserves control flow semantics, making it easier to manipulate programmatically.
• Trace: A runtime recording of mathematical operations (often a Wengert List). Traces unroll and inline all control flow, functions, and data structures, which can lead to a loss of original program semantics.

When a function f returns a real number (e.g., abs2(c)), Zygote treats the complex input c = x + yi as a pair of reals (x, y). The resulting gradient is defined as the adjoint $\bar{c} = \frac{\partial z}{\partial x} + \frac{\partial z}{\partial y}i$. This pragmatic definition is suitable for gradient descent applications.

Note that this definition is not equivalent to the mathematical complex derivative unless the function is holomorphic.

You can use try/catch blocks in code being differentiated, provided that no exception is actually thrown during the execution. If the control flow enters a catch block, Zygote will fail because it cannot differentiate through the exception handling mechanism.

Solution: Use more graceful error handling, such as returning nothing or a sentinel value instead of relying on exceptions.

function safe_sqrt(x)
  try
    sqrt(x)
  catch
    0.
  end
end

# Works if no error is thrown
gradient(safe_sqrt, 4.) 

# Fails if error is thrown (control flow enters catch)
pullback(safe_sqrt, -1.)

While Zygote theoretically supports derivatives of derivatives, it faces several challenges:

• Many existing rules are not themselves differentiable (they use forbidden mutation).
• Complexity grows rapidly as Zygote differentiates its own un-optimized output.
• Reverse-over-reverse is often inefficient.

Recommendation: For second-order derivatives (like the Hessian), it is often better to use a different AD system. For example, the hessian function uses ForwardDiff over Zygote.

In Zygote, a pullback is the core mechanism for reverse-mode differentiation. Given a function $y = f(x)$, the pullback is the function $\bar{x} = \text{back}(\bar{y})$ returned by y, back = Zygote.pullback(f, x). This operation is mathematically equivalent to a Vector-Jacobian product.

When defining new pullbacks (for example, using the @adjoint macro), you are defining the adjoint of the Jacobian. This process is used to determine the sensitivity—how a change in a value $x$ affects a scalar loss $l$.

Zygote cannot differentiate expressions that call external libraries (C or Fortran) via ccall because the underlying code is not in Julia.

Solution: The only way to differentiate a foreign call is to write a custom ChainRulesCore.rrule that manually defines the gradient (pullback).

# A function calling a C function
jclock(x) = ccall(:clock, Int32, ()) * x

# Custom rrule to enable differentiation
function ChainRulesCore.rrule(::typeof(jclock), x)
  y = jclock(x)
  pb(ȳ) = (ChainRulesCore.NoTangent(), ȳ * y)
  return y, pb
end

# Now this works
gradient(jclock, rand())

Zygote has limited and buggy support for mutating fields in mutable structs (e.g., x.a = val).

Recommendation: Use immutable structs instead. If you need to 'modify' a value, use tools like @set from Accessors.jl, which returns a new object without side-effects.

For holomorphic (analytic) functions—functions that can be represented by a Taylor series and do not explicitly use real, imag, or conj (e.g., exp, log)—the complex derivative is the conjugate of the gradient of the real part.

To obtain the mathematically correct complex derivative for a holomorphic function, take the gradient of the real part and apply conj to the result.