Gradient calculation bug re-introduced in Flux v0.10.4 and Zygote v0.4.22

[alecokas]

I am seeing what looks like a bug which has been re-introduced between Zygote versions 0.4.20 and 0.4.22. By using the newer version of Zygote (which is automatically installed with Flux 0.10.4), I am unable to obtain gradients from a convolutional VAE Flux model. I am running Julia 1.4.1.

Here is a piece of code which you can use to reproduce the bug, taken from the more general discussions at here and further discussion here.

module FLUXVAE                                                                   
                                                                                 
using Flux                                                                       
using Flux: @epochs, binarycrossentropy                                          
using DistributionsAD                                                            
using Distributions                                                              
                                                                                 
# dummy data                                                                     
function dummy_data()                                                            
    d = Array{Float32}(zeros((796, 512, 1, 10))) .+ 1                            
    batches = [reshape(d[:,:,:, i:i+4], (796, 512, 1, 5)) for i in 1:5]          
end                                                                              
                                                                                 
struct Reshape                                                                   
    shape                                                                        
end                                                                              
                                                                                 
Reshape(args...) = Reshape(args)                                                 
(r::Reshape)(x) = reshape(x, r.shape)                                            
Flux.@functor Reshape ()                                                         
                                                                                 
                                                                                 
# convolutional encoder                                                          
function encoder()                                                               
    conv1 = Conv((14, 10), 1 => 4, relu, stride = (10, 10), pad = 4)             
    pool1 = MaxPool((8, 8), stride = (4, 4), pad = 2)                            
    conv2 = Conv((4, 3), 4 => 4, stride = (2, 2), pad = 1)                       
    res = Reshape(280, :)                                                        
    # enc1(X) = reshape(conv2(pool1(conv1(X))), (280, :))                        
    # Chain(res, conv2, pool1, conv1)                                            
    Chain(conv1, pool1, conv2, res)                                              
end                                                                              
                                                                                 
# decoder, I am using the one with transposed convolutions                       
function decoder(;dense_decoder = false)                                         
    if dense_decoder                                                             
        dec = Dense(4, 796*512, sigmoid)                                         
        dec1(X) = reshape(dec(X),  (796, 512, 1, :))                             
    else                                                                         
        interaction1 = Dense(4, 280)  # specific to my setup                     
        res = Reshape(10, 7, 4, :)                                               
        # int1(X) = reshape(interaction1(X), (10, 7, 4, :))                      
        tc1 = ConvTranspose((4, 3), 4 => 4, relu, stride = (2, 2), pad = 1)         
        tc2 = ConvTranspose((8, 8), 4 => 4, relu, stride = (4, 4), pad = 2)         
        tc3 = ConvTranspose((14, 10), 4 => 1, sigmoid, stride = (10, 10), pad = 4)
        dec = Chain(interaction1, tc1, tc2, tc3) # for params                    
        dec1 = Chain(interaction1, res, tc1, tc2, tc3)                           
    end                                                                          
    return (dec, dec1)                                                           
end                                                                              
                                                                                 
# sample from z-distribution                                                     
z(μ::T, logσ) where {T} = μ + exp(logσ) * randn(T)                               
z(μ, logσ, eps) = μ + exp(logσ) * eps                                            
                                                                                 
# log(p(x|z)), log(p(z)), log(q(z|x))                                            
logp_x_z1(X, z, dec1) = -sum(binarycrossentropy.(dec1(z), X))                    
logp_z(z::AbstractArray{T}) where {T} = sum((logpdf.(Normal(zero(T), one(T)), z)))
log_q_z_x(ϵ, log_sigma) = logpdf(Normal(zero(ϵ), one(ϵ)), ϵ) - log_sigma         
                                                                                 
# vae loss estimator                                                             
function vae_loss(enc1, dec1, μ1, logσ1)                                         
    mu(X) = μ1(enc1(X))                                                          
    l(X) =  logσ1(enc1(X))                                                       
    e(X) = randn(eltype(X), size(l(X))) # latentdim1                             
    z_(X) = z.(mu(X), l(X), e(X))                                                
    return X->-(logp_x_z1(X, z_(X), dec1) + logp_z(z_(X)) - sum(log_q_z_x.(e(X), l(X)))) * 1//5
end                                                                              
                                                                                 
# train vae1                                                                     
function train!()                                                                
    enc1 = encoder()                                                             
    dec, dec1 = decoder()                                                        
    # mean and log-variance of vae1's z-variable/latent space                    
    μ1 = Dense(280, 4)                                                           
    logσ1 = Dense(280, 4)                                                        
    L1 = vae_loss(enc1, dec1, μ1, logσ1)                                         
    ps1 = Flux.params(enc1, μ1, logσ1, dec1)                                     
    batches = dummy_data()                                                       
    @epochs 3 Flux.train!(L1, ps1, zip(batches), ADAM())                         
