Error using Enzyme to autodiff wrt Flux neural net parameters

General Usage
3

Some more minimal cases which hit bugs, for @wsmoses … possibly even the same bugs?

julia> VERSION
v"1.10.4"

julia> v2, re2 = Flux.destructure((a=[1 2; 3 4.], b=2.0))  # easy case, involves reshape
([1.0, 3.0, 2.0, 4.0], Restructure(NamedTuple, ..., 4))

julia> Enzyme.jacobian(set_runtime_activity(Reverse), v -> re2(v).a[1,:], v2)[1]
2×4 transpose(::Matrix{Float64}) with eltype Float64:
 1.0  0.0  0.0  0.0
 0.0  0.0  1.0  0.0

julia> Enzyme.jacobian(Reverse, v -> re2(v).a[1,:], v2)[1]
ERROR: Constant memory is stored (or returned) to a differentiable variable.

julia> v3, re3 = Flux.destructure((a=[1 2; 3 4.], b=[5, 6.]))  # harder case
([1.0, 3.0, 2.0, 4.0, 5.0, 6.0], Restructure(NamedTuple, ..., 6))

julia> Enzyme.jacobian(set_runtime_activity(Reverse), v -> begin nt = re3(v); nt.a[1,:] .^ nt.b[1] end, v3)[1]
{[0]:Pointer, [0,-1]:Float@double, [8]:Integer, [9]:Integer, [10]:Integer, [11]:Integer, [12]:Integer, [13]:Integer,
...
 canonicalizing 8
ERROR: LLVM error: Canonicalization failed
Stacktrace:
 [1] handle_error(reason::Cstring)
   @ LLVM ~/.julia/packages/LLVM/wMjUU/src/core/context.jl:194

On 1.11, all give this error:

julia> VERSION
v"1.11.0"

julia> v2, re2 = Flux.destructure((a=[1 2; 3 4.], b=2.0))  # easy case
([1.0, 3.0, 2.0, 4.0], Restructure(NamedTuple, ..., 4))

julia> Enzyme.jacobian(set_runtime_activity(Reverse), v -> re2(v).a[1,:], v2)[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]
1 Like