DifferentiationInterface.jl

DifferentiationInterfaceTest provides testing and benchmarking utilities for automatic differentiation (AD) in Julia. It is built on top of DifferentiationInterface.jl and is designed to help developers evaluate a specific function against various AD backends.

Key capabilities include:

• Backend Compatibility: Determining which AD backends can successfully differentiate a given function.
• Performance Benchmarking: Measuring how fast different backends perform differentiation, including runtime and memory allocations.
• Correctness & Stability: Running correctness tests and type stability tests.
• Execution Tracking: Counting the number of calls made to the function during differentiation.
• Custom Scenarios: Defining custom test scenarios for specialized evaluation.

Source transformation backends (like Zygote.jl, Enzyme.jl, and Mooncake.jl) transform the code itself to compute derivatives. Compatibility requirements vary:

• Zygote.jl: Generally avoids mutation. If you encounter ERROR: Mutating arrays is not supported, you must define a custom reverse rule (rrule) using ChainRulesCore.jl.
• Enzyme.jl: Supports a wider range of code by targeting lower-level representations. It has its own extensible rule system that uses input activity annotations.
• Mooncake.jl: Handles a large subset of Julia programs out-of-the-box with a rule system that is generally considered less expressive (and thus potentially easier to use) than Enzyme's.

DifferentiationInterface provides argument wrappers to control how specific arguments are treated during the differentiation process. Use these to specify constants or manage memory/computation via caching.

• Context: Provides context for differentiation.
• Constant: Marks an argument as a constant that should not be differentiated.
• Cache: Used to store intermediate values to optimize differentiation.
• ConstantOrCache: A hybrid wrapper that can act as either a constant or a cache.

DifferentiationInterface supports two specific function signatures for computing derivatives. In both cases, the first argument x is the only 'active' argument (the one being differentiated), and all subsequent arguments are treated as 'contexts'.

1. Out-of-place functions: y = f(x, contexts...) where y is the returned output.
2. In-place functions: f!(y, x, contexts...) where y is the output to be overwritten and the function returns nothing.

All differentiation operators return the (partial) derivative of y with respect to x.

# Out-of-place
y = f(x, contexts...)

# In-place
f!(y, x, contexts...)

In forward mode, DI computes the Jacobian-Vector Product (JVP) $\dot{y} = (\frac{\partial f}{\partial x}) \dot{x}$.

To achieve this, the AD backend is run on the state transformation $\phi$ with input tangents $(\dot{x}_0, \dot{c}_0, \dot{s}_0, \dot{y}_0)$. Because of the implementation hypotheses, the output tangent $\dot{y}_1$ simplifies to: $$\dot{y}_1 = \frac{\partial y_1}{\partial x_0} \dot{x}_0 + \frac{\partial y_1}{\partial c_0} \dot{c}_0$$

Requirement: To obtain the correct JVP, you must ensure the constant argument tangent is zeroed: $\dot{c}_0 = 0$.

If your function has multiple arguments but you only want the derivative with respect to one (the active argument x), you can pass the other arguments as Contexts. This is often more performant than creating closures.

Every context argument must:

1. Be wrapped in a subtype of Context.
2. Follow the active argument x in the function signature.

Example function signatures:

• y = f(x, contexts...)
• f!(y, x, contexts...)

Support for complex derivatives varies significantly across AD backends and conventions.

DifferentiationInterface operates under the assumption that if complex numbers are involved, the function being differentiated is holomorphic.

Important Notes:

• This functionality is experimental and does not carry public API guarantees.
• If your function is non-holomorphic, you must manually separate the real and imaginary parts of the function before differentiating.

For second-order operators like second_derivative, hessian, and hvp, you have two options:

1. Single Backend: Use a single backend that natively supports second-order derivatives.
2. Combined Backends: Use the SecondOrder(outer_backend, inner_backend) struct to compose two backends. The inner_backend is called first, and the outer_backend differentiates the code generated by the inner one.

Best Practice: For optimal performance, use a forward-mode outer backend with a reverse-mode inner backend.

Warning: Second-order AD is complex; many backend combinations (including a backend paired with itself) may fail. Experimentation is often required.

backend = SecondOrder(outer_backend, inner_backend)

DifferentiationInterfaceTest.jl uses Scenario objects to encapsulate all information required for a test. A Scenario includes:

• The operator category (e.g., :gradient or :jacobian).
• The behavior of the operator: :in (in-place) or :out (out-of-place).
• The function f to be differentiated.
• The input x (and optional tangents/contexts).
• The reference result res1 (the ground truth first-order result).
• Optional second-order results res2.
• Arguments prep_args used during preparation.

Example of creating scenarios for a gradient test:

f(x::AbstractArray) = sum(sin, x)
∇f(x::AbstractArray) = cos.(x)

xv = rand(Float32, 3)
scenarios = [
    Scenario{:gradient, :out}(f, xv; res1=∇f(xv))
]

In reverse mode, DI computes the Vector-Jacobian Product (VJP) $\bar{x} = (\frac{\partial f}{\partial x})^* \bar{y}$.

To achieve this, the AD backend is run on the state transformation $\phi$ with output sensitivities $(\bar{x}_1, \bar{c}_1, \bar{s}_1, \bar{y}_1)$. Because of the implementation hypotheses, the input sensitivity $\bar{x}_0$ simplifies to: $$\bar{x}_0 = \bar{x}_1 + (\frac{\partial s_1}{\partial x_0})^* \bar{s}_1 + (\frac{\partial y_1}{\partial x_0})^* \bar{y}_1$$

Requirement: To obtain the correct VJP, you must ensure the active and scratch sensitivities are zeroed: $\bar{x}_1 = 0$ and $\bar{s}_1 = 0$.

DifferentiationInterface.jl provides a unified syntax to swap between different Automatic Differentiation (AD) backends. This allows you to easily compare performance (e.g., comparing Forward-mode vs. Reverse-mode AD) without changing your core logic.

Example of switching from ForwardDiff to Zygote:

import Zygote: Zygote
backend2 = AutoZygote()
g = gradient(f, backend2, x)