Flux.destructure gives DimensionMismatch error in backward pass due to Chain of mutable struct(s)

[ngphuoc]

I was trying a hypernetwork of a custom network (mutable struct) when I encountered the following error:

using Flux

mutable struct Linear  # Will give an error
# struct Linear  # no error
    w
end
Flux.@functor Linear

Linear(i, o) = Linear(Flux.glorot_normal(o, i))

(m::Linear)(x) = m.w*x

n = Chain(Linear(2, 2), )
# n = Linear(2, 2)  # no error for mutable struct
θ, re = Flux.destructure(n)
m = Chain(Dense(2, length(θ)), re)
x, y = randn(2), randn(2)
loss(x, y) = Flux.logitcrossentropy(m(x)(x), y)
loss(x, y)
l, back = Flux.pullback(params(m)) do
    loss(x, y)
end
grad = back(1f0)
ERROR: DimensionMismatch("new dimensions (4, 1) must be consistent with array size 0")
Stacktrace:
 [1] (::Base.var"#throw_dmrsa#213")(::Tuple{Int64,Int64}, ::Int64) at ./reshapedarray.jl:41
 [2] reshape at ./reshapedarray.jl:45 [inlined]
 [3] #446 at /home/ngphuoc/.julia/dev/Zygote/src/lib/array.jl:94 [inlined]
 [4] #2418#back at /home/ngphuoc/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
 [5] #163 at /home/ngphuoc/.julia/dev/Zygote/src/lib/lib.jl:191 [inlined]
 [6] #1695#back at /home/ngphuoc/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
 [7] Dense at /home/ngphuoc/.julia/dev/Flux/src/layers/basic.jl:127 [inlined]
 [8] (::typeof(∂(λ)))(::Array{Any,1}) at /home/ngphuoc/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [9] applychain at /home/ngphuoc/.julia/dev/Flux/src/layers/basic.jl:36 [inlined]
 [10] (::typeof(∂(applychain)))(::NamedTuple{(:layers,),Tuple{Tuple{Base.RefValue{Any}}}}) at /home/ngphuoc/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [11] Chain at /home/ngphuoc/.julia/dev/Flux/src/layers/basic.jl:38 [inlined]
 [12] (::typeof(∂(λ)))(::NamedTuple{(:layers,),Tuple{Tuple{Base.RefValue{Any}}}}) at /home/ngphuoc/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [13] loss at ./REPL[11]:1 [inlined]
 [14] (::typeof(∂(loss)))(::Float32) at /home/ngphuoc/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [15] #3 at ./REPL[13]:2 [inlined]
 [16] (::typeof(∂(#3)))(::Float32) at /home/ngphuoc/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [17] (::Zygote.var"#54#55"{Zygote.Params,Zygote.Context,typeof(∂(#3))})(::Float32) at /home/ngphuoc/.julia/dev/Zygote/src/compiler/interface.jl:181
 [18] top-level scope at REPL[14]:1

[ToucheSir]

This works without erroring on v0.12.8.