end                                                                              
                                                                                 
end                                                                               

When I manually downgrade to using Zygote 0.4.20, the above code is run successfully, however I see the following stack trace using Zygote 0.4.22:

julia> FLUXVAE.train!()
[ Info: Epoch 1
ERROR: MethodError: no method matching ∇maxpool(::Array{Float64,4}, ::Array{Float32,4}, ::Array{Float32,4}, ::PoolDims{2,(8, 8),(4, 4),(2, 2, 2, 2),(1, 1)})
Closest candidates are:
  ∇maxpool(::AbstractArray{T,N}, ::AbstractArray{T,N}, ::AbstractArray{T,N}, ::PoolDims; kwargs...) where {T, N} at /home/aleco/.julia/packages/NNlib/FAI3o/src/pooling.jl:123
Stacktrace:
 [1] (::Zygote.var"#1239#1240"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Array{Float32,4},PoolDims{2,(8, 8),(4, 4),(2, 2, 2, 2),(1, 1)},Array{Float32,4}})(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/lib/nnlib.jl:74
 [2] (::Zygote.var"#2759#back#1241"{Zygote.var"#1239#1240"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Array{Float32,4},PoolDims{2,(8, 8),(4, 4),(2, 2, 2, 2),(1, 1)},Array{Float32,4}}})(::Array{Float64,4}) at /home/aleco/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
 [3] MaxPool at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/conv.jl:377 [inlined]
 [4] (::typeof(∂(λ)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [5] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
 [6] (::typeof(∂(applychain)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [7] applychain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:36 [inlined]
 [8] (::typeof(∂(applychain)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [9] Chain at /home/aleco/.julia/packages/Flux/Fj3bt/src/layers/basic.jl:38 [inlined]
 [10] (::typeof(∂(λ)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [11] l at ./REPL[11]:66 [inlined]
 [12] (::typeof(∂(λ)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [13] z_ at ./REPL[11]:68 [inlined]
 [14] (::typeof(∂(λ)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [15] #7 at ./REPL[11]:69 [inlined]
 [16] (::typeof(∂(λ)))(::Float32) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [17] #175 at /home/aleco/.julia/packages/Zygote/1GXzF/src/lib/lib.jl:182 [inlined]
 [18] #347#back at /home/aleco/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
 [19] #17 at /home/aleco/.julia/packages/Flux/Fj3bt/src/optimise/train.jl:89 [inlined]
 [20] (::typeof(∂(λ)))(::Float32) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface2.jl:0
 [21] (::Zygote.var"#50#51"{Params,Zygote.Context,typeof(∂(λ))})(::Float32) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface.jl:177
 [22] gradient(::Function, ::Params) at /home/aleco/.julia/packages/Zygote/1GXzF/src/compiler/interface.jl:54
 [23] macro expansion at /home/aleco/.julia/packages/Flux/Fj3bt/src/optimise/train.jl:88 [inlined]
 [24] macro expansion at /home/aleco/.julia/packages/Juno/tLMZd/src/progress.jl:134 [inlined]
 [25] train!(::Main.FLUXVAE.var"#7#12"{Chain{Tuple{Dense{typeof(identity),Array{Float32,2},Array{Float32,1}},Main.FLUXVAE.Reshape,ConvTranspose{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},ConvTranspose{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},ConvTranspose{2,4,typeof(σ),Array{Float32,4},Array{Float32,1}}}},Main.FLUXVAE.var"#l#9"{Chain{Tuple{Conv{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},MaxPool{2,4},Conv{2,4,typeof(identity),Array{Float32,4},Array{Float32,1}},Main.FLUXVAE.Reshape}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}},Main.FLUXVAE.var"#e#10"{Main.FLUXVAE.var"#l#9"{Chain{Tuple{Conv{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},MaxPool{2,4},Conv{2,4,typeof(identity),Array{Float32,4},Array{Float32,1}},Main.FLUXVAE.Reshape}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}}},Main.FLUXVAE.var"#z_#11"{Main.FLUXVAE.var"#mu#8"{Chain{Tuple{Conv{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},MaxPool{2,4},Conv{2,4,typeof(identity),Array{Float32,4},Array{Float32,1}},Main.FLUXVAE.Reshape}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}},Main.FLUXVAE.var"#l#9"{Chain{Tuple{Conv{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},MaxPool{2,4},Conv{2,4,typeof(identity),Array{Float32,4},Array{Float32,1}},Main.FLUXVAE.Reshape}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}},Main.FLUXVAE.var"#e#10"{Main.FLUXVAE.var"#l#9"{Chain{Tuple{Conv{2,4,typeof(relu),Array{Float32,4},Array{Float32,1}},MaxPool{2,4},Conv{2,4,typeof(identity),Array{Float32,4},Array{Float32,1}},Main.FLUXVAE.Reshape}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}}}}}, ::Params, ::Base.Iterators.Zip{Tuple{Array{Array{Float32,4},1}}}, ::ADAM; cb::Flux.Optimise.var"#18#26") at /home/aleco/.julia/packages/Flux/Fj3bt/src/optimise/train.jl:81
 [26] train!(::Function, ::Params, ::Base.Iterators.Zip{Tuple{Array{Array{Float32,4},1}}}, ::ADAM) at /home/aleco/.julia/packages/Flux/Fj3bt/src/optimise/train.jl:79
 [27] macro expansion at /home/aleco/.julia/packages/Flux/Fj3bt/src/optimise/train.jl:122 [inlined]
 [28] macro expansion at /home/aleco/.julia/packages/Juno/tLMZd/src/progress.jl:134 [inlined]
 [29] train!() at ./REPL[11]:82
 [30] top-level scope at REPL[12]:1

As a side note, I wasn't sure where this issue should be raised since it effects both Flux.jl and Zygote.jl. I'm happy to cross-post if appropriate.

[CarloLucibello]

The problem is an incorrect gradient type for logp_z. Previous NNlib versions were more tolerant to type mismatches.

julia> using Zygote, Distributions, DistributionsAD

julia> logp_z(z::AbstractArray{T}) where {T} = sum((logpdf.(Normal(zero(T), one(T)), z)))

# correct type for the forward
julia> logp_z(rand(Float32,3))
-2.9049828f0

# wrong type in the backward
julia> gradient(x-> logp_z(x), rand(Float32, 3))[1]
3-element Array{Float64,1}:
 -0.8972270488739014
 -0.448053240776062
 -0.06319701671600342

[CarloLucibello]

This can be further reduced to a bug with division, that is hit when computing the logpdf here

julia> gradient(x -> x/2, 1f0)
(0.5,)

This is a very serious bug with Zygote that incredibly got through our CI

[CarloLucibello]

Could you test if this is fixed on latest releases?

[alecokas]

I have updated so that my package status looks as follows:

  [31c24e10] Distributions v0.23.4
  [ced4e74d] DistributionsAD v0.6.2
  [587475ba] Flux v0.11.0
  [6a3955dd] ImageFiltering v0.6.14
  [eb30cadb] MLDatasets v0.5.2
  [872c559c] NNlib v0.7.3

The Zygote dependency is now on version 0.5.1.

I have updated the loss function from the example above to use Losses.binarycrossentropy(ŷ, y, agg=identity).
Here is the new output:

julia> FLUXVAE.train!()
[ Info: Epoch 1
┌ Warning: Slow fallback implementation invoked for ∇conv_data!  You probably don't want this; check your datatypes.
│   yT = Float64
│   T1 = Float64
│   T2 = Float32
└ @ NNlib ~/.julia/packages/NNlib/sSn9M/src/conv.jl:206
┌ Warning: Slow fallback implementation invoked for ∇conv_filter!  You probably don't want this; check your datatypes.
│   yT = Float64
│   T1 = Float32
│   T2 = Float64
└ @ NNlib ~/.julia/packages/NNlib/sSn9M/src/conv.jl:206
ERROR: MethodError: no method matching ∇maxpool(::Array{Float64,4}, ::Array{Float32,4}, ::Array{Float32,4}, ::NNlib.PoolDims{2,(8, 8),(4, 4),(2, 2, 2, 2),(1, 1)})
Closest candidates are:
  ∇maxpool(::AbstractArray{T,N}, ::AbstractArray{T,N}, ::AbstractArray{T,N}, ::NNlib.PoolDims; kwargs...) where {T, N} at /home/aleco/.julia/packages/NNlib/sSn9M/src/pooling.jl:123
Stacktrace:
 [1] (::Zygote.var"#1239#1240"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Array{Float32,4},NNlib.PoolDims{2,(8, 8),(4, 4),(2, 2, 2, 2),(1, 1)},Array{Float32,4}})(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/lib/nnlib.jl:74
 [2] (::Zygote.var"#2759#back#1241"{Zygote.var"#1239#1240"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Array{Float32,4},NNlib.PoolDims{2,(8, 8),(4, 4),(2, 2, 2, 2),(1, 1)},Array{Float32,4}}})(::Array{Float64,4}) at /home/aleco/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
 [3] MaxPool at /home/aleco/.julia/packages/Flux/IjMZL/src/layers/conv.jl:643 [inlined]
 [4] (::typeof(∂(λ)))(::Array{Float64,4}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [5] applychain at /home/aleco/.julia/packages/Flux/IjMZL/src/layers/basic.jl:36 [inlined]
 [6] (::typeof(∂(applychain)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [7] applychain at /home/aleco/.julia/packages/Flux/IjMZL/src/layers/basic.jl:36 [inlined]
 [8] (::typeof(∂(applychain)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [9] Chain at /home/aleco/.julia/packages/Flux/IjMZL/src/layers/basic.jl:38 [inlined]
 [10] (::typeof(∂(λ)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [11] l at ./REPL[2]:65 [inlined]
 [12] (::typeof(∂(λ)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [13] z_ at ./REPL[2]:67 [inlined]
 [14] (::typeof(∂(λ)))(::Array{Float64,2}) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [15] #7 at ./REPL[2]:68 [inlined]
 [16] (::typeof(∂(λ)))(::Float32) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [17] #175 at /home/aleco/.julia/packages/Zygote/pmW1K/src/lib/lib.jl:182 [inlined]
 [18] #347#back at /home/aleco/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
 [19] #15 at /home/aleco/.julia/packages/Flux/IjMZL/src/optimise/train.jl:83 [inlined]
 [20] (::typeof(∂(λ)))(::Float32) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface2.jl:0
 [21] (::Zygote.var"#50#51"{Zygote.Params,Zygote.Context,typeof(∂(λ))})(::Float32) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface.jl:177
 [22] gradient(::Function, ::Zygote.Params) at /home/aleco/.julia/packages/Zygote/pmW1K/src/compiler/interface.jl:54
 [23] macro expansion at /home/aleco/.julia/packages/Flux/IjMZL/src/optimise/train.jl:82 [inlined]
 [24] macro expansion at /home/aleco/.julia/packages/Juno/tLMZd/src/progress.jl:134 [inlined]
 [25] train!(::Function, ::Zygote.Params, ::Base.Iterators.Zip{Tuple{Array{Array{Float32,4},1}}}, ::Flux.Optimise.ADAM; cb::Flux.Optimise.var"#16#22") at /home/aleco/.julia/packages/Flux/IjMZL/src/optimise/train.jl:80
 [26] train! at /home/aleco/.julia/packages/Flux/IjMZL/src/optimise/train.jl:78 [inlined]
 [27] macro expansion at /home/aleco/.julia/packages/Flux/IjMZL/src/optimise/train.jl:115 [inlined]
 [28] macro expansion at /home/aleco/.julia/packages/Juno/tLMZd/src/progress.jl:134 [inlined]
 [29] train!() at ./REPL[2]:81
 [30] top-level scope at REPL[3]:1

It looks like the original bug still persists.

[ShuhuaGao]

I have the same issue with

• Flux v0.11.3
• Zygote v0.5.17

I am new to Flux. When playing with the MNIST example, I simply set the input image type to Float64, and then encountered the following error after the first Conv + MaxPool:

┌ Info: Epoch 1
└ @ Main C:\Users\Gao Shuhua\.julia\packages\Flux\sY3yx\src\optimise\train.jl:136
┌ Warning: Slow fallback implementation invoked for conv!  You probably don't want this; check your datatypes.
│   yT = Float64
│   T1 = Float64
│   T2 = Float32
└ @ NNlib C:\Users\Gao Shuhua\.julia\packages\NNlib\1WXEb\src\conv.jl:206
MethodError: no method matching ∇maxpool(::Array{Float32,4}, ::Array{Float64,4}, ::Array{Float64,4}, ::PoolDims{2,(2, 2),(2, 2),(0, 0, 0, 0),(1, 1)})
Closest candidates are:
  ∇maxpool(::AbstractArray{T,N}, !Matched::AbstractArray{T,N}, !Matched::AbstractArray{T,N}, ::PoolDims; kwargs...) where {T, N} at C:\Users\Gao Shuhua\.julia\packages\NNlib\1WXEb\src\pooling.jl:119

It seems that there is no fallback for mixed-type computation in ∇maxpool.

[DhairyaLGandhi]

You could cast the model to use float 64 with Flux.f64

[ShuhuaGao]

You could cast the model to use float 64 with Flux.f64

Thanks @DhairyaLGandhi . I have cast the input to Float32 for efficiency. My key point above is to show that ∇maxpool cannot handle mixed types, though I am not sure whether it is a feature or a bug.

[DhairyaLGandhi]

We want to try to prevent people from foot gunning in most cases, but also need to respect user choices. It gets a bit more complicated with mixed precision training that we are working on.

[ToucheSir]

Both of Carlo's MWEs now output the correct type. I tried running the large example up top but killed it after 5+ minutes of unresponsiveness. Given that Zygote is now far more robust to mixed types with the recent projection changes, this is likely a non-issue now/