Skip to content

Forward-over-reverse: Not yet implemented, mixed activity for jl_new_struct in runtime_generic_fwd when outer forward differentiates a dynamically-dispatched inner reverse rule (ODE adjoint) #3277

Description

@ChrisRackauckas-Claude

Summary

Forward-over-reverse (nested Enzyme: outer Forward over an inner Reverse) through an ODE solve with a continuous-adjoint sensealg fails at runtime with:

EnzymeRuntimeException: Enzyme execution failed.
Enzyme: Not yet implemented, mixed activity for jl_new_struct
  constants=Bool[0, 0, 1, 0, 0, 1]
  ... @jl_f_tuple( ... )  Tuple{Bool, Any, Any}

The outer forward pass, running through its generic runtime (runtime_generic_fwd), can't construct a mixed active/const aggregate (jl_new_struct / jl_f_tuple) while differentiating the dynamically-dispatched inner-reverse rule. The inner reverse itself succeeds (the reverse gradient is correct); only the outer-forward-over-it fails.

This is the same error family as #3246 (mixed activity for jl_new_struct, fixed by #3251/#3254 for ODEProblem construction), but at a new site — a Tuple{Bool, Any, Any} bundling the reverse-adjoint's seed closure, the sensealg, and the init algorithm.

MWE (SciMLSensitivity + OrdinaryDiffEq + Enzyme; no DI, manual nesting)

import SciMLSensitivity as SMS
import OrdinaryDiffEq as ODE
import Enzyme

f(u, p, t) = -p[1] .* u
u0    = [1.0]
tspan = (0.0, 1.0)
ts    = range(0.0, 1.0, length = 5)

function loss(p)
    prob = ODE.ODEProblem(f, u0, tspan, p)
    sol  = ODE.solve(prob, ODE.Tsit5(); saveat = ts,
                     sensealg = SMS.GaussAdjoint(autojacvec = SMS.EnzymeVJP()))
    sum(abs2, Array(sol))
end

const RA_R = Enzyme.set_runtime_activity(Enzyme.Reverse)
const RA_F = Enzyme.set_runtime_activity(Enzyme.Forward)

@noinline function grad(p)                 # inner reverse gradient — works
    dp = zero(p)
    Enzyme.autodiff(RA_R, Enzyme.Const(loss), Enzyme.Active, Enzyme.Duplicated(p, dp))
    dp
end

function hvp(p, v)                          # outer forward over the reverse gradient
    Enzyme.autodiff(RA_F, Enzyme.Const(p_ -> sum(grad(p_))), Enzyme.Duplicated(p, v))
end

grad([0.5])          # [-2.440239558871761]   (inner reverse is fine)
hvp([0.5], [1.0])    # EnzymeRuntimeException: Not yet implemented, mixed activity for jl_new_struct

Reproduces identically with InterpolatingAdjoint, GaussAdjoint, QuadratureAdjoint, and BacksolveAdjoint (all with EnzymeVJP()), so it's not specific to one sensealg.

Where it fails (stack)

The outer forward's generic runtime is differentiating DiffEqBase's Enzyme reverse rule for solve_up:

 [1] adjoint_sensitivity_backpass   @ SciMLSensitivity/src/concrete_solve.jl:945   (calls adjoint_sensitivities)
 [2] #reverse#4                     @ DiffEqBase/ext/DiffEqBaseEnzymeExt.jl:113
 ...
 [9] runtime_generic_fwd(...)       @ Enzyme/src/rules/jitrules.jl:377
       differentiating ::typeof(EnzymeRules.reverse), returning ::Type{Duplicated{Any}}
 [10] reverse                       @ DiffEqBaseEnzymeExt.jl:99
 [11] #solve#38                     @ OrdinaryDiffEqCore/.../solve.jl:603
 ...
 [26] runtime_generic_fwd(...)      @ jitrules.jl:377

The jl_new_struct being built is a Tuple{Bool, Any, Any} whose elements (from the Enzyme metadata) are: the discrete-adjoint seed closure (SciMLSensitivity var"#df_iip#340" / #df_oop#341, which capture the active ODESolution), the sensealg (GaussAdjoint{…EnzymeVJP}), and BrownFullBasicInit{Float64,Nothing} — a mix of active (MUT_REF/BITS_REF) and const elements (constants=Bool[0,0,1,0,0,1]). Note the inner reverse rule is dynamically dispatched and returns Duplicated{Any}, so the outer forward goes through runtime_generic_fwd rather than a typed forward path.

Reduction attempts (why no dependency-free MWE)

I tried to reduce this to an Enzyme-only MWE without success (~10 shapes): building mixed const/active Tuple{Bool,Any,Any} of boxed closures; closures capturing active mutable Vectors; mixed MUT_REF/BITS_REF element kinds; and forcing the generic runtime via Base.inferencebarrier / ::Any-returning inner functions / dynamic dispatch — under both plain forward mode and forward-over-reverse. All return the correct derivative — the typed forward path handles mixed activity fine, and I couldn't get the synthetic cases onto the runtime_generic_fwd jl_new_struct path the way the real dynamically-dispatched EnzymeRules.reverse-returning-Duplicated{Any} call does. Filing with the SciML-stack MWE per maintainer preference; happy to keep reducing if a dependency-free repro is required.

Versions

  • Enzyme v0.13.173
  • SciMLSensitivity v7.112.2, OrdinaryDiffEq v7.1.1 / OrdinaryDiffEqCore v4.5.0, DiffEqBase v7.6.0, SciMLBase v3.30.1
  • Julia 1.11.9

Context

This is the next blocker for SecondOrder(AutoEnzyme(Forward), AutoEnzyme(Reverse)) forward-over-reverse Hessians/HVPs through SciMLSensitivity + OrdinaryDiffEq (SciML/SciMLSensitivity.jl#1427), after the _ode_init Float32 type-analysis blocker #3276. (I used a Float64 problem here specifically to get past #3276 and reach this one; the Float32 docs workload hits #3276 first.) Prior cleared blockers in this chain: #3135#3137, #3213, #3246#3251, #3253#3254, #3258#3261.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions