Closed avik-pal closed 1 month ago
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.
@avik-pal with https://github.com/EnzymeAD/Enzyme.jl/pull/1445 merged can you see if this is fixed?
Yes it works on Enzyme master.
gives