Tutorial

This tutorial introduces Mooncake's API one step at a time, starting with gradients and building up to Jacobians, Hessian-vector products, and full Hessians. Each derivative follows the same two-step pattern: prepare a cache once, then call a fast run function as many times as you like.

import Mooncake

Computing gradients

Suppose you want to differentiate the function

f(x) = sum(abs2, x)
x = float.(1:3)

3-element Vector{Float64}:
 1.0
 2.0
 3.0

The simplest entry point is Mooncake.value_and_gradient!!. Before evaluating it, prepare a cache once on a typical input — this is where Mooncake compiles the differentiation rule:

typical_x = rand(3)
cache = Mooncake.prepare_gradient_cache(f, typical_x)

Mooncake.Cache
  mode: reverse
  friendly_tangents: false
  inputs: 1
  input_1: Vector{Float64} (size (3,))
  output: Float64 (scalar)

The contents of typical_x do not matter; only its type and shape. Subsequent calls on inputs with matching shape are fast:

val, grad = Mooncake.value_and_gradient!!(cache, f, x)
(val, grad)

(14.0, (Mooncake.NoTangent(), [2.0, 4.0, 6.0]))

The returned grad has one entry per argument, preceded by the entry for f itself: here (df, dx), where df is the gradient with respect to any differentiable fields of f (NoTangent() since f is not a callable struct), and dx is the gradient with respect to x.

The cache owns the gradient buffers, so grad aliases storage inside cache. If you need to keep it across calls, take a copy or deepcopy first; otherwise the next call to value_and_gradient!! will overwrite it.

Functions of several arguments work the same way: pass each argument to both prepare_gradient_cache and value_and_gradient!!, and grad gains one entry per argument, again led by the entry for f.

g(x, a, b) = a * f(x) + b
typical_a, typical_b = 1.0, 1.0
a, b = 42.0, 3.14

cache = Mooncake.prepare_gradient_cache(g, typical_x, typical_a, typical_b)
val, grad = Mooncake.value_and_gradient!!(cache, g, x, a, b)
(val, grad)

(591.14, (Mooncake.NoTangent(), [84.0, 168.0, 252.0], 14.0, 1.0))

Friendly tangents

By default, Mooncake represents tangents using internal types such as Mooncake.Tangent for structs (see Mooncake.jl's Rule System). To return tangents in the same shape as the primal — for example a Symmetric tangent for a Symmetric matrix, or a NamedTuple mirroring a custom struct — set friendly_tangents=true in the Mooncake.Config:

config = Mooncake.Config(; friendly_tangents=true)
cache = Mooncake.prepare_gradient_cache(f, typical_x; config)
val, grad = Mooncake.value_and_gradient!!(cache, f, x)
(val, grad)

(14.0, (Mooncake.NoTangent(), [2.0, 4.0, 6.0]))

The performance impact of friendly_tangents=true should be negligible. If it is noticeable, something is likely wrong — please open an issue.

Beyond gradients

The same prepare-once, call-many-times pattern extends to the other derivatives below.

Forward mode

Mooncake.prepare_derivative_cache prepares a forward-mode cache. For a scalar-valued function it also backs value_and_gradient!!, computing the gradient in forward mode. See Interface for the chunked forward-mode controls.

Jacobians

Mooncake.value_and_jacobian!! computes the full Jacobian of a vector-valued function of a single dense vector input. It is not tied to a single mode: the cache it uses can come from either forward mode (Mooncake.prepare_derivative_cache) or reverse mode (Mooncake.prepare_pullback_cache). Either way it returns the primal output together with a dense Jacobian whose columns correspond to input coordinates:

h(x) = cos.(x) .* sin.(reverse(x))
cache = Mooncake.prepare_derivative_cache(h, x)
Mooncake.value_and_jacobian!!(cache, h, x)

([0.07624746575887673, -0.37840124765396416, -0.833049961066805], [-0.11874839215823475 -0.0 -0.5348952287053772; 0.0 -0.6536436208636119 0.0; -0.5348952287053772 -0.0 -0.11874839215823475])

Pullbacks

For outputs that are not scalars or vectors — for example a matrix or a custom struct — use Mooncake.prepare_pullback_cache and Mooncake.value_and_pullback!!, supplying a cotangent ȳ that matches the shape of f(x...).

Hessian-vector products

For a scalar-valued function with vector inputs, Mooncake.prepare_hvp_cache sets up forward-over-reverse AD:

q(x) = sum(x .* x)
cache = Mooncake.prepare_hvp_cache(q, x)
v = [1.0, 0.0, 0.0]
Mooncake.value_and_hvp!!(cache, q, v, x)

(14.0, [2.0, 4.0, 6.0], [2.0, 0.0, 0.0])

The returned tuple is (value, gradient, Hv).

Hessians

To materialise the full Hessian, use Mooncake.prepare_hessian_cache and Mooncake.value_gradient_and_hessian!!:

cache = Mooncake.prepare_hessian_cache(q, x)
Mooncake.value_gradient_and_hessian!!(cache, q, x)

(14.0, [2.0, 4.0, 6.0], [2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 2.0])

Terminology

Mooncake.jl is built around Fréchet derivatives and their adjoints, described in detail in Algorithmic Differentiation.

General cases:

• Fréchet derivative: In forward mode, Mooncake computes the Fréchet derivative D f[x], which maps tangent vectors to tangent vectors. It is implemented in Mooncake.value_and_derivative!!.

• Adjoint of the derivative (pullback): In reverse mode, Mooncake computes the adjoint D f[x]* of the Fréchet derivative, which maps cotangent vectors backwards through the computation. It is implemented in Mooncake.value_and_pullback!!.

Special cases (scalar input/output):

• Derivative: When the input is scalar, the Fréchet derivative f'(x) = D f[x](v) with v = 1 gives the ordinary derivative, handled as a special case of Mooncake.value_and_derivative!!.

• Gradient: When the output is scalar, the adjoint of the derivative applied to 1 gives the gradient ∇f, implemented in Mooncake.value_and_gradient!!.