ConvTranspose produces nonsense results for 3D (aka 5D) inputs

[bcdarwin]

E.g.,

julia> ConvTranspose((2, 2, 2), stride=2, 4=>8)(rand(16, 16, 16, 4, 3))
# ERROR: DimensionMismatch("arrays could not be broadcast to a common size")

When the batch size is 1, the application 'succeeds' by coincidence due to broadcasting but the output is of the wrong size.

The PR (#311) that added this operation explicitly noted that only 2D support was added but @doc doesn't mention this and anyway an error should be signalled if the functionality is missing.

Examining the definition reveals that the channel and batch dims somehow end up swapped in the 2d vs 3d case:

ct2 = ConvTranspose((2, 2),    4=>5, stride=2)
ct3 = ConvTranspose((2, 2, 2), 4=>5, stride=2)
mat2 = rand(    16, 16, 4, 3);
mat3 = rand(16, 16, 16, 4, 3);
∇conv_data(mat3, ct3.weight, conv_transpose_dims(ct3, mat3)) |> size
# (32, 32, 32, 3, 5)
∇conv_data(mat2, ct2.weight, conv_transpose_dims(ct2, mat2)) |> size
# (32, 32, 5, 3)

[bcdarwin]

On some investigation, it looks like the issue is that the 2D case throws away the return value of the ∇conv_data call and thus misses the final transpose_swapbatch that the 3D case undergoes. However, the 2D result is the 'correct' one so I'm not immediately sure how to restructure the code to fix this ...

[bcdarwin]

Closing here since this is an NNlib issue.

[haampie]

Pinging @tejank10 for this issue as well, maybe he has some insights :)

Ok, this is indeed NNlib internals. For anyone coming across this issue, it might have happened because you mixed Float64 with Float32 (e.g. Float32 weights, Float64 input), because that takes another code path. Everything Float32 does not hit this issue.

[bcdarwin]

Moved here.