DARNN error

Error

ERROR: LoadError: MethodError: no method matching JuliennedArrays.Slices(::ChainRulesCore.NoTangent, ::Int64)
Closest candidates are:
  JuliennedArrays.Slices(::AbstractArray{Item, Dimensions}, ::Int64...) where {Item, Dimensions} at /Users/sdobber/.julia/packages/JuliennedArrays/fUHr7/src/JuliennedArrays.jl:230
Stacktrace:
  [1] (::SliceMap.var"#67#69"{Tuple{Int64}})(Δ::ChainRulesCore.NoTangent)
    @ SliceMap ~/.julia/packages/SliceMap/iZBr6/src/SliceMap.jl:191
  [2] (::SliceMap.var"#69#back#71"{SliceMap.var"#67#69"{Tuple{Int64}}})(Δ::ChainRulesCore.NoTangent)
    @ SliceMap ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
  [3] Pullback
    @ ./none:0 [inlined]
  [4] (::typeof(∂(λ)))(Δ::ChainRulesCore.NoTangent)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
  [5] (::Zygote.var"#547#552")(::Tuple{JuliennedArrays.Align{Float32, 2, Vector{Vector{Float32}}, Tuple{JuliennedArrays.True, JuliennedArrays.False}}, typeof(∂(λ))}, δ::ChainRulesCore.NoTangent)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/lib/array.jl:202
  [6] (::Base.var"#4#5"{Zygote.var"#547#552"})(a::Tuple{Tuple{JuliennedArrays.Align{Float32, 2, Vector{Vector{Float32}}, Tuple{JuliennedArrays.True, JuliennedArrays.False}}, typeof(∂(λ))}, ChainRulesCore.NoTangent})
    @ Base ./generator.jl:36
  [7] iterate
    @ ./generator.jl:47 [inlined]
  [8] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{Vector{Tuple{JuliennedArrays.Align{Float32, 2, Vector{Vector{Float32}}, Tuple{JuliennedArrays.True, JuliennedArrays.False}}, typeof(∂(λ))}}, Vector{Any}}}, Base.var"#4#5"{Zygote.var"#547#552"}})
    @ Base ./array.jl:678
  [9] map
    @ ./abstractarray.jl:2383 [inlined]
 [10] (::Zygote.var"#544#549"{FluxArchitectures.var"#6#8"{Vector{Tuple{Matrix{Float32}, Matrix{Float32}}}}, 1, Tuple{UnitRange{Int64}}, Vector{Tuple{JuliennedArrays.Align{Float32, 2, Vector{Vector{Float32}}, Tuple{JuliennedArrays.True, JuliennedArrays.False}}, typeof(∂(λ))}}})(Δ::Vector{Any})
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/lib/array.jl:202
 [11] (::Zygote.var"#back#580"{Zygote.var"#544#549"{FluxArchitectures.var"#6#8"{Vector{Tuple{Matrix{Float32}, Matrix{Float32}}}}, 1, Tuple{UnitRange{Int64}}, Vector{Tuple{JuliennedArrays.Align{Float32, 2, Vector{Vector{Float32}}, Tuple{JuliennedArrays.True, JuliennedArrays.False}}, typeof(∂(λ))}}}})(ȳ::Vector{Any})
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/lib/array.jl:235
 [12] Pullback
    @ ~/Uni/Code/FluxArchitectures/src/shared/Sequentialize.jl:69 [inlined]
 [13] (::typeof(∂(λ)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
 [14] Pullback
    @ ~/Uni/Code/FluxArchitectures/src/shared/Sequentialize.jl:51 [inlined]
 [15] (::typeof(∂(λ)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
 [16] Pullback
    @ ~/Uni/Code/FluxArchitectures/src/DARNN.jl:169 [inlined]
 [17] (::typeof(∂(darnn_decoder)))(Δ::Matrix{Float32})
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
 [18] Pullback
    @ ~/Uni/Code/FluxArchitectures/src/DARNN.jl:123 [inlined]
 [19] (::typeof(∂(λ)))(Δ::Matrix{Float32})
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
 [20] Pullback
    @ ~/Uni/Code/FluxArchitectures/examples/Example_DARNN.jl:33 [inlined]
 [21] (::typeof(∂(loss)))(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
 [22] #203
    @ ~/.julia/packages/Zygote/Lw5Kf/src/lib/lib.jl:203 [inlined]
 [23] #1734#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [24] Pullback
    @ ~/.julia/packages/Flux/ZnXxS/src/optimise/train.jl:105 [inlined]
 [25] (::typeof(∂(λ)))(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface2.jl:0
 [26] (::Zygote.var"#84#85"{Zygote.Params, typeof(∂(λ)), Zygote.Context})(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface.jl:343
 [27] gradient(f::Function, args::Zygote.Params)
    @ Zygote ~/.julia/packages/Zygote/Lw5Kf/src/compiler/interface.jl:76
 [28] macro expansion
    @ ~/.julia/packages/Flux/ZnXxS/src/optimise/train.jl:104 [inlined]
 [29] macro expansion
    @ ~/.julia/packages/Juno/n6wyj/src/progress.jl:134 [inlined]
 [30] train!(loss::Function, ps::Zygote.Params, data::Base.Iterators.Take{Base.Iterators.Repeated{Tuple{Array{Float32, 4}, Vector{Float32}}}}, opt::ADAM; cb::Flux.Optimise.var"#40#46")
    @ Flux.Optimise ~/.julia/packages/Flux/ZnXxS/src/optimise/train.jl:102
 [31] train!(loss::Function, ps::Zygote.Params, data::Base.Iterators.Take{Base.Iterators.Repeated{Tuple{Array{Float32, 4}, Vector{Float32}}}}, opt::ADAM)
    @ Flux.Optimise ~/.julia/packages/Flux/ZnXxS/src/optimise/train.jl:100
 [32] top-level scope
    @ ~/Uni/Code/FluxArchitectures/examples/Example_DARNN.jl:49
in expression starting at /Users/sdobber/Uni/Code/FluxArchitectures/examples/Example_DARNN.jl:49

Example.jl

using Flux
using Flux.Zygote

# Define a version of `Recur` that returns the whole hidden state
mutable struct HiddenRecur{T,S}
    cell::T
    state::S
end

HiddenRecur(m) = HiddenRecur(m, m.state0)

function (m::HiddenRecur)(xs...)
    h, y = m.cell(m.state, xs...)
    m.state = h
    return h  # collect(h)
end
function (m::HiddenRecur)(x::AbstractArray{T,3}) where T
	h = [m(view(x, :, :, i)) for i in 1:size(x, 3)]
	sze = size(h[1][1])
	return [reshape(reduce(hcat, [h[j][i] for j = 1:length(h)]), sze[1], sze[2], length(h)) for i = 1:length(h[1])]
end


# Feed arrays sequentially to recurrent neural nets
mutable struct Seq{T}
    chain::T
    state
end
Seq(chain) = Seq(chain, [0.0f0])
(l::Seq)(x) = l(l.chain.state, l.chain, x)
_add_dims(x::AbstractArray{T,2}) where T = reshape(x, size(x, 1), 1, size(x, 2))
_remove_dims(x::AbstractArray{T,3}) where T = dropdims(x; dims=2)
function (l::Seq)(::Tuple, _, x)
	l.state = _remove_dims.(l.chain(_add_dims(x)))
	return l.state
end


# Quick gradient
function bw(m, ip)
    gs = gradient((m, x) -> sum(m(x)), m, ip)
end

# Actual Code
inp = rand(Float32, 10, 50)
encoder = Seq(HiddenRecur(Flux.LSTMCell(10, 5)))

encoder(inp)  # forward pass works
bw(x -> encoder(x)[1], inp)  # works in Zygote 0.6.21, errors in 0.6.22