Pullback within pullback throws error when using swish activation function

[IsaacBreen]

This code works.

using Flux, Zygote

# Generate a rand input
x = rand(1,1) |> gpu

# Construct a neural net
m = Dense(1,1,tanh) |> gpu

# Get ∇m⋅m
function ∇m_m(x)
    out, back = Zygote.pullback(m, x)
    return back(out)
end

pullback(∇m_m, x)

But when I change tanh to swish, I get the following error at out, back = Zygote.pullback(m, x).

ERROR: CuArray only supports bits types
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] CUDA.CuArray{Tuple{Float32, typeof(∂(#1047))}, 2}(#unused#::UndefInitializer, dims::Tuple{Int64, Int64})
    @ CUDA ~/.julia/packages/CUDA/Zmd60/src/array.jl:19
  [3] (CUDA.CuArray{Tuple{Float32, typeof(∂(#1047))}, N} where N)(#unused#::UndefInitializer, dims::Tuple{Int64, Int64})
    @ CUDA ~/.julia/packages/CUDA/Zmd60/src/array.jl:76
  [4] similar(#unused#::Type{CUDA.CuArray{Tuple{Float32, typeof(∂(#1047))}, N} where N}, dims::Tuple{Int64, Int64})
    @ Base ./abstractarray.jl:779
  [5] similar(#unused#::Type{CUDA.CuArray{Tuple{Float32, typeof(∂(#1047))}, N} where N}, shape::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}})
    @ Base ./abstractarray.jl:778
  [6] similar(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{2}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, Zygote.var"#507#511"{Zygote.Context, Zygote.var"#1047#1051"}, Tuple{Base.Broadcast.Extruded{CUDA.CuArray{ForwardDiff.Dual{Nothing, Float32, 1}, 2}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}}}, #unused#::Type{Tuple{Float32, typeof(∂(#1047))}})
    @ CUDA ~/.julia/packages/CUDA/Zmd60/src/broadcast.jl:11
  [7] copy
    @ ./broadcast.jl:923 [inlined]
  [8] materialize(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{2}, Nothing, Zygote.var"#507#511"{Zygote.Context, Zygote.var"#1047#1051"}, Tuple{CUDA.CuArray{ForwardDiff.Dual{Nothing, Float32, 1}, 2}}})
    @ Base.Broadcast ./broadcast.jl:883
  [9] map(::Function, ::CUDA.CuArray{ForwardDiff.Dual{Nothing, Float32, 1}, 2})
    @ GPUArrays ~/.julia/packages/GPUArrays/WV76E/src/host/broadcast.jl:89
 [10] ∇map(cx::Zygote.Context, f::Function, args::CUDA.CuArray{ForwardDiff.Dual{Nothing, Float32, 1}, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/lib/array.jl:180
 [11] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/array.jl:196 [inlined]
 [12] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [13] _pullback
    @ ~/.julia/packages/Zygote/KpME9/src/lib/broadcast.jl:216 [inlined]
 [14] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [15] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [16] _pullback
    @ ~/.julia/packages/Zygote/KpME9/src/lib/broadcast.jl:226 [inlined]
 [17] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [18] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [19] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:177 [inlined]
 [20] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [21] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [22] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [23] _pullback
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [24] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [25] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [26] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Zygote.Context, ::typeof(Core._apply_iterate), ::typeof(iterate), ::typeof(Base.Broadcast.broadcasted), ::Tuple{CUDA.CuArrayStyle{2}, typeof(swish), CUDA.CuArray{Float32, 2}}, ::Tuple{})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [27] _pullback
    @ ./broadcast.jl:1307 [inlined]
 [28] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Zygote.Context, ::typeof(Base.Broadcast.broadcasted), ::typeof(swish), ::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [29] _pullback
    @ ~/.julia/packages/Flux/goUGu/src/layers/basic.jl:123 [inlined]
 [30] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Zygote.Context, ::typeof(invoke), ::Dense{typeof(swish), CUDA.CuArray{Float32, 2}, CUDA.CuArray{Float32, 1}}, ::Type{Tuple{AbstractArray}}, ::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [31] _pullback
    @ ~/.julia/packages/Flux/goUGu/src/layers/basic.jl:134 [inlined]
 [32] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Zygote.Context, ::Dense{typeof(swish), CUDA.CuArray{Float32, 2}, CUDA.CuArray{Float32, 1}}, ::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [33] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [34] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [35] _pullback
    @ ~/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:33 [inlined]
 [36] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Dense{typeof(swish), CUDA.CuArray{Float32, 2}, CUDA.CuArray{Float32, 1}}, ::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [37] adjoint
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:188 [inlined]
 [38] _pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]
 [39] _pullback
    @ ~/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:39 [inlined]
 [40] _pullback(::Zygote.Context, ::typeof(pullback), ::Dense{typeof(swish), CUDA.CuArray{Float32, 2}, CUDA.CuArray{Float32, 1}}, ::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [41] _pullback
    @ ./REPL[4]:3 [inlined]
 [42] _pullback(ctx::Zygote.Context, f::typeof(∇m_m), args::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [43] _pullback(f::Function, args::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:33
 [44] pullback(f::Function, args::CUDA.CuArray{Float32, 2})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:39
 [45] top-level scope
    @ REPL[5]:1

I'm running on Debian GNU/Linux 10 and I encounter the same error on both Julia 1.5.3 and 1.6.0-beta1.

[DhairyaLGandhi]

I am guessing swish itself might have issues working on the GPU, do you think you can try with a cpu? We should get help from #1472

[IsaacBreen]

I don't have access to the machine I originally encountered this error on, but when I run the same code on a Colab GPU instance it works fine. In three lines:

!wget -q https://julialang-s3.julialang.org/bin/linux/x64/1.6/julia-1.6.0-linux-x86_64.tar.gz
!tar zxf julia-1.6.0-linux-x86_64.tar.gz
!/content/julia-1.6.0/bin/julia -e "import Pkg; Pkg.add([\"Flux\",\"Zygote\",\"CUDA\"]); using Flux, Zygote; x = rand(1,1) |> gpu; m = Dense(1,1,tanh) |> gpu; function ∇m_m(x) ;out, back = Zygote.pullback(m, x); return back(out); end ;pullback(∇m_m, x)"

So I guess we could close this for now as long as nobody else is reporting the same issue.