Batch Norm should have in_dim

[albertz]

We need to know the dimensions which should be reduced to gather the statistics.

It would reduce over batch dim and spatial dims (in_spatial_dims).

Or, we specify the dims which it should not reduce, e.g. in_dim.

Or in_dims but that is not supported currently. Note that while conceptually the order of in_dims does not matter, it matters for the definition of the parameter. So it must be a sequence in case of multiple dims. Also, in_spatial_dims would not be enough as it would not define the order of in_dims. So that implies that specifying in_dims makes more sense than in_spatial_dims.

But anyway, maybe in_dim makes most sense for now.

Like #597. Or like #789 for ConvLayer.