Historical AD validation

This archive preserves the rewrite's dated correctness checks, measurements, and implementation investigations. References to “current” or “previous” describe the checkout used for that snapshot, not necessarily today's implementation. Earlier passing results are not a compatibility guarantee. Use the current AD guide for maintained regression coverage and recommended configurations.

This page records the configurations tested with ShootingLayer, their known limitations, and the design decisions behind differentiable trajectory assembly. Compatibility depends on both the outer differentiation backend and the sensealg used by the underlying SciML problem.

Maintained regression groups

The Differentiation test group exercises full ODE solves with ForwardDiff, ReverseDiff, and Zygote against an analytic reference at two inputs. The loss includes saved states and continuity residuals; inputs include shared parameters, a free global initial value, interior shooting values, and numerical values in Lux state. ForwardDiff and Zygote cover serial and threaded layers. ReverseDiff uses serial execution with FullSpecialize and SensitivityADPassThrough().

The Distributed group compares two-worker ForwardDiff gradients with serial execution and finite differences. It uses a named callable objective loaded on every worker; anonymous objective closures can produce mismatched ForwardDiff tag types after serialization. Reverse-mode differentiation of the ensemble map still records blocks serially. These checks do not establish compiled ReverseDiff tapes, parallel reverse execution, Mooncake, or Enzyme support.

Run either group from the repository root:

GROUP=Differentiation julia --project=. -e 'using Pkg; Pkg.test()'
GROUP=Distributed julia --project=. -e 'using Pkg; Pkg.test()'

The dated measurements below remain historical snapshots rather than a fresh benchmark of the current checkout.

Rewrite validation: 28 October 2026

The root test groups passed 1,630 checks, with the existing implicit-import check marked broken. Separate serial Lotka benchmarks checked controls alone and joint state/parameter shooting values at two inputs, including preparation reuse and unchanged inputs. Maximum gradient errors against central differences were below 3.5e-10.

ReverseDiff pass-through, Mooncake, and Enzyme use FullSpecialize and SensitivityADPassThrough(). Mooncake passes without caller-side initial-state promotion. Enzyme forward still hits LuxCore's layer-annotation restriction. Default ReverseDiff returns a tracked state array rather than a full solution; the next segment fails at internal_solve. Pass-through retains the solution interface. FiniteDiff is a numerical reference.

The reproducible scripts and full tables, including first-use costs and forward allocations, are in benchmark/README.md and benchmark/results/2026-10-28.md. These runs use the regular layer(nothing, ps, st) entry point. Older measurements below used a direct (problem, tstops) call and different first-call timing boundaries, so they are historical snapshots rather than matched comparisons. The benchmark results do not establish parallel reverse execution, compiled ReverseDiff tapes, or broader Mooncake coverage beyond the two serial cases.

Scalar-packing validation: 26 October 2026

State and parameter packing now uses vector literals to retain ordinary arrays of tracked scalars. Initial-condition resets return the same promoted scalar type as pass-through values, keeping preallocated solution buffers compatible across segments. The first-solve remake and paired NamedTuple worker-state traversal are also in place. The full Core suite passes 1,063 checks, including full-solver ReverseDiff regressions without caller-side initial-state promotion. QA retains its existing broken implicit-import check.

Full-solver checks use an in-place Lotka ODE on [0, 4], eight original control intervals, and abstol = reltol = 1e-10. The controls case has eight inputs. The joint case lifts all three states at 1.25 and 2.75, splitting crossing control intervals and giving 16 inputs. The loss is the terminal third state plus 0.1 * sum(abs2, shooting_constraints(trajectory)). Two inputs and reused gradient preparations are compared against central differences, with discrepancies below 3.5e-10; numerical inputs remain unchanged.

Enzyme's failure existed before scalar packing changed. These ReverseDiff results cover serial execution with uncompiled tapes. The default ReverseDiff sensitivity adapter still returns tracked state arrays instead of the full solution interface required by trajectory assembly; it is a different route from pass-through. Compiled tapes and threaded ReverseDiff have not been established by these checks.

Prepared in-place gradient measurements use one Julia and one BLAS thread:

