Error using Enzyme to autodiff wrt Flux neural net parameters

General Usage
2

I get a different error, below is 1.10 (which I believe is still better-tested) but see the same on 1.11.0:

julia> Enzyme.jacobian(Reverse, p -> enzSL_0f0(eval_model(re(p), X)', L, μ, k, r, Ecm), params)
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:   store {} addrspace(10)* %.fca.0.0.0.0.extract, {} addrspace(10)* addrspace(10)* %.repack.repack, align 8, !dbg !69, !tbaa !43, !alias.scope !47, !noalias !48 const val:   %.fca.0.0.0.0.extract = extractvalue { [1 x [6 x { {} addrspace(10)*, {} addrspace(10)* }]], [1 x [6 x { i64, i64 }]], i64 } %0, 0, 0, 0, 0, !dbg !8, !enzyme_type !51, !enzymejl_byref_MUT_REF !0, !enzymejl_source_type_Matrix\7BFloat64\7D !0
 value=Unknown object of type Matrix{Float64}
 llvalue=  %.fca.0.0.0.0.extract = extractvalue { [1 x [6 x { {} addrspace(10)*, {} addrspace(10)* }]], [1 x [6 x { i64, i64 }]], i64 } %0, 0, 0, 0, 0, !dbg !8, !enzyme_type !51, !enzymejl_byref_MUT_REF !0, !enzymejl_source_type_Matrix\7BFloat64\7D !0

Stacktrace:
  [1] _trainmap
    @ ~/.julia/packages/Optimisers/a4OnF/src/destructure.jl:114
  [2] _Trainable_biwalk
    @ ~/.julia/packages/Optimisers/a4OnF/src/destructure.jl:110
  [3] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126

(jl_vmDBZd) pkg> st
Status `/private/var/folders/yq/4p2zwd614y59gszh7y9ypyhh0000gn/T/jl_vmDBZd/Project.toml`
  [7da242da] Enzyme v0.13.24
  [587475ba] Flux v0.16.0
  [276daf66] SpecialFunctions v2.5.0
⌃ [c489a379] SphericalHarmonics v0.1.14
Info Packages marked with ⌃ have new versions available and may be upgradable.

julia> versioninfo()
Julia Version 1.10.4
Commit 48d4fd48430 (2024-06-04 10:41 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 11 × Apple M3 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)

As suggested, set_runtime_activity(Reverse) does avoid it, at least on 1.10:

julia> Enzyme.jacobian(set_runtime_activity(Reverse), p -> enzSL_0f0(eval_model(re(p), X)', L, μ, k, r, Ecm), params)[1]
2×1202 transpose(::Matrix{Float64}) with eltype Float64:
 -2.7278e-17   -4.22341e-18  -3.93169e-18  0.0  …  0.0  0.0  0.0  0.0   1.44062e-18  1.86651e-17
  2.95072e-17   2.51867e-17  -7.68234e-18  0.0     0.0  0.0  0.0  0.0  -1.86646e-17  1.34197e-18

julia> VERSION
v"1.10.4"

But on 1.11 there’s another error about reshape:

julia> Enzyme.jacobian(set_runtime_activity(Reverse), p -> enzSL_0f0(eval_model(re(p), X)', L, μ, k, r, Ecm), params)[1]
ERROR: 
No augmented forward pass found for jl_genericmemory_slice
 at context:   %93 = call nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @jl_genericmemory_slice({} addrspace(10)* nonnull %76, i64 %92, i64 %40) #28, !dbg !234

Stacktrace:
 [1] reshape
   @ ./reshapedarray.jl:55
 [2] reshape
   @ ./reshapedarray.jl:121
 [3] _getat
   @ ~/.julia/packages/Optimisers/a4OnF/src/destructure.jl:102


Stacktrace:
  [1] reshape
    @ ./reshapedarray.jl:55 [inlined]
  [2] reshape
    @ ./reshapedarray.jl:121 [inlined]
  [3] _getat
    @ ~/.julia/packages/Optimisers/a4OnF/src/destructure.jl:102
  [4] #57
    @ ~/.julia/packages/Optimisers/a4OnF/src/destructure.jl:97 [inlined]
  [5] ExcludeWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:126 [inlined]
  [6] CachedWalk
    @ ~/.julia/packages/Functors/LbNAu/src/walks.jl:177 [inlined]

julia> VERSION
v"1.11.0"
1 Like