Allow user to change the channel axis for BatchNorm function and the likes

[vboussange]

Proposition for issue #1664

This PR allows the user to customize the channel axis for normalisation functions (BatchNorm, GroupNorm and InstanceNorm).

Example

using Flux
channel_size = 3
channel_axis = 1
BN = BatchNorm(channel_size, dim = channel_axis)
x = randn(channel_size, 10, 10)
BN(x)

[DhairyaLGandhi]

This should never be nothing, but use N - 1 as default.

[ToucheSir]

How do you know N - 1 a priori though? The question is what is an appropriate sentinel value. Perhaps a negative offset from the end, since by default channels are at dim end - 1?

[DhairyaLGandhi]

Right so, it will need to be defined at runtime.

[DhairyaLGandhi]

How well is cudnn able to handle this?

[ToucheSir]

Looping back to answer the cuDNN question: it supports one other configuration for channel dim via CUDNN_TENSOR_NHWC. I think adding that to NNlibCUDA at https://github.com/FluxML/NNlibCUDA.jl/blob/96a334633ef3a3707c85fc1754c2c7eb8849db4e/src/cudnn/batchnorm.jl#L27 would be a good first step for getting this going again. Here are a couple example implementation PRs from MXNet and PyTorch.