Differentiation: AD backend strategies

The CTBase.Differentiation submodule provides an automatic differentiation (AD) backend abstraction built on the Strategies contract. It exposes a small set of differentiation primitives — gradients, derivatives, partial derivatives and Jacobian–vector products — behind a single strategy type, so that consumers (Hamiltonian systems, flows, differential geometry) never depend on a concrete AD package directly.

Prerequisites

Read the Implementing a Strategy and Options System guides first: an AD backend is a strategy with a single :ad_backend option.

Overview

The module is organised around a two-level strategy contract:

Level	Symbol	Role
Family	Differentiation.AbstractADBackend	abstract strategy; defines the AD contract
Concrete	Differentiation.DifferentiationInterface	wraps a DifferentiationInterface.jl backend

The concrete strategy stores a single option, :ad_backend, holding an ADTypes.jl backend. The default is device-dependent: AutoForwardDiff() on CPU and AutoMooncake() on GPU (see Device parameterization below). ADTypes is a hard dependency of CTBase, so the default backend markers are always available.

The contract has eleven methods. Two are resolved in core and always available: Differentiation.ad_backend, the accessor returning the wrapped ADTypes backend, and Differentiation.check_backend, which checks that the package behind that backend is loaded (see Checking that the AD package is loaded). The nine differentiation primitives below live in an extension.

The differentiation methods live in an extension

The contract methods (gradient, derivative, differentiate, pushforward, hamiltonian_gradient, variable_gradient, pseudo_hamiltonian_gradient, pseudo_hamiltonian_control_gradient, pseudo_variable_gradient) are implemented in the CTBaseDifferentiationInterface package extension. They become available only once DifferentiationInterface and a concrete AD package (e.g. ForwardDiff) are loaded. Without the extension, every contract method throws a CTBase.Exceptions.NotImplemented with a helpful message; with the extension but without the package of the selected AD backend, it throws a CTBase.Exceptions.ExtensionError naming that package.

Building a backend

Construct a Differentiation.DifferentiationInterface strategy with the default AD type:

backend = Differentiation.DifferentiationInterface()

DifferentiationInterface{CPU} (instance, id=:di)
└─ ad_backend = AutoForwardDiff  [computed]
Tip: use describe(DifferentiationInterface{CPU}) to see all available options.

Pass an explicit ADTypes backend through the ad_backend option (the aliases backend and ad resolve to the same option):

Differentiation.DifferentiationInterface(;
    ad_backend = ADTypes.AutoForwardDiff(),
)

DifferentiationInterface{CPU} (instance, id=:di)
└─ ad_backend = AutoForwardDiff  [user]
Tip: use describe(DifferentiationInterface{CPU}) to see all available options.

The wrapped AD type is recovered with Differentiation.ad_backend:

Differentiation.ad_backend(backend)

AutoForwardDiff()

Because it is a regular strategy, its metadata is introspectable at the type level:

Strategies.metadata(Differentiation.DifferentiationInterface)

StrategyMetadata with 1 option:
│  
└─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: AutoForwardDiff() [computed])
   description: DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU).

Device parameterization

DifferentiationInterface is parameterized on the execution device P via the Strategy Parameters system. Two built-in parameters are available: Strategies.CPU and Strategies.GPU.

The :ad_backend default is computed from the parameter — each device gets a different default AD backend:

Parameter	Default :ad_backend	Notes
CPU	AutoForwardDiff()	General-purpose, well-tested
GPU	AutoMooncake()	Validated end-to-end on CuArray, including through a mutating in-place RHS

DifferentiationInterface() (no parameter) defaults to CPU, so existing call sites are unaffected:

backend_cpu = Differentiation.DifferentiationInterface{Strategies.CPU}()

DifferentiationInterface{CPU} (instance, id=:di)
└─ ad_backend = AutoForwardDiff  [computed]
Tip: use describe(DifferentiationInterface{CPU}) to see all available options.

backend_gpu = Differentiation.DifferentiationInterface{Strategies.GPU}()

DifferentiationInterface{GPU} (instance, id=:di)
└─ ad_backend = AutoMooncake  [computed]
Tip: use describe(DifferentiationInterface{GPU}) to see all available options.

The wrapped AD backend reflects the device default:

Differentiation.ad_backend(backend_cpu)

AutoForwardDiff()

Differentiation.ad_backend(backend_gpu)

AutoMooncake()

The per-parameter metadata shows the different defaults side by side:

Strategies.metadata(Differentiation.DifferentiationInterface{Strategies.CPU})

StrategyMetadata with 1 option:
│  
└─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: AutoForwardDiff() [computed])
   description: DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU).

Strategies.metadata(Differentiation.DifferentiationInterface{Strategies.GPU})

