         self.padding = padding
     }
 
+    /// Creates a average pooling layer.
+    ///
+    /// - Parameters:
+    ///   - poolSize: Vertical and horizontal factors by which to downscale.
+    ///   - strides: The strides.
+    ///   - padding: The padding.
     public init(poolSize: (Int, Int), strides: (Int, Int), padding: Padding = .valid) {
         self.poolSize = (1, Int32(poolSize.0), Int32(poolSize.1), 1)
         self.strides = (1, Int32(strides.0), Int32(strides.1), 1)
         self.padding = padding
     }
 
+    /// Returns the output obtained from applying the layer to the given input.
+    ///
+    /// - Parameters:
+    ///   - input: The input to the layer.
+    ///   - context: The contextual information for the layer application, e.g. the current learning
+    ///     phase.
+    /// - Returns: The output.
     @differentiable
     public func applied(to input: Tensor<Scalar>, in _: Context) -> Tensor<Scalar> {
-        return input.averagePooled(
-          kernelSize: poolSize, strides: strides, padding: padding)
-    }
-}
-
+        return input.averagePooled(kernelSize: poolSize, strides: strides, padding: padding)
+    }
+}
+
+/// A layer that applies layer normalization over a mini-batch of inputs.
+///
+/// Reference: [Layer Normalization](https://arxiv.org/abs/1607.06450).
 @_fixed_layout
 public struct LayerNorm<Scalar: TensorFlowFloatingPoint>: Layer {
     /// The offset value, also known as beta.

    }
}

@_fixed_layout
public struct BatchNorm<Scalar: TensorFlowFloatingPoint>: Layer {
    /// The batch dimension.

    }
}

/// A flatten layer.
///
/// A flatten layer flattens the input when applied without affecting the batch size.
@_fixed_layout
public struct Flatten<Scalar: TensorFlowFloatingPoint>: Layer {
    /// Creates a flatten layer.
    public init() {}

    /// Returns the output obtained from applying the layer to the given input.
    ///
    /// - Parameters:
    ///   - input: The input to the layer.
    ///   - context: The contextual information for the layer application, e.g. the current learning
    ///     phase.
    /// - Returns: The output.
    @differentiable
    public func applied(to input: Tensor<Scalar>, in _: Context) -> Tensor<Scalar> {
        let batchSize = input.shape[0]