ForwardDiff and Zygote values are medians of five alternating before/after rounds of 100 prepared gradient! calls. Other rows use medians of up to 20 warmed calls. Allocations are per call in decimal MB. Strategies differ between backends, so these exploratory measurements are not a universal backend ranking.

The matched before/after gradient comparison isolates scalar packing and the promoted reset return type; the baseline already contains both inference fixes:

First preparation plus gradient measurements, excluding package loading:

These include preparation and JIT; the joint case reuses compilation from controls. They are not pure compilation timings.

Whole-layer forward measurements use the 120-interval Lotka problem on [0, 12], tolerances 1e-9, four Julia threads, and no optional AD backend explicitly loaded. Multiple shooting lifts all states at 3.05, 6.05, and 9.05. Values below are medians of five alternating before/after rounds of 1,000 warmed calls, including trajectory construction. The baseline already includes the first-solve remake:

Bytes and allocation counts are identical before and after. Serial timing changes are small. Threaded round medians ranged approximately 128–166 µs, so the median difference does not establish a systematic regression. Fresh-process first calls measured 3.542 → 3.508 s for single shooting, 0.513 → 0.549 s for serial multiple shooting, and 1.261 → 1.273 s for threaded multiple shooting. Package loading and setup are excluded; later workloads reuse compilation. These are single observations.

These checks used Julia 1.12.7, SciMLBase 3.58.0, DiffEqBase 7.21.4, OrdinaryDiffEqTsit5 2.1.5, SciMLSensitivity 7.119.13, DifferentiationInterface 0.7.21, ForwardDiff 1.4.6, ReverseDiff 1.18.6, Zygote 0.7.14, Mooncake 0.5.64, and Enzyme 0.13.216.

Historical validation snapshot: 22 October 2026

The following results were obtained on 22 October 2026 with the full layer returning (trajectory, state), including trajectory.residuals. They describe the tested configurations, not unconditional backend support.

FiniteDiff is numerical differentiation, rather than automatic differentiation. "Not tested" does not indicate a failure. Compiled ReverseDiff tapes, Enzyme reverse mode, Mooncake forward mode, Tracker, and other configurations are not established by this snapshot.

All nine full-solver checks passed. They used an eight-interval Lotka ODE with Tsit5() over [0, 12], with abstol = reltol = 1e-11. The shooting-state cases also varied physical parameters and initial/interior shooting states. The loss combined the final third state with 0.3 * sum(r -> sum(abs2, r), trajectory.residuals).

Gradients were compared with central differences at two distinct inputs, using steps 1e-5 and 5e-6 to check reference stability. The second input reused the DifferentiationInterface preparation and exercised gradient!. Inputs were checked for unintended mutation. These were correctness checks, not gradient runtime benchmarks.

The snapshot predates the subsequent change to extract terminal states from tracked matrices. After that change, the Core suite still passed all 183 assertions, but the default ReverseDiff sensitivity path still failed on sol.alg. Rerun the backend matrix after further changes; the historical passing matrix does not replace regression checks on a new implementation.

The per-variable InitialConditions implementation has focused regression tests in test/core/initial_conditions_ad.jl for ForwardDiff, ReverseDiff, and Zygote. These check stored values, pass-through state values, residuals, and numeric values held in Lux state. A full-solver ForwardDiff regression covers the first lifted value occurring after the simulation starts: pass-through must retain the numeric type required by later lifted values so the preallocated buffers remain compatible. The other full-solver configurations in the historical matrix have not been rerun for this representation.

Configuring the solver and backend

Full specialization and pass-through

The ReverseDiff, Mooncake, and Enzyme configurations above used the following problem construction for an in-place right-hand side rhs!:

using SciMLBase, SciMLSensitivity

f = ODEFunction{true, SciMLBase.FullSpecialize}(rhs!)
prob = ODEProblem(f, u0, tspan, p;
    abstol = 1e-11, reltol = 1e-11,
    sensealg = SciMLSensitivity.SensitivityADPassThrough())