StrategyMetadata with 1 option:
│  
└─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: AutoMooncake() [computed])
   description: DifferentiationInterface.jl backend for GPU execution. Default: AutoMooncake(), validated end-to-end on CuArray including through a mutating in-place RHS. AutoForwardDiff() does not work on GPU (scalar-indexes a CuArray). AutoZygote() can be selected explicitly; it was found unreliable in some device call contexts (an unexpected KernelException was observed on a non-mutating call).

GPU backend selection

AutoForwardDiff() does not work on GPU because it scalar-indexes a CuArray. AutoMooncake() is the validated default for GPU execution. AutoZygote() can be selected explicitly but was found unreliable in some device call contexts (an unexpected KernelException was observed on a non-mutating call).

The device parameter can be overridden explicitly at construction time:

backend_gpu_fd = Differentiation.DifferentiationInterface{Strategies.GPU}(;
    ad_backend = ADTypes.AutoForwardDiff(),
)
Differentiation.ad_backend(backend_gpu_fd)

AutoForwardDiff()

Checking that the AD package is loaded

Selecting an ADTypes backend does not load the package that implements it: AutoMooncake() needs using Mooncake, and even the default AutoForwardDiff() needs using ForwardDiff. Every differentiation primitive checks this before differentiating and throws a CTBase.Exceptions.ExtensionError naming the missing package. The check is a compile-time constant for a concrete backend, so it costs nothing on the hot path.

Differentiation.check_backend runs the same check explicitly — call it to fail early, before a long computation such as an integration. It returns nothing when the package is loaded:

Differentiation.check_backend(backend)   # ForwardDiff is loaded

Mooncake is not loaded in this session, so the GPU default fails with an actionable message:

julia> Differentiation.check_backend(backend_gpu)

Sparse backends are checked through their dense backend, and SecondOrder backends through both their outer and inner backends.

Differentiation primitives

Gradient and derivative

Differentiation.gradient differentiates a scalar function of a vector, and Differentiation.derivative a scalar function of a scalar:

f(x) = sum(x .^ 2)          # ∇f = 2x
Differentiation.gradient(backend, f, [1.0, 2.0])

2-element Vector{Float64}:
 2.0
 4.0

g(t) = t^3                  # g'(t) = 3t²
Differentiation.derivative(backend, g, 2.0)

12.0

Partial derivative at a slot

Differentiation.differentiate computes the partial derivative of f with respect to the argument at a compile-time slot Val(Slot), holding the other arguments fixed. The active argument comes first, the frozen ones follow in slot order:

# H(x, p) = ½‖p‖² + ‖x‖²  →  ∂H/∂x = 2x,  ∂H/∂p = p
H(x, p) = 0.5 * sum(p .^ 2) + sum(x .^ 2)
x = [1.0, 2.0]; p = [3.0, 4.0]

# active = x (slot 1), const = p
∂x = Differentiation.differentiate(
    backend, H, Val(1), x, p,
)
# active = p (slot 2), const = x
∂p = Differentiation.differentiate(
    backend, H, Val(2), p, x,
)
(∂x, ∂p)

([2.0, 4.0], [3.0, 4.0])

A scalar slot yields a scalar derivative:

# f(t, x) = t * x[1]  →  ∂f/∂t = x[1]
ft(t, x) = t * x[1]
Differentiation.differentiate(backend, ft, Val(1), 2.0, [3.0, 4.0])

3.0

Pushforward (Jacobian–vector product)

Differentiation.pushforward computes the directional derivative of f at x along dx, freezing the other arguments:

A = [0.0 1.0; -1.0 0.0]
X(x) = A * x                 # J_X = A  →  pushforward = A·dx
Differentiation.pushforward(backend, X, Val(1), [1.0, 2.0], [5.0, 6.0])

2-element Vector{Float64}:
  6.0
 -5.0

Hamiltonian gradients

The two domain-specific methods operate directly on a Data.Hamiltonian (see the Data guide). Differentiation.hamiltonian_gradient returns (∂H/∂x, ∂H/∂p), non-negated — the caller applies the signs of Hamilton's equations (ẋ = ∂H/∂p, ṗ = -∂H/∂x):

# H(x, p) = ½(‖x‖² + ‖p‖²)  →  ∂H/∂x = x,  ∂H/∂p = p
h = Data.Hamiltonian((x, p) -> 0.5 * (sum(abs2, x) + sum(abs2, p)))
∂x, ∂p = Differentiation.hamiltonian_gradient(
    backend, h, 0.0, [1.0, 2.0], [3.0, 4.0], nothing,
)
(∂x, ∂p)

([1.0, 2.0], [3.0, 4.0])

For a non-fixed Hamiltonian, Differentiation.variable_gradient returns ∂H/∂v:

