feat(differentiation): GPU default ad_backend AutoZygote -> AutoMooncake

BREAKING.md

@@ -2,6 +2,18 @@
 
 This document outlines all breaking changes introduced in CTBase v0.18.0-beta compared to v0.17.4. Use this guide to migrate your code and understand the impact of these changes.
 
+## Non-breaking note (0.28.2-beta)
+
+- **`Differentiation`: GPU default `:ad_backend` changed from `AutoZygote()`
+  to `AutoMooncake()`** (roadmap-v4 §5 phase 4e, measured on an H200 — see
+  `CHANGELOG.md`). The option name, override mechanism, and `computed=true`
+  flag are all unchanged; only the *computed default value* on the opt-in
+  `GPU` parameter changes. **No breaking change**: any caller already
+  overriding `:ad_backend` explicitly — including to `AutoZygote()` — is
+  unaffected, and the change is invisible to every CPU (default) caller. No
+  new dependency (`ADTypes.AutoMooncake()` is a marker type from `ADTypes`,
+  already a hard dep, same as `AutoZygote()` before it).
+
 ## Non-breaking note (0.28.1-beta)
 
 - **`Differentiation`: `DifferentiationInterface` parameterized on the

CHANGELOG.md

@@ -5,6 +5,40 @@ All notable changes to CTBase will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.28.2-beta] - 2026-07-22
+
+### 🛠 Enhancements
+
+#### **Differentiation** — GPU default `ad_backend` swapped from `AutoZygote()` to `AutoMooncake()`
+
+- **`__ad_backend(::Type{GPU})` now returns `ADTypes.AutoMooncake()`**, not
+  `AutoZygote()`. Driven by roadmap-v4 §5 **phase 4e** (CTFlows PR #353), an
+  H200 probe measuring every reverse-mode DI candidate end to end:
+  `AutoMooncake` cleared the entire matrix, including the real `Flow(ocp)`
+  Hamiltonian call on a `CuArray` with a numerically correct result.
+  `AutoZygote` turned out to fail *unexpectedly* on a non-mutating
+  `hamiltonian_gradient` device call (`KernelException`) — using the exact
+  same Hamiltonian formula that succeeds through a different call shape
+  elsewhere — meaning its device behaviour is call-context dependent rather
+  than a stable pass/fail. `AutoZygote` remains selectable explicitly; it is
+  simply no longer the default.
+- **`describe(:di, registry)` now shows a GPU-specific message**: the
+  `:ad_backend` option's `description` is P-dependent (new
+  `_ad_backend_description(P)` in `differentiation_interface.jl`). The `GPU`
+  variant explains that `AutoForwardDiff()` does not work there
+  (scalar-indexes a `CuArray`) and that `AutoZygote()` is unreliable in some
+  device call contexts — the `CPU` variant is unchanged.
+- **No new dependency**: `ADTypes.AutoMooncake()` is a marker type from
+  `ADTypes` (already a CTBase hard dep), same pattern as `AutoZygote()`
+  before it — constructing it needs no `Mooncake` package loaded.
+
+### 🧪 Testing
+
+- `test/suite/differentiation/test_di_parameter.jl`: the GPU default
+  assertions now check `AutoMooncake()`; new assertions pin the GPU-specific
+  description text (mentions `AutoMooncake`/the ForwardDiff-on-GPU warning)
+  and confirm the CPU description is unaffected.
+
 ## [0.28.1-beta] - 2026-07-19
 
 ### 🛠 Enhancements

Project.toml

@@ -1,6 +1,6 @@
 name = "CTBase"
 uuid = "54762871-cc72-4466-b8e8-f6c8b58076cd"
-version = "0.28.1-beta"
+version = "0.28.2-beta"
 authors = ["Olivier Cots <olivier.cots@irit.fr>", "Jean-Baptiste Caillau <caillau@univ-cotedazur.fr>"]
 
 [deps]

src/Differentiation/default.jl