Pass-through bypasses SciML's sensitivity overrides and lets the outer backend differentiate the numerical solver directly. It retains the solution object and its saved metadata. It is a different differentiation strategy from using a SciML adjoint algorithm, and the two strategies need not have the same memory or runtime behavior. See the SciML sensitivity interface.

ForwardDiff and FiniteDiff

The tested DifferentiationInterface backends were DI.AutoForwardDiff() and DI.AutoFiniteDiff(; fdtype = Val(:central)), where import DifferentiationInterface as DI.

ForwardDiff passed both control-only and shooting-state cases with the default problem configuration. The FiniteDiff backend check covered controls; central differences were also used as the independent reference for the other cases.

ReverseDiff

The passing configuration used DI.AutoReverseDiff() without a compiled tape, together with full specialization and pass-through sensitivity.

In the historical control-only check, fixed initial states were promoted inside the differentiated loss using the value-preserving expression [u + 0 * z[1] for u in prob.u0], where z is the active optimization vector. The problem was remade with those initial states before applying the layer. This makes the initial state compatible with the tracked arithmetic introduced by the controls. The shooting-state case did not need this separate promotion. The current scalar-packing implementation passes both cases without caller-side promotion.

Pass-through retains the full solution interface, but current tracked-array storage compatibility must also be checked; see the current validation above. For the relevant upstream conversion, see the DiffEqBase 7.21.4 ReverseDiff extension.

Zygote

DI.AutoZygote() passed both control-only and shooting-state cases with the default problem configuration.

The original trajectory constructor collided with a LuxCore constructor rule that dispatched broadly on UnionAll. The layer now calls a named builder, and the public constructor has a constructor-specific configured ChainRules rule. Rules for parametric constructors should dispatch on an appropriate Type{...} signature, rather than typeof(Constructor). See the ChainRules constructor guidance.

Trajectory gathering uses a dedicated pullback instead of relying on Zygote.Buffer. The existing segment-solver buffers still use the Zygote extension when it is loaded. Zygote's limitations documentation describes mutation restrictions and cautions about Buffer.

Mooncake and Enzyme

These cases computed the segment times before differentiation and reused them while changing numerical parameters on the same grid. The benchmark used these internal entry points:

segment_times = Corleone.get_timestops(layer, ps, st)
trajectory, next_st = layer((layer.problem, segment_times), ps, st)

This is distinct from preparing a gradient: all backend checks used DI.prepare_gradient, while these two also prepared the segment schedule. The current layer also caches segment times in st.timestops during setup; ordinary forward calls reuse that cache. After changing schedules, refresh it with st = merge(st, (; timestops = Corleone.get_timestops(layer, nothing, st))). Gradient preparations may also need rebuilding when dimensions, types, or traced control flow change.

Mooncake used DI.AutoMooncake() and passed the shooting-state case. An initial 180-second worker limit was too short and was inconclusive. The full-trajectory follow-up passed in approximately 290 seconds with a 900-second limit, including package loading, compilation, and validation. These worker times are not warmed gradient timings. In a separate run without trajectory construction, preparation alone took approximately 158 seconds.

Enzyme used the following backend and passed the control-only case:

using Enzyme
import DifferentiationInterface as DI

backend = DI.AutoEnzyme(;
    mode = Enzyme.set_runtime_activity(Enzyme.Forward),
    function_annotation = Enzyme.Const)

The original trajectory construction raised an IllegalTypeAnalysisException. Resolving the symbolic container once at layer construction and retaining its concrete type removed that failure in the tested configuration. Runtime activity was already enabled; enabling it alone was not the fix. See the Enzyme FAQ for activity and memory-layout requirements.

Trajectory differentiation contract

The trajectory retains layer, ps, and st by reference. Saved state arrays are also shared with the segment solutions. Parameter histories are evaluated through layer.parameter and the retained ps/st; no separate parameter-history struct or detached copy is stored. Mutating these objects changes later parameter queries, so they must not be mutated between a forward evaluation and its pullback.

Assembly separates fixed-grid layout metadata from numerical state gathering. The gathering pullback accumulates contributions into source samples, including when an index occurs more than once. State values, residuals, and parameter values remain differentiable. Symbolic metadata and the discrete layout are excluded from the ChainRules differentiation path.