# H(x, p, v) = ½(‖x‖² + ‖p‖² + ‖v‖²)  →  ∂H/∂v = v
hv = Data.Hamiltonian(
    (x, p, v) -> 0.5 * (sum(abs2, x) + sum(abs2, p) + sum(abs2, v));
    is_variable=true,
)
Differentiation.variable_gradient(
    backend, hv, 0.0, [1.0, 2.0], [3.0, 4.0], [5.0, 6.0],
)

2-element Vector{Float64}:
 5.0
 6.0

Pseudo-Hamiltonian gradients

The two pseudo-Hamiltonian methods operate directly on a Data.PseudoHamiltonian (see the Data guide). Differentiation.pseudo_hamiltonian_gradient returns (∂H̃/∂x, ∂H̃/∂p), non-negated — the caller applies the signs of Hamilton's equations:

# H̃(x, p, u) = ½(‖x‖² + ‖p‖²) + u²  →  ∂H̃/∂x = x, ∂H̃/∂p = p
ph = Data.PseudoHamiltonian(
    (x, p, u) -> 0.5 * (sum(abs2, x) + sum(abs2, p)) + u^2,
)
∂x, ∂p = Differentiation.pseudo_hamiltonian_gradient(
    backend, ph, 0.0, [1.0, 2.0], [3.0, 4.0], 2.0, nothing,
)
(∂x, ∂p)

([1.0, 2.0], [3.0, 4.0])

Differentiation.pseudo_hamiltonian_control_gradient returns ∂H̃/∂u, which vanishes along optimal trajectories by the Pontryagin Maximum Principle stationarity condition:

# ∂H̃/∂u = 2u = 4.0
Differentiation.pseudo_hamiltonian_control_gradient(
    backend, ph, 0.0, [1.0, 2.0], [3.0, 4.0], 2.0, nothing,
)

4.0

For a non-fixed pseudo-Hamiltonian, the v argument carries the variable (here one-dimensional, hence a scalar):

phv = Data.PseudoHamiltonian(
    (x, p, u, v) -> 0.5 * (sum(abs2, x) + sum(abs2, p)) + u^2 + v;
    is_variable=true,
)
Differentiation.pseudo_hamiltonian_control_gradient(
    backend, phv, 0.0, [1.0, 2.0], [3.0, 4.0], 2.0, 5.0,
)

4.0

Pseudo-Hamiltonian variable gradient

For a non-fixed pseudo-Hamiltonian, Differentiation.pseudo_variable_gradient returns ∂H̃/∂v with the control u held constant — a partial derivative, not the total derivative ∂/∂v[H̃(t, x, p, u(t,x,p,v), v)]:

# H̃(x, p, u, v) = ½(‖x‖² + ‖p‖²) + u² + v  →  ∂H̃/∂v = 1.0
phv = Data.PseudoHamiltonian(
    (x, p, u, v) -> 0.5 * (sum(abs2, x) + sum(abs2, p)) + u^2 + v;
    is_variable=true,
)
Differentiation.pseudo_variable_gradient(
    backend, phv, 0.0, [1.0, 2.0], [3.0, 4.0], 2.0, 5.0,
)

1.0

This differs from Differentiation.variable_gradient whenever the feedback is not stationary (∂H̃/∂u ≠ 0), because the total derivative would also account for ∂H̃/∂u · ∂u/∂v.

The contract without the extension

Every differentiation primitive has a fallback on AbstractADBackend that throws CTBase.Exceptions.NotImplemented (see the Exceptions guide). This is what users hit when DifferentiationInterface is not loaded, and what custom backends must override (check_backend defaults to a no-op, so a custom backend overrides it only if it can detect a missing package):

julia> struct MyBackend <: Differentiation.AbstractADBackend
           options::Strategies.StrategyOptions
       end

julia> Differentiation.gradient(
           MyBackend(Strategies.StrategyOptions()), x -> sum(x), [1.0],
       )
ERROR: NotImplemented
│
│  gradient not implemented for Main.MyBackend
│
│  Method   gradient(backend::Main.MyBackend, f::Function, x)
│
│  Context  AD backend contract
│  Hint     Load CTBaseDifferentiationInterface (load DifferentiationInterface)
└─

See also

• Differentiation.AbstractADBackend, Differentiation.DifferentiationInterface — the strategy types.

• Differentiation.check_backend — fail early when the package of the AD backend is not loaded.

• Implementing a Strategy, Options System — the contract these build on.

• Data — Hamiltonian and PseudoHamiltonian wrappers consumed by the gradient methods.

• Exceptions — NotImplemented, raised by the contract fallbacks, and ExtensionError, raised when the AD package is not loaded.