MixedReturnException in runtime_generic_augfwd with custom rule on function returning complex type

[ChrisRackauckas-Claude]

Description

Enzyme.gradient through a function with a custom augmented_primal/reverse rule throws MixedReturnException at runtime when the return type involves a @generated PolyAlgorithm with 3+ mixed sub-algorithm types. Works fine with 1-2 mixed types or any number of identical types.

The error is thrown from runtime_generic_augfwd, not at compile time — autodiff_thunk compiles successfully and has_rrule_from_sig returns true. Julia's Base.return_types infers a single concrete type for all cases.

MWE

using NonlinearSolve, Enzyme

f(u, p) = u .^ 2 .- p

# ✓ Works: 2 mixed algorithm types
poly2 = NonlinearSolvePolyAlgorithm([NewtonRaphson(), TrustRegion()])
Enzyme.gradient(Enzyme.Reverse,
    p -> sum(solve(NonlinearProblem(f, [1.0], p), poly2).u), [2.0])
# => ([0.353...],)

# ✗ Fails: 3 mixed algorithm types
poly3 = NonlinearSolvePolyAlgorithm([NewtonRaphson(), TrustRegion(), LevenbergMarquardt()])
Enzyme.gradient(Enzyme.Reverse,
    p -> sum(solve(NonlinearProblem(f, [1.0], p), poly3).u), [2.0])
# => MixedReturnException

# ✓ Works: 6 identical types (same count, no type diversity)
poly6 = NonlinearSolvePolyAlgorithm([NewtonRaphson() for _ in 1:6])
Enzyme.gradient(Enzyme.Reverse,
    p -> sum(solve(NonlinearProblem(f, [1.0], p), poly6).u), [2.0])
# => ([0.353...],)

Key observations

• autodiff_thunk compiles: the error is runtime, not compile-time
• has_rrule_from_sig returns true for all cases
• Base.return_types returns a single concrete type for all cases
• Without any custom rule, Enzyme differentiates through all cases natively (including 6 mixed types)
• The custom rule in NonlinearSolveBaseEnzymeExt (on solve_up) handles Union{Type{Duplicated{RT}}, Type{MixedDuplicated{RT}}}
• Could not reproduce with pure Julia types — the trigger seems specific to the complexity of the NonlinearSolve dispatch chain inside the rule body

Versions

• Julia 1.10.11
• Enzyme 0.13.134
• NonlinearSolve 4.16.0
• NonlinearSolveBase 2.16.0

Related

• NonlinearSolve.jl#878
• SciMLSensitivity.jl#1358

🤖 Generated with Claude Code

Co-Authored-By: Chris Rackauckas accounts@chrisrackauckas.com

[wsmoses]

@ChrisRackauckas-Claude this must be possible to reproduce without a dependency on NonlinearSolve. worst case inline everything then simplify it

[wsmoses]

in any case the error should say: "MixedReturnException: Custom Rule for method returns type $(RT), which has mixed internal activity types. This is not presently supported.\n"

why are you writing a rule for a mixed return type?

[ChrisRackauckas]

It seemed to be due to an Any.

[wsmoses]

in which case is all resolved?

[ChrisRackauckas]

I believe this one is resolved, yes. The other one is still open.