         padding: padding.raw)
 }
 
+@usableFromInline
+func _vjpMaxPool2D<Scalar: TensorFlowFloatingPoint>(
+    _ x: Tensor<Scalar>,
+    filterSize: (Int, Int, Int, Int),
+    strides: (Int, Int, Int, Int),
+    padding: Padding
+) -> (Tensor<Scalar>, (Tensor<Scalar>) -> Tensor<Scalar>) {
+    // TODO: Currently this is not higher order differentiable. Redefine in
+    // closed form.
+    let value = maxPool2D(x, filterSize: filterSize, strides: strides, padding: padding)
+    return (value, { v in
+        Raw.maxPoolGradV2(
+            origInput: x,
+            origOutput: value,
+            grad: v,
+            ksize: Tensor<Int32>([Int32(filterSize.0), Int32(filterSize.1),
+                                  Int32(filterSize.2), Int32(filterSize.3)]),
+            strides: Tensor<Int32>([Int32(strides.0), Int32(strides.1),
+                                    Int32(strides.2), Int32(strides.3)]),
+            padding: padding.raw)
+    })
+}
+
 /// Returns a 3-D max pooling, with the specified filter sizes, strides, and
 /// padding.
 ///

        padding: padding.raw)
}

/// Returns a 3-D average pooling, with the specified filter sizes, strides,
/// and padding.
///

                  Int32(strides.4)],
        padding: padding.raw)
}

@usableFromInline
func _vjpAvgPool3D<Scalar: TensorFlowFloatingPoint>(
    _ x: Tensor<Scalar>,
    filterSize: (Int, Int, Int, Int, Int),
    strides: (Int, Int, Int, Int, Int),
    padding: Padding
) -> (Tensor<Scalar>, (Tensor<Scalar>) -> Tensor<Scalar>) {
    // TODO: Currently this is not higher order differentiable. Redefine in
    // closed form.
    let value = avgPool3D(x, filterSize: filterSize, strides: strides, padding: padding)
    return (value, { v in
        return Raw.avgPool3DGrad(
            origInputShape: x.shapeTensor,
            grad: v,
            ksize: [Int32(filterSize.0), Int32(filterSize.1), Int32(filterSize.2),
                    Int32(filterSize.3), Int32(filterSize.4)],
            strides: [Int32(strides.0), Int32(strides.1), Int32(strides.2), Int32(strides.3),
                      Int32(strides.4)],
            padding: padding.raw
        )
    })
}