Validation covers numerical gradients on fixed time grids. It does not establish gradients with respect to event times, control-knot positions, or changes in grid topology. The full-solver matrix covers ODEs; it does not establish DAE, SDE, GPU, or higher-order derivative support.

Preserving metadata for array-returning solve rules

Actual saved times cannot generally be reconstructed from the outer control grid or segment boundaries. saveat, callbacks, and other saving settings can affect the samples saved by the solver. Internal tstops affect integration and need not each appear as a saved sample. States and timestamps must correspond to the same actual solve. DAE derivative values may need preserving as well.

If support for the default ReverseDiff sensitivity path becomes necessary, the adapter belongs at the segment-solve boundary. It should perform one primal solve, preserve that solve's actual metadata before array conversion, expose differentiable numerical values, and connect those values to the solve pullback. Metadata must remain consistent on subsequent evaluations and tape replay. An extra solve solely to recover timestamps is not the intended design.

This adapter is deferred until a workload requires the array-returning adjoint path. Prefer documenting the working configuration over maintaining a custom solver differentiation rule without that requirement.

Forward-pass measurements

The AD refactor was also compared with the preceding trajectory assembly on a 120-interval Lotka problem. Five rounds of 5,000 samples used one Julia thread; these are warmed forward timings with no gradient evaluation.

The previous assembly was reconstructed on the same current solver. In a separate assembly-only comparison with symbolic metadata cached in both versions, the old loop took 1.72 µs and the new layout/gather took 1.68 µs. The new version allocated 6,768 bytes in 21 allocations, versus 6,048 bytes in 25 allocations for the old loop. Thus the extra layout storage cost 720 bytes, with essentially equal runtime. Allocated bytes are not peak or retained memory.

These measurements showed no full-forward regression; they do not rank backend gradient performance. Loading Zygote changes the existing solver-buffer path, so comparisons should record which optional backends are loaded.

Regression checks and future validation

Run the Core test group from the repository root:

GROUP=Core julia --project -e 'using Pkg; Pkg.test()'

test/core/trajectory_ad.jl checks analytical gradients through the public constructor and named builder with ReverseDiff and Zygote. It covers state, residual, parameter-history, and combined losses at two inputs, continuous/reset boundaries, and repeated-source-index accumulation.

The full backend matrix above was exploratory and is not part of the committed Core suite. To repeat or extend it, use isolated processes per backend, reconstruct the ODE/loss described above, and record the exact solver and backend configuration. Compare central differences at multiple inputs, verify reference stability, reuse gradient preparations at a second input, and check that inputs are unchanged. For new solve-boundary adapters, include interior saveat samples and callbacks and verify state/time correspondence. Record failures separately from timeouts and untested combinations.

Versions used in the historical validation snapshot

Shared-grid scalar-buffer validation

On 28 October 2026, Julia 1.12.7 checks exercised the shared-lookup implementation with two tunable controls and one fixed control on the same grid. A second case also grouped three shooting states and inserted parameter intervals at full-state resets. Each backend was checked at two inputs against central differences at two step sizes; prepared gradients were reused and input vectors remained unchanged.

Successful serial backend gradients differed from finite differences by less than 5e-10 in these cases. FiniteDiff also passed as a numerical reference. This does not extend Mooncake or ReverseDiff validation to parallel execution, or establish Enzyme reverse support.

The maintained Differentiation group now has 70 checks: 30 existing analytic checks and 40 additional checks for grouped parameter/reset containers, shared parameters, and numerical values in Lux state. The new analytic fixture varies both a fixed reset value and a fixed parameter through the input vector; their gradients remain part of the solve. The Distributed group has 25 checks, including this shared-grid fixture on two workers. Both fixtures use two inputs and verify that calls do not mutate their input vectors.

Set CORLEONE_AD_SHARED_GRIDS=true when running benchmark/shooting_ad.jl to repeat the shared-control cases. The correctness audit used one benchmark sample per case and ran backend processes concurrently; its timings are not a performance comparison. The compatibility statements above apply to the tested configurations.