 public struct AvgPool2D<Scalar: TensorFlowFloatingPoint>: Layer {
     /// The size of the sliding reduction window for pooling.
     @noDerivative let poolSize: (Int32, Int32, Int32, Int32)
-
     /// The strides of the sliding window for each dimension of a 4-D input.
-    /// Strides in non-spatial dimensions must be 1.
+    /// Strides in non-spatial dimensions must be `1`.
     @noDerivative let strides: (Int32, Int32, Int32, Int32)
-
     /// The padding algorithm for pooling.
     @noDerivative let padding: Padding
 
-    // strides are just for the spatial dimensions (H and W)
     public init(poolSize: (Int, Int), strides: (Int, Int), padding: Padding) {
         self.poolSize = (1, Int32(poolSize.0), Int32(poolSize.1), 1)
         self.strides = (1, Int32(strides.0), Int32(strides.1), 1)

            (Dropout<Scalar>.CotangentVector, Tensor<Scalar>)) {
        switch context.learningPhase {
        case .training:
            return self.valueWithPullback(at: input) {
                $0.applyingTraining(to: $1)
            }
        case .inference:
            return self.valueWithPullback(at: input) {
                $0.applyingInference(to: $1)
            }
        }

}

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