EnzymeAD / EnzymeAD/Enzyme.jl

Cannot store inactive but differentiable variable into active tuple

Open
#2,212 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

Full discussion is here: https://discourse.julialang.org/t/lux-enzyme-and-zygote-neuralode-segmentation-fault/123645/4

MWE (with full error below the code):

```
println("Load libraries")
using OrdinaryDiffEq, DifferentialEquations, ComponentArrays
using Enzyme, Lux
using Zygote, SciMLSensitivity
using StableRNGs

using Optimization, OptimizationOptimisers

using Plots

rng = StableRNG(1111)

println("Setup NODE")

in_size = 4
# This script works with Zygote.gradient when layer_size=5, but julia crashes on type inference when layer_size=10 when calling Zygote.gradient()
layer_size = 10

const sc = Lux.Chain(Lux.Dense(in_size,layer_size,tanh),
Lux.Dense(layer_size,layer_size,tanh),
Lux.Dense(layer_size,2))
# Get the initial parameters and state variables of the model
p_nn, st = Lux.setup(rng, sc)
const _st = st

function NODE!(du,u,p,t)
NN = sc([u;vs],p,_st)[1]
du[1] = u[2] + NN[1]
du[2] = NN[2]
end

println("Test NODE")
u0_test = [1.0,2.0]
vs = [3.,4.]

theta = ComponentArray(p=p_nn, u0=u0_test)
sc([theta.u0;vs], theta.p, _st)

println("Check if ODE solves")

ts = [0,.1,.2,.3,.4,.5]

prob = ODEProblem(NODE!, theta.u0, (0.,1.))
pred = solve(prob, Tsit5(), u0=theta.u0, p=theta.p, saveat=ts)[1,:]

R = rand(length(ts))

function predict_neuralode(theta)
Array(solve(prob, Tsit5(), u0 = theta.u0, p = theta.p, saveat = ts))
end

function loss_neuralode(theta)
pred = predict_neuralode(theta)[1,:]
loss = sum(abs2, R .- pred)
return loss
end

loss_neuralode(theta)

println("Check if works with Enzyme autodiff")

Enzyme.gradient(Reverse, loss_neuralode, theta)

# Testing Enzyme autodiff based on: https://docs.sciml.ai/SciMLSensitivity/dev/faq/

prob = ODEProblem(NODE!, theta.u0, (0.,1.), theta.p)
u0 = prob.u0
p = prob.p
tmp2 = Enzyme.make_zero(p)
t = prob.tspan[1]
du = zero(u0)

if DiffEqBase.isinplace(prob)
_f = prob.f
else
_f = (du, u, p, t) -> (du .= prob.f(u, p, t); nothing)
end

_tmp6 = Enzyme.make_zero(_f)
tmp3 = zero(u0)
tmp4 = zero(u0)
ytmp = zero(u0)
tmp1 = zero(u0)

# Error here
Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(_f, _tmp6),
Enzyme.Const, Enzyme.Duplicated(tmp3, tmp4),
Enzyme.Duplicated(ytmp, tmp1),
Enzyme.Duplicated(p, tmp2),
Enzyme.Const(t))
```

Error message:

```
ERROR: LoadError: "Error cannot store inactive but differentiable variable [0.0, 0.1, 0.2, 0.3, 0.4, 0.5] into active tuple"
Stacktrace:
[1] runtime_newstruct_augfwd(activity::Type{Val{(true, true, false)}}, width::Val{1}, ModifiedBetween::Val{(true, true, true, true)}, ::Type{@NamedTuple{u0::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, p::ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))}}}, saveat::Vector{Float64}}}, RT::Val{Any}, primal_1::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, shadow_1_1::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, primal_2::ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))}}}, shadow_2_1::ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))}}}, primal_3::Vector{Float64}, shadow_3_1::Nothing)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/aViNX/src/rules/typeunstablerules.jl:33
[2] NamedTuple
@ ./boot.jl:727 [inlined]
[3] predict_neuralode
@ ~/DropBounce/src/MWE.jl:52
[4] loss_neuralode
@ ~/DropBounce/src/MWE.jl:56 [inlined]
[5] loss_neuralode
@ ~/DropBounce/src/MWE.jl:0 [inlined]
[6] augmented_julia_loss_neuralode_17527_inner_1wrap
@ ~/DropBounce/src/MWE.jl:0
[7] macro expansion
@ ~/.julia/packages/Enzyme/aViNX/src/compiler.jl:5201 [inlined]
[8] enzyme_call
@ ~/.julia/packages/Enzyme/aViNX/src/compiler.jl:4747 [inlined]
[9] AugmentedForwardThunk
@ ~/.julia/packages/Enzyme/aViNX/src/compiler.jl:4683 [inlined]
[10] autodiff
@ ~/.julia/packages/Enzyme/aViNX/src/Enzyme.jl:396 [inlined]
[11] autodiff
@ ~/.julia/packages/Enzyme/aViNX/src/Enzyme.jl:524 [inlined]
[12] macro expansion
@ ~/.julia/packages/Enzyme/aViNX/src/sugar.jl:319 [inlined]
[13] gradient(::ReverseMode{false, false, FFIABI, false, false}, ::typeof(loss_neuralode), ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(p = ViewAxis(1:182, Axis(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))), u0 = 183:184)}}})
@ Enzyme ~/.julia/packages/Enzyme/aViNX/src/sugar.jl:258
[14] top-level scope
@ ~/DropBounce/src/MWE.jl:115
in expression starting at /home/vleon/DropBounce/src/MWE.jl:115
```

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.