Enzyme.jl
Enzyme.jl copied to clipboard
EnzymeAD
[Julia 1.11] Broadcast fails with Const arguments
Given a function like:
function lagrangian_elastic_rope(q, qdot, p)
N = length(q)
(;m, k, g) = p
# kinetic energy is just 0.5*mass*velocity^2
Eₖ = 0.5 * sum(m .* qdot.^2)
# potential energy is the elastic energy of the springs plus the gravitational potential
Eₚ = k[1] * q[1]^2
for i in 2:N
Eₚ += k[i] * (q[i] - q[i-1])^2
end
Eₚ = 0.5*Eₚ - g * sum(m .* q)
return Eₖ - Eₚ # lagrangian = kinetic_energy - potential_energy
end
function parameters_elastic_rope(N=3)
# a helper to setup the parameters of the problem
N = min(N, 5)
# time domain
t0 = 0
tend = 50
# initial condition
q0 = [1.0, 2.0, 3.0, 4.0, 5.0]
qdot0 = [0.2, 0.3, 0.7, 0.4, 0.5]
q0, qdot0 = q0[1:N], qdot0[1:N]
# parameters
m = [1.1, 1.2, 1.3, 1.4, 1.5]
k = [1.5, 1.3, 1.7, 1.9, 2.1]
g = 1.
m, k = m[1:N], k[1:N]
params = (;m, k, g)
return params, q0, qdot0, t0, tend
end
p, q0, qdot0, t0, tend = parameters_elastic_rope(3)
lagrangian_elastic_rope(q0, qdot0, p)
I want to be able to take the derivative of lagrangian_elastic_rope solely w.r.t q0
dq = Enzyme.make_zero(q0)
autodiff(Reverse, lagrangian_elastic_rope, Active, Duplicated(q0, dq), Const(qdot0), Const(p))
dq
This used to work in 1.10, but fails in 1.11 with a:
ERROR: Constant memory is stored (or returned) to a differentiable variable.
As a result, Enzyme cannot provably ensure correctness and throws this error.
This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Runtime-Activity).
If Enzyme should be able to prove this use non-differentable, open an issue!
To work around this issue, either:
a) rewrite this variable to not be conditionally active (fastest, but requires a code change), or
b) set the Enzyme mode to turn on runtime activity (e.g. autodiff(set_runtime_activity(Reverse), ...) ). This will maintain correctness, but may slightly reduce performance.
Mismatched activity for: %nodecayed..pre-phi = phi {} addrspace(10)* [ %94, %L131.i ], [ %.fca.0.extract, %guard_exit160.i ], [ %.fca.0.extract, %L56.i ], [ %.fca.0.extract, %L89.i ], !dbg !174 const val: %.fca.0.extract = extractvalue { {} addrspace(10)*, {} addrspace(10)*, double } %2, 0, !dbg !15, !enzyme_type !16, !enzymejl_source_type_Vector\7BFloat64\7D !0, !enzymejl_byref_MUT_REF !0
value=Unknown object of type Vector{Float64}
llvalue= %.fca.0.extract = extractvalue { {} addrspace(10)*, {} addrspace(10)*, double } %2, 0, !dbg !15, !enzyme_type !16, !enzymejl_source_type_Vector\7BFloat64\7D !0, !enzymejl_byref_MUT_REF !0
Stacktrace:
[1] size
@ ./array.jl:194
[2] axes
@ ./abstractarray.jl:98
[3] newindexer
@ ./broadcast.jl:598
[4] extrude
@ ./broadcast.jl:645
[5] preprocess
@ ./broadcast.jl:953
[6] preprocess_args
@ ./broadcast.jl:955
[7] preprocess
@ ./broadcast.jl:952
[8] override_bc_copyto!
@ ~/.julia/packages/Enzyme/RIEYk/src/compiler/interpreter.jl:798
[9] copyto!
@ ./broadcast.jl:925
[10] copy
@ ./broadcast.jl:897
[11] materialize
@ ./broadcast.jl:872
[12] lagrangian_elastic_rope
@ ./REPL[6]:5
[13] lagrangian_elastic_rope
@ ./REPL[6]:0
Stacktrace:
[1] broadcast_unalias
@ ./broadcast.jl:946 [inlined]
[2] preprocess
@ ./broadcast.jl:953 [inlined]
[3] preprocess_args
@ ./broadcast.jl:955 [inlined]
[4] preprocess
@ ./broadcast.jl:952 [inlined]
[5] override_bc_copyto!
@ ~/.julia/packages/Enzyme/RIEYk/src/compiler/interpreter.jl:798 [inlined]
[6] copyto!
@ ./broadcast.jl:925 [inlined]
[7] copy
@ ./broadcast.jl:897 [inlined]
[8] materialize
@ ./broadcast.jl:872 [inlined]
[9] lagrangian_elastic_rope
@ ./REPL[6]:5 [inlined]
[10] lagrangian_elastic_rope
@ ./REPL[6]:0 [inlined]
[11] diffejulia_lagrangian_elastic_rope_11596_inner_3wrap
@ ./REPL[6]:0
[12] macro expansion
@ ~/.julia/packages/Enzyme/RIEYk/src/compiler.jl:5463 [inlined]
[13] enzyme_call
@ ~/.julia/packages/Enzyme/RIEYk/src/compiler.jl:4997 [inlined]
[14] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/RIEYk/src/compiler.jl:4869 [inlined]
[15] autodiff
@ ~/.julia/packages/Enzyme/RIEYk/src/Enzyme.jl:504 [inlined]
[16] autodiff(::ReverseMode{false, false, FFIABI, false, false}, ::typeof(lagrangian_elastic_rope), ::Type{Active}, ::Duplicated{Vector{…}}, ::Const{Vector{…}}, ::Const{@NamedTuple{…}})
@ Enzyme ~/.julia/packages/Enzyme/RIEYk/src/Enzyme.jl:525
[17] top-level scope
@ REPL[11]:1
To "fix" this I must pass in:
dq = Enzyme.make_zero(q0)
autodiff(Reverse, lagrangian_elastic_rope, Active, Duplicated(q0, dq), Duplicated(qdot0, Enzyme.make_zero(qdot0)), Duplicated(p, Enzyme.make_zero(p)))
dq
julia> dq
3-element Vector{Float64}:
0.9000000000000004
1.5999999999999999
-0.3999999999999999
I think this is a 1.11 not a regression. In particular our alias check fixes can't handle the extra pointer by 1.11
