SciML-free: runtime activity error through FunctionWrappersWrapper with Lux ComponentArray
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
## Summary
A reverse-mode Enzyme call through `FunctionWrappersWrappers.FunctionWrappersWrapper`, with a Lux RHS and `ComponentArray` parameters, raises `EnzymeRuntimeActivityError` in static-activity mode.
This is a **SciML-free partial reduction** found while reducing #3260. It isolates an activity mismatch at the function-wrapper boundary. It does **not** reproduce #3260's deeper `Array has no field layer_1` / Julia 1.12 GC-corruption failure: enabling runtime activity makes this smaller case pass.
## Environment
- Julia 1.10.11
- Enzyme current main: `23ff63088cd6532ea0d58bcdc279573bec6e2731` (package version 0.13.182)
- Enzyme_jll 0.0.282+0
- Lux 1.31.4
- LuxLib 1.15.9
- ComponentArrays 0.15.42
- FunctionWrappersWrappers 1.10.0
## Reproducer
```julia
using ComponentArrays
using Enzyme
using FunctionWrappersWrappers
using Lux
using Random
const model = Chain(Dense(4, 50, tanh), Dense(50, 2))
const parameters, state = Lux.setup(Xoshiro(0), model)
const p = ComponentArray{Float32}(parameters)
function rhs!(du, u, p, t)
du[1:2] .= first(model(u, p, state))
du[3:4] .= 0.0f0
return nothing
end
struct Void{F}
f::F
end
function (f::Void)(args...)
f.f(args...)
return nothing
end
argtypes = (Tuple{Vector{Float32}, Vector{Float32}, typeof(p), Float32},)
wrapped_rhs = Void(FunctionWrappersWrapper(rhs!, argtypes, (Nothing,)))
du = zeros(Float32, 4)
ddu = Float32[1, 1, 0, 0]
u = Float32[2, 0, 0, 0]
dp = zero(p)
mode = get(ENV, "RUNTIME_ACTIVITY", "false") == "true" ?
Enzyme.set_runtime_activity(Enzyme.Reverse) : Enzyme.Reverse
Enzyme.autodiff(
mode, Enzyme.Const(wrapped_rhs), Enzyme.Const,
Enzyme.Duplicated(du, ddu), Enzyme.Const(u),
Enzyme.Duplicated(p, dp), Enzyme.Const(0.5f0)
)
println("gradient norm: ", sqrt(sum(abs2, dp)))
```
With Enzyme checked out at the main commit above:
```sh
julia +1.10 --startup-file=no --project=. mwe.jl
```
fails with exit 1:
```text
EnzymeRuntimeActivityError: Detected potential need for runtime activity.
...
Failure within method:
(::Void{FunctionWrappersWrapper{...}})(
::Vector{Float32}, ::Vector{Float32},
::ComponentVector{Float32, Vector{Float32}, ...}, ::Float32)
...
Mismatched activity for:
store {} addrspace(10)* %2, {} addrspace(10)** %.fca.1.gep
Julia value causing error: Unknown object of type Vector{Float32}
```
The runtime-activity control:
```sh
RUNTIME_ACTIVITY=true julia +1.10 --startup-file=no --project=. mwe.jl
```
passes with exit 0:
```text
gradient norm: 8.4620905
```
## Expected behavior
Static activity should either differentiate this wrapper correctly or identify a concrete user-level alias that must be annotated differently. The wrapped function itself and the wrapper object are constant; only the output and parameter arrays are duplicated.
## Relation to #3260
This removes SciMLSensitivity, OrdinaryDiffEq, callbacks, QuadGK, Optimization, and Zygote from the first activity failure. Because runtime activity fixes this reduced case, it should be treated as a narrower wrapper/activity problem and not yet as a complete reproducer for #3260's wrapper-loss or GC-corruption signatures.
Contributor guide
Assessment
This issue has not been assessed yet.