The goal of the Mooncake.jl project is to produce an automatic differentiation (AD) package written entirely in Julia that improves on ForwardDiff.jl, ReverseDiff.jl, and Zygote.jl in several ways. Applying AD to Julia's type-inferred, optimised intermediate representation helps produce efficient derivative code. Support for mutation allows Mooncake to differentiate a wide range of numerical Julia code without hand-written rules. See the documentation for details.

Performance

On one system, Flux benchmarks found Mooncake gradient evaluations after warm-up were 2.03 times as fast as Zygote on CPU across 19 models, with comparable GPU performance. First evaluations were substantially slower. See also the DynamicPPL benchmarks. Performance varies by workload.

Getting started

Check whether Mooncake's support policy covers your Julia version.

Mooncake uses reusable caches for repeated gradient and Hessian evaluations:

import Mooncake as MC

f(x) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2  # Rosenbrock
x = [1.2, 1.2]

# Reverse mode
grad_cache = MC.prepare_gradient_cache(f, x);
value, (_, gradient) = MC.value_and_gradient!!(grad_cache, f, x)

# Forward mode
fwd_cache = MC.prepare_derivative_cache(f, x);
value_fwd, (_, gradient_fwd) = MC.value_and_gradient!!(fwd_cache, f, x)

# Hessian
hess_cache = MC.prepare_hessian_cache(f, x);
value, gradient, hessian = MC.value_gradient_and_hessian!!(hess_cache, f, x)

Cache preparation takes some time, but calls that reuse the cache are fast. Each cache is tied to its inputs' types and sizes; passing a differently sized input raises an error. See the tutorial for a walkthrough and the interface for details.

Contributing

In the spirit of long-lived projects such as R and TeX, we take a conservative approach to development: correctness, stability, and tightly scoped fixes take precedence over expanding the package's scope.

Contributions are most welcome when they address reproducible defects, such as incorrect results, unexpected failures, or behaviour inconsistent with the documented scope.

If you wish to extend Mooncake's coverage, we encourage you to organise independently. See the support policy for guidance and examples.

Licensing

Mooncake is licensed under the MIT License. Its required and optional dependencies are licensed separately and may impose additional terms on redistributed applications or binaries. See Project.toml for the dependency list.