Skip to content

AD backend ​

ad, Poisson, ∂ₜ and @Lie compute their derivatives by automatic differentiation (AD); Lift computes none. This page shows which AD backend they use, and how to change it.

julia
using OptimalControl

The default ​

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

The default backend calls DifferentiationInterface with ForwardDiff, on the CPU. OptimalControl depends on both packages, so nothing else needs to be installed. describe lists the options of the backend, and their defaults on the CPU and on the GPU:

describe(:di)
julia
describe(:di)
DifferentiationInterface (strategy)
├─ id: :di
├─ hierarchy: DifferentiationInterface → AbstractADBackend → AbstractStrategy
├─ description: AD backend wrapping DifferentiationInterface.jl backends (e.g., AutoForwardDiff).
├─ family: AbstractADBackend
├─ default parameter: CPU
├─ parameters: CPU, GPU
│
├─ computed options for CPU:
│  └─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: AutoForwardDiff() [computed])
│     description: DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU).
│
└─ computed options for GPU:
   └─ ad_backend (backend, ad)::ADTypes.AbstractADType (default: ADTypes.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).

Choosing a backend ​

A backend is built with CTBase.Differentiation.DifferentiationInterface, whose option ad_backend is a backend type of ADTypes. DifferentiationInterface re-exports these types; add it to your environment to use them. For example, ForwardDiff with a chunk size of one:

julia
using DifferentiationInterface: AutoForwardDiff

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

Every operation accepts ad_backend= for one call:

julia
X(x) = [x[2], -x[1]]
f(x) = x[1]^2 + x[2]^2 + x[1]

ad(X, f; ad_backend=backend)([1.0, 2.0])   # 2x₁x₂ - 2x₂x₁ + x₂
2.0

dg_ad_backend! changes the default for the rest of the session. Keep the previous one to restore it:

julia
previous = dg_ad_backend()
dg_ad_backend!(backend)
ad(X, f)([1.0, 2.0])
2.0
julia
dg_ad_backend!(previous)
dg_ad_backend()
DifferentiationInterface{CPU} (instance, id=:di)
└─ ad_backend = AutoForwardDiff  [computed]
Tip: use describe(DifferentiationInterface{CPU}) to see all available options.

GPU ​

The GPU backend is built with the parameter GPU. Its default is Mooncake, since ForwardDiff does not run on GPU arrays:

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

It differentiates functions of CuArrays, and needs CUDA and Mooncake to be loaded, as Flow(…; method=:gpu) does (see Flows overview).

See also ​