Enzyme: MixedReturnException with PolyAlgorithm + NamedTuple broadcast in reverse rule

[ChrisRackauckas-Claude]

Two remaining Enzyme issues after #869 / #870

1. MixedReturnException with default PolyAlgorithm

The MixedReturnException persists with Enzyme 0.13.134 and NonlinearSolveBase 2.16.0. It occurs even without SciMLSensitivity loaded — it's a compile-time type check failure.

Root cause: NonlinearSolution has Any as the 7th type parameter (the original/cache field) when using PolyAlgorithm, because different sub-algorithms have different cache types. Enzyme can't determine activity for Any-typed fields and rejects the custom rule's return type.

With explicit NewtonRaphson(), the solution has all concrete type parameters and works fine.

using NonlinearSolve, Enzyme
f(u, p) = u .^ 2 .- p

# Works (concrete types in solution)
loss1(p) = sum(solve(NonlinearProblem(f, [1.0], p), NewtonRaphson()).u)
Enzyme.gradient(Enzyme.Reverse, loss1, [2.0])  # ✓

# Fails (Any in solution type from PolyAlgorithm)
loss2(p) = sum(solve(NonlinearProblem(f, [1.0], p)).u)
Enzyme.gradient(Enzyme.Reverse, loss2, [2.0])  # MixedReturnException

Fix needed: Make the original/cache field type concrete for PolyAlgorithm solutions, or find a way to tell Enzyme to ignore inactive fields.

2. NamedTuple broadcasting error in reverse rule with MTKParameters

When p is MTKParameters, the reverse rule at line 62 (ptr.dval .+= darg) fails because darg (from SciMLSensitivity's steadystatebackpass) is a NamedTuple:

ArgumentError: broadcasting over dictionaries and `NamedTuple`s is reserved

The tangent from steadystatebackpass uses a Zygote pullback that produces a NamedTuple{(:tunable, :initials, ...)} when the parameter type is MTKParameters. The reverse rule needs to handle this case.

Proposed fix: Use SciMLStructures to accumulate the tangent, or convert the NamedTuple to the appropriate parameter type before accumulating.

Versions

• Julia 1.10.11
• Enzyme 0.13.134
• NonlinearSolveBase 2.16.0

Related

• SciMLSensitivity.jl#1358
• Enzyme.jl#3020 (EnzymeMutabilityException with remake closure)

🤖 Generated with Claude Code

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

[ChrisRackauckas-Claude]

Update: MixedReturnException root cause identified

The MixedReturnException is not caused by the Any type parameter in build_solution_less_specialize. Setting original=nothing (concrete Nothing type) doesn't fix it.

Actual root cause

It's an Enzyme limitation with @generated functions that have many code branches returning different types. Testing shows:

PolyAlgorithm config	Enzyme
1-6 identical NewtonRaphson()	✓ works
2 mixed types (NR + TrustRegion)	✓ works
3+ mixed types (NR + TR + LM)	✗ MixedReturnException

The @generated solve! function produces N branches, each with a different sub-algorithm cache type. Even though all branches call build_solution_less_specialize which returns the same outer type (with alg=PolyAlgorithm), Enzyme's static analysis of the generated function body gets confused by 3+ divergent code paths.

This is an Enzyme bug with @generated function analysis. The NonlinearSolve side can't fix it without restructuring the PolyAlgorithm to avoid @generated functions, or reducing the default algorithm count.

🤖 Generated with Claude Code

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

[ChrisRackauckas-Claude]

Correction: MixedReturnException is a runtime error in Enzyme's fallback dispatch

Further investigation shows:

1. autodiff_thunk compiles successfully — the error is NOT at compile time
2. The Enzyme rule on solve_up matches at compile time for explicit algorithms (e.g. NewtonRaphson()) but NOT for the default PolyAlgorithm case
3. For PolyAlgorithm, Enzyme falls through to runtime_generic_augfwd (its runtime fallback), which then throws MixedReturnException

The original=nothing change doesn't help because the issue isn't the return type — it's that Enzyme can't resolve the rule at compile time when no algorithm is passed, and the runtime fallback doesn't handle the rule correctly.

Root cause: Enzyme's compile-time rule matching fails for solve_up(prob, nothing, u0, p) (no algorithm arg) even though the rule signature matches. This is specific to Enzyme's runtime_generic_augfwd codepath.

This needs an Enzyme fix, not a NonlinearSolve fix. The NonlinearSolve rule is correctly defined and works when Enzyme can resolve it at compile time (explicit algorithm case).

🤖 Generated with Claude Code

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