@@ -29,12 +29,16 @@ __ad_backend(::Type{Strategies.CPU}) = ADTypes.AutoForwardDiff()
 """
 $(TYPEDSIGNATURES)
 
-Default AD backend for GPU execution: `AutoZygote()`.
+Default AD backend for GPU execution: `AutoMooncake()`.
 
-`AutoZygote` is a GPU-capable AD backend (measured correct on `CuArray` inputs). The marker
-type comes from `ADTypes` (a hard dependency) and needs no Zygote loaded to construct; Zygote
-is required only when a gradient is actually evaluated.
+`AutoMooncake` was validated end-to-end on `CuArray` inputs, including through a mutating
+in-place right-hand side (the OCP Hamiltonian's `dynamics!` shape) — see
+[CTFlows PR #353](https://github.com/control-toolbox/CTFlows.jl/pull/353). `AutoZygote` was the
+prior default; it is GPU-capable for some call shapes but was found to fail unexpectedly on
+others (a non-mutating `hamiltonian_gradient` call), so it is no longer the default (still
+selectable explicitly). The marker type comes from `ADTypes` (a hard dependency) and needs no
+Mooncake loaded to construct; Mooncake is required only when a gradient is actually evaluated.
 
 See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref), [`CTBase.Strategies.GPU`](@ref).
 """
-__ad_backend(::Type{Strategies.GPU}) = ADTypes.AutoZygote()
+__ad_backend(::Type{Strategies.GPU}) = ADTypes.AutoMooncake()

src/Differentiation/differentiation_interface.jl

@@ -13,19 +13,20 @@ ecosystem.
 
 Parameterized on the execution device `P`:
 - `DifferentiationInterface{CPU}`: default backend `AutoForwardDiff()` (default device);
-- `DifferentiationInterface{GPU}`: default backend `AutoZygote()` (GPU-capable AD).
+- `DifferentiationInterface{GPU}`: default backend `AutoMooncake()` (validated end-to-end on
+  `CuArray`, including through a mutating in-place right-hand side).
 
 `DifferentiationInterface(...)` builds a `DifferentiationInterface{CPU}` — the device
 parameterization is fully backward compatible with existing call sites.
 
 # Arguments
 - `backend`: The DifferentiationInterface.jl backend to use. Defaults to the device default
-  (`AutoForwardDiff()` on CPU, `AutoZygote()` on GPU) from `ADTypes.jl` (a hard dependency).
+  (`AutoForwardDiff()` on CPU, `AutoMooncake()` on GPU) from `ADTypes.jl` (a hard dependency).
 - `kwargs...`: Additional options passed to `StrategyOptions`.
 
 # Notes
  - `ADTypes.jl` is a hard dependency, so the default backend markers are always available in core;
-   `AutoZygote()` is a marker type and needs no Zygote loaded to construct.
+   `AutoMooncake()` is a marker type and needs no Mooncake loaded to construct.
  - Gradient computation requires the `CTBaseDifferentiationInterface` extension.
  - Without the extension, the gradient methods throw `NotImplemented` with a helpful message.
 
@@ -65,7 +66,7 @@ $(TYPEDSIGNATURES)
 
 Construct a parameterized `DifferentiationInterface{P}` for the execution device `P`
 (`CPU` or `GPU`). The default `:ad_backend` is device-aware (`AutoForwardDiff()` on CPU,
-`AutoZygote()` on GPU) and can be overridden through the `:ad_backend` option.
+`AutoMooncake()` on GPU) and can be overridden through the `:ad_backend` option.
 
 # Arguments
 - `mode::Symbol=:strict`: Validation mode forwarded to `build_strategy_options`.
@@ -134,14 +135,30 @@ function Strategies.description(::Type{<:DifferentiationInterface})
     return "AD backend wrapping DifferentiationInterface.jl backends (e.g., AutoForwardDiff)."
 end
 
