Flux.jl: Conv with non quadratic size broken

[jw3126]

The following used to work in older versions of Flux, but errors now:

julia> using Flux

julia> randn(Float32, 10,10,1,1) |> Conv((6,6), 1=>1, Flux.σ)
Tracked 5×5×1×1 Array{Float32,4}:
[:, :, 1, 1] =
 0.295454   0.541948  0.0864369  0.959286   0.891815 
 0.41697    0.218508  0.4821     0.746911   0.499474 
 0.500112   0.898844  0.592479   0.637593   0.874677 
 0.0124033  0.168184  0.0669264  0.024803   0.109923 
 0.286872   0.219004  0.0657297  0.0456955  0.0598319

julia> randn(Float32, 10,10,1,1) |> Conv((6,1), 1=>1, Flux.σ)
ERROR: DimensionMismatch("new dimensions (5, 5, 6, 1, 1) must be consistent with array size 300")

[MikeInnes]

@tejank10 can you take a look at this?

[thebhatman]

After quite a bit of debugging, I found out the error lies in https://github.com/FluxML/NNlib.jl/blob/d07ac0bfd3c71c3a29bc9c22becbba19227bbeb5/src/impl/conv.jl#L21. Here, as you can see, ouput height dimension is using K[1], instead of K[2], S[1] instead of S[2]. Which is why, it was working only for kernel sizes where width and height were equal.

[thebhatman]

Changing the output size calculation to this, fixes the issue.

function output_size(c::ConvDims{I,K,C,S,P,D,F}) where {I, K, C, S, P, D, F}
    O_w = div(I[1] + P[1] + P[2] - (K[1] - 1) * D[1] - 1, S[1]) + 1
    O_h = div(I[2] + P[3] + P[4] - (K[2] - 1) * D[2] - 1, S[2]) + 1
    return (O_w, O_h)
end

[tejank10]

Thanks, @jw3126 for pointing this out and adding tests, and @thebhatman for a quick fix!

[MikeInnes]

This is fixed on newer NNlibs. We'll tag a release with that soon.

[nmoran]

Also possible to override library version in local file until fix makes it to stable version.

For example:

using Flux: NNlib.ConvDims

function Flux.NNlib.output_size(c::ConvDims{I,K,C,S,P,D,F}) where {I, K, C, S, P, D, F}
    O_w = div(I[1] + P[1] + P[2] - (K[1] - 1) * D[1] - 1, S[1]) + 1
    O_h = div(I[2] + P[3] + P[4] - (K[2] - 1) * D[2] - 1, S[2]) + 1
    return (O_w, O_h)
end