 
     @differentiable(wrt: (self, input))
     public func applied(to input: Tensor<Scalar>, in _: Context) -> Tensor<Scalar> {
-        let batchSize = input.shape[0]
-        let height = input.shape[1]
-        let width = input.shape[2]
-        let channels = input.shape[3]
+        let shape = input.shape
+        let (batchSize, height, width, channels) = (shape[0], shape[1], shape[2], shape[3])
         let reshapeSize = Tensor<Int32>([batchSize, height, 1, width, 1, channels])
         let scaleOnes = Tensor<Scalar>(ones: [1, 1, size, 1, size, 1])
         let upSampling = input.reshaped(toShape: reshapeSize) * scaleOnes
         let upSampledShape = Tensor<Int32>([batchSize, height * size, width * size, channels])
-        let upSampled = upSampling.reshaped(toShape: upSampledShape)
-        return upSampled
-    }
-}
+        return upSampling.reshaped(toShape: upSampledShape)
+    }
+}

}

public extension Conv2D where Scalar.RawSignificand: FixedWidthInteger {
    init<G : RandomNumberGenerator>(
        filterShape: (Int, Int, Int, Int),
        strides: (Int, Int) = (1, 1),
        padding: Padding,

public struct BatchNorm<Scalar: TensorFlowFloatingPoint>: Layer {
    /// The batch dimension.
    @noDerivative public let axis: Int32
    /// The momentum for the running mean and running variance.
    @noDerivative public let momentum: Tensor<Scalar>
    /// The offset value, also known as beta.
    public var offset: Tensor<Scalar>
    /// The scale value, also known as gamma.
    public var scale: Tensor<Scalar>
    /// The variance epsilon value.
    @noDerivative public let epsilon: Tensor<Scalar>
    /// The running mean.
    @noDerivative public let runningMean: Parameter<Scalar>
    /// The running variance.
    @noDerivative public let runningVariance: Parameter<Scalar>

