Simple MLP requires Enzyme runtimeActivity

[avik-pal]

using Lux, ComponentArrays, Random, Enzyme
import Zygote

rng = Random.MersenneTwister(1234)

NN = Lux.Chain(Lux.Dense(5 => 5, tanh), Lux.Dense(5 => 1))

ps, st = Lux.setup(rng, NN)

xtest = [0.1, 0.2, 0.3, 0.4, 0.5]
NN(xtest, ps, st)[1][1]
dx = zeros(size(xtest)[1])

# Enzyme.API.runtimeActivity!(true)

function test_function(NN, x, ps, st)
    y, _ = NN(x, ps, st)
    return sum(y)
end

@time autodiff(
    Reverse, test_function, Active, Const(NN), Duplicated(xtest, dx), Const(ps), Const(st))

gives

ERROR: Enzyme execution failed.
Mismatched activity for:   store {} addrspace(10)* %9, {} addrspace(10)* addrspace(10)* %.repack.repack191, align 8, !dbg !336, !tbaa !99, !alias.scope !103, !noalias !104 const val:   %9 = call nonnull {} addrspace(10)* @ijl_reshape_array({} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 133482621319440 to {}*) to {} addrspace(10)*), {} addrspace(10)* noundef nonnull %getfield8, {} addrspace(10)* noundef nonnull %box18) #49, !dbg !97
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@float, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
  [1] copyto!
    @ ./broadcast.jl:956
  [2] materialize!
    @ ./broadcast.jl:914
  [3] materialize!
    @ ./broadcast.jl:911
  [4] __nonuniform_fast_broadcast!
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/utils.jl:130
  [5] __apply_bias_activation!!
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/utils.jl:100
  [6] __fused_dense_bias_activation_impl
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/impl/fused_dense.jl:31
  [7] fused_dense_bias_activation
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/api/dense.jl:46
  [8] fused_dense_bias_activation
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/api/dense.jl:38
  [9] Dense
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/basic.jl:218
 [10] Dense
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/basic.jl:214
 [11] apply
    @ ~/.julia/packages/LuxCore/qiHPC/src/LuxCore.jl:179

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/2FwRI/src/compiler.jl:1323
  [2] copyto!
    @ ./broadcast.jl:956 [inlined]
  [3] materialize!
    @ ./broadcast.jl:914 [inlined]
  [4] materialize!
    @ ./broadcast.jl:911 [inlined]
  [5] __nonuniform_fast_broadcast!
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/utils.jl:130 [inlined]
  [6] __apply_bias_activation!!
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/utils.jl:100 [inlined]
  [7] __fused_dense_bias_activation_impl
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/impl/fused_dense.jl:31 [inlined]
  [8] fused_dense_bias_activation
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/api/dense.jl:46 [inlined]
  [9] fused_dense_bias_activation
    @ /mnt/research/ongoing/lux/LuxLib.jl/src/api/dense.jl:38 [inlined]
 [10] Dense
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/basic.jl:218 [inlined]
 [11] Dense
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/basic.jl:214 [inlined]
 [12] apply
    @ ~/.julia/packages/LuxCore/qiHPC/src/LuxCore.jl:179
 [13] macro expansion
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/containers.jl:0 [inlined]
 [14] applychain
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/containers.jl:478 [inlined]
 [15] Chain
    @ /mnt/research/ongoing/lux/Lux.jl/src/layers/containers.jl:476 [inlined]
 [16] test_function
    @ /mnt/research/ongoing/lux/enzyme_ca.jl:20 [inlined]
 [17] test_function
    @ /mnt/research/ongoing/lux/enzyme_ca.jl:0 [inlined]
 [18] diffejulia_test_function_9960_inner_1wrap
    @ /mnt/research/ongoing/lux/enzyme_ca.jl:0
 [19] macro expansion
    @ ~/.julia/packages/Enzyme/2FwRI/src/compiler.jl:5703 [inlined]
 [20] enzyme_call
    @ ~/.julia/packages/Enzyme/2FwRI/src/compiler.jl:5381 [inlined]
 [21] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/2FwRI/src/compiler.jl:5260 [inlined]
 [22] autodiff
    @ ~/.julia/packages/Enzyme/2FwRI/src/Enzyme.jl:291 [inlined]
 [23] autodiff(::ReverseMode{…}, ::typeof(test_function), ::Type{…}, ::Const{…}, ::Duplicated{…}, ::Const{…}, ::Const{…})
    @ Enzyme ~/.julia/packages/Enzyme/2FwRI/src/Enzyme.jl:303
 [24] macro expansion
    @ ./timing.jl:279 [inlined]
 [25] top-level scope
    @ /mnt/research/ongoing/lux/enzyme_ca.jl:24
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

@avik-pal with EnzymeAD/Enzyme.jl#1445 merged can you see if this is fixed?

[avik-pal]

Yes it works on Enzyme master.