Consider the simple setup
using Lux, LinearAlgebra
import Optimization, Random, Enzyme, ComponentArrays
model = Chain(
Dense(2 => 16, Lux.tanh),
Dense(16, 1),
)
parameters, states = Lux.setup(Random.Xoshiro(42), model) |> f64
point = [0.0, 0.0]
adtype = AutoEnzyme()
function f(parameters, model, states, point)
return only(first(model(point, parameters, states)))
end
print("Testing f: $(f(parameters, model, states, point))\n")
manual_gradient = Enzyme.gradient(Enzyme.Reverse, f, parameters, Enzyme.Const(model), Enzyme.Const(states),Enzyme.Const(point))
print("Manual gradient: $manual_gradient\n")
On Julia v1.11.6 with Enzyme v0.13.66 this raises
Testing f: 0.4700911923782194
ERROR:
If you are using Enzyme by selecting the `AutoEnzyme` object from ADTypes, you may want to try setting the `mode` option as follows:
AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Forward))
AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))
This hint appears because DifferentiationInterface and Enzyme are both loaded. It does not necessarily imply that Enzyme is being called through DifferentiationInterface.
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: %.pn46 = phi {} addrspace(10)* [ %38, %L49 ], [ %27, %L42 ] const val: %27 = load {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %26, align 8, !dbg !188, !tbaa !190, !alias.scope !168, !noalias !169, !dereferenceable_or_null !193, !align !141, !enzyme_type !194, !enzymejl_source_type_Memory\7BFloat64\7D !0, !enzymejl_byref_MUT_REF !0
value=Unknown object of type Memory{Float64}
llvalue= %27 = load {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %26, align 8, !dbg !188, !tbaa !190, !alias.scope !168, !noalias !169, !dereferenceable_or_null !193, !align !141, !enzyme_type !194, !enzymejl_source_type_Memory\7BFloat64\7D !0, !enzymejl_byref_MUT_REF !0
Stacktrace:
[1] reshape
@ ./reshapedarray.jl:60
[2] reshape
@ ./reshapedarray.jl:129
[3] reshape
@ ./reshapedarray.jl:128
[4] make_abstract_matrix
@ ~/.julia/packages/Lux/H3WdN/src/utils.jl:204
[5] Dense
@ ~/.julia/packages/Lux/H3WdN/src/layers/basic.jl:343
Stacktrace:
[1] reshape
@ ./reshapedarray.jl:54 [inlined]
[2] reshape
@ ./reshapedarray.jl:129 [inlined]
[3] reshape
@ ./reshapedarray.jl:128 [inlined]
[4] make_abstract_matrix
@ ~/.julia/packages/Lux/H3WdN/src/utils.jl:204 [inlined]
[5] Dense
@ ~/.julia/packages/Lux/H3WdN/src/layers/basic.jl:343
[6] apply
@ ~/.julia/packages/LuxCore/XUV80/src/LuxCore.jl:155 [inlined]
[7] macro expansion
@ ~/.julia/packages/Lux/H3WdN/src/layers/containers.jl:0 [inlined]
[8] applychain
@ ~/.julia/packages/Lux/H3WdN/src/layers/containers.jl:482 [inlined]
[9] Chain
@ ~/.julia/packages/Lux/H3WdN/src/layers/containers.jl:480 [inlined]
Adding Const to f as well only makes the error log less readable, giving
ERROR: Enzyme.Compiler.EnzymeRuntimeActivityError(Cstring(0x00007fc704148521))
What’s the proper way to use Lux models in e.g. a loss function with Enzyme?