+"""
+Per-device `:ad_backend` description surfaced by `Strategies.describe(:di, registry)` (each
+parameter renders its own `computed options for P` section, see
+`CTBase.Strategies._describe_multi_param_metadata`). The GPU variant is deliberately more
+detailed: it is the one place a user is warned, before hitting a device, that `AutoForwardDiff`
+does not work there and that `AutoZygote` is unreliable in some call contexts.
+"""
+_ad_backend_description(::Type{Strategies.CPU}) = "DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU)."
+function _ad_backend_description(::Type{Strategies.GPU})
+    return "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)."
+end
+
 """
 $(TYPEDSIGNATURES)
 
 Return metadata defining `DifferentiationInterface{P}` options and their specifications.
 
-The `:ad_backend` default is device-aware: `AutoForwardDiff()` on `CPU`, `AutoZygote()` on `GPU`
-(via [`CTBase.Differentiation.__ad_backend`](@ref)). The bare `metadata(DifferentiationInterface)`
-delegates here through `DifferentiationInterface{CPU}`.
+The `:ad_backend` default is device-aware: `AutoForwardDiff()` on `CPU`, `AutoMooncake()` on
+`GPU` (via [`CTBase.Differentiation.__ad_backend`](@ref)). The bare
+`metadata(DifferentiationInterface)` delegates here through `DifferentiationInterface{CPU}`.
 """
 function Strategies.metadata(
     ::Type{<:DifferentiationInterface{P}}
@@ -151,8 +168,8 @@ function Strategies.metadata(
             name=:ad_backend,
             type=ADTypes.AbstractADType,
             default=__ad_backend(P),
-            computed=true,  # Default is computed from parameter P (CPU→ForwardDiff, GPU→Zygote)
-            description="DifferentiationInterface.jl backend (e.g. AutoForwardDiff() on CPU, AutoZygote() on GPU).",
+            computed=true,  # Default is computed from parameter P (CPU→ForwardDiff, GPU→Mooncake)
+            description=_ad_backend_description(P),
             aliases=(:backend, :ad),
         ),
     )

test/suite/differentiation/test_di_parameter.jl

@@ -45,14 +45,20 @@ function test_di_parameter()
             md_gpu = Strategies.metadata(DI{Strategies.GPU})
 
             Test.@test md_cpu[:ad_backend].default === ADTypes.AutoForwardDiff()
-            Test.@test md_gpu[:ad_backend].default === ADTypes.AutoZygote()
+            Test.@test md_gpu[:ad_backend].default === ADTypes.AutoMooncake()
 
             # The default is computed from P, so it is flagged `computed` (for `describe`/provenance).
             Test.@test md_cpu[:ad_backend].computed == true
             Test.@test md_gpu[:ad_backend].computed == true
 
             # Bare `metadata(DI)` delegates to `metadata(DI{CPU})` (back-compat).
             Test.@test md_bare[:ad_backend].default === ADTypes.AutoForwardDiff()
+
+            # GPU description is P-specific (phase 4e): warns that ForwardDiff doesn't work on
+            # GPU and that Zygote is unreliable there, pointing users at the Mooncake default.
+            Test.@test occursin("AutoMooncake", md_gpu[:ad_backend].description)
+            Test.@test occursin("AutoForwardDiff() does not work on GPU", md_gpu[:ad_backend].description)
+            Test.@test !occursin("AutoMooncake", md_cpu[:ad_backend].description)
         end
 
         # ====================================================================
@@ -71,8 +77,8 @@ function test_di_parameter()
             di_gpu = DI{Strategies.GPU}()
             Test.@test di_gpu isa DI{Strategies.GPU}
             Test.@test Strategies.parameter(typeof(di_gpu)) == Strategies.GPU
-            # GPU default backend is AutoZygote (resolved from the per-device metadata).
-            Test.@test Differentiation.ad_backend(di_gpu) === ADTypes.AutoZygote()
+            # GPU default backend is AutoMooncake (resolved from the per-device metadata).
+            Test.@test Differentiation.ad_backend(di_gpu) === ADTypes.AutoMooncake()
 
             # Explicit override still wins on either device.
             di_gpu_fd = DI{Strategies.GPU}(; backend=ADTypes.AutoForwardDiff())