EnzymeAD / EnzymeAD/Enzyme.jl

SciML-free: runtime activity error through FunctionWrappersWrapper with Lux ComponentArray

Open
#3,330 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.