use .moments() in LayerNorm and BatchNorm layers

[mikowals]

Mean and variance in the layers are now calculated using Tensor.moments(). I also added tests for both BatchNorm and LayerNorm layers.

The tests turned up a flaw in how the shape of the scale and offset which were always of shape [featureCount] irrespective of the input shape or axis for normalisation. That shape leads to incorrect broadcasting when the axis being normalized along is not the last axis.

I have fixed this by always reshaping scale and offset before they are used. This seems hacky in that I get the shapes from the calculated mean and variance. Without the input shape being known at initialization time though I couldn't see a better way to do this.

I think the axis argument is probably there to be consistent with Keras but most of the Swift api layers assume inputs and activations are NHWC. So requiring NHWC, eliminating the axis argument, and the setting the correct shapes in init() would be another option.

[rxwei]

Nice!

[BradLarson]

@mikowals - The inference changes introduced here were triggering shape mismatches within the MiniGo sample project and its tests on tensorflow/swift-models. I've submitted a PR to roll those back here: #426 , but if you have a better way to approach this, I'm all ears.