Repository: coreylowman/dfdx (language: rust)

Issue:
Add Conv2D groups
See the `groups` parameter on PyTorch's Conv2D module: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html

This turns out to be necessary to implement the efficient depthwise/pointwise convolution used by MobileNet.


Discussion / hints:
Was trying to implement MobileNet, hit this today. Seems like a popular technique for fast/small convnets.
Added to tracking issue for next release. Will be good to bundle this & the dilation option together

Fix patch:
diff --git a/src/nn/conv.rs b/src/nn/conv.rs
index 07214ec30..a88a83f32 100644
--- a/src/nn/conv.rs
+++ b/src/nn/conv.rs
@@ -13,17 +13,28 @@ pub mod builder {
         const KERNEL_SIZE: usize,
         const STRIDE: usize = 1,
         const PADDING: usize = 0,
+        const DILATION: usize = 1,
+        const GROUPS: usize = 1,
     >;
 }
 
-impl<const I: usize, const O: usize, const K: usize, const S: usize, const P: usize, E, D>
-    BuildOnDevice<D, E> for builder::Conv2D<I, O, K, S, P>
+impl<
+        const I: usize,
+        const O: usize,
+        const K: usize,
+        const S: usize,
+        const P: usize,
+        const L: usize,
+        const G: usize,
+        E,
+        D,
+    > BuildOnDevice<D, E> for builder::Conv2D<I, O, K, S, P, L, G>
 where
     E: Dtype,
     D: Device<E>,
-    Conv2D<I, O, K, S, P, E, D>: BuildModule<D, E>,
+    Conv2D<I, O, K, S, P, L, G, E, D>: BuildModule<D, E>,
 {
-    type Built = Conv2D<I, O, K, S, P, E, D>;
+    type Built = Conv2D<I, O, K, S, P, L, G, E, D>;
     fn try_build_on_device(device: &D) -> Result<Self::Built, <D>::Err> {
         Self::Built::try_build(device)
     }
@@ -45,6 +56,12 @@ where
 /// - `KERNEL_SIZE`: The size of the kernel applied to both width and height of the images.
 /// - `STRIDE`: How far to move the kernel each step. Defaults to `1`
 /// - `PADDING`: How much zero padding to add around the images. Defaults to `0`.
+/// - `DILATION`: Controls the spacing between kernel points. Defaults to `1`.
+/// - `GROUPS`: Controls the connections between inputs and outputs.
+///     `IN_CHAN` and `OUT_CHAN` must both be divisible by `GROUPS`. For example,
+///
+/// See [conv animations](https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md) for helpful
+/// visualization of all of these parameters.
 #[derive(Debug, Clone)]
 pub struct Conv2D<
     const IN_CHAN: usize,
@@ -52,19 +69,30 @@ pub struct Conv2D<
     const KERNEL_SIZE: usize,
     const STRIDE: usize,
     const PADDING: usize,
+    const DILATION: usize,
+    const GROUPS: usize,
     E: Dtype,
     D: DeviceStorage,
 > {
     pub weight: Tensor<Rank4<OUT_CHAN, IN_CHAN, KERNEL_SIZE, KERNEL_SIZE>, E, D>,
 }
 
-impl<const I: usize, const O: usize, const K: usize, const S: usize, const P: usize, E, D>
-    TensorCollection<E, D> for Conv2D<I, O, K, S, P, E, D>
+impl<
+        const I: usize,
+        const O: usize,
+        const K: usize,
+        const S: usize,
+        const P: usize,
+        const L: usize,
+        const G: usize,
+        E,
+        D,
+    > TensorCollection<E, D> for Conv2D<I, O, K, S, P, L, G, E, D>
 where
     E: Dtype + Float + SampleUniform,
     D: Device<E>,
 {
-    type To<E2: Dtype, D2: Device<E2>> = Conv2D<I, O, K, S, P, E2, D2>;
+    type To<E2: Dtype, D2: Device<E2>> = Conv2D<I, O, K, S, P, L, G, E2, D2>;
 
     fn iter_tensors<V: ModuleVisitor<Self, E, D>>(
         visitor: &mut V,
@@ -85,23 +113,52 @@ where
 }
 
 #[cfg(feature = "nightly")]
-impl<const C: usize, const O: usize, const K: usize, const S: usize, const P: usize, E, D, Img>
-    Module<Img> for Conv2D<C, O, K, S, P, E, D>
+impl<
+        const C: usize,
+        const O: usize,
+        const K: usize,
+        const S: usize,
+        const P: usize,
+        const L: usize,
+        const G: usize,
+        E,
+        D,
+        Img,
+    > Module<Img> for Conv2D<C, O, K, S, P, L, G, E, D>
 where
     E: Dtype,
     D: Device<E>,
-    Img: TryConv2DTo<Tensor<Rank4<O, C, K, K>, E, D>, S, P> + HasErr<Err = D::Err>,
+    (Img, Tensor<Rank4<O, C, K, K>, E, D>): TryConv2D<Const<S>, Const<P>, Const<L>, Const<G>>,
 {
-    type Output = Img::Output;
-    type Error = D::Err;
-
-    fn try_forward(&self, x: Img) -> Result<Self::Output, D::Err> {
-        x.try_conv2d_to(self.weight.clone())
+    type Output = <(Img, Tensor<Rank4<O, C, K, K>, E, D>) as TryConv2D<
+        Const<S>,
+        Const<P>,
+        Const<L>,
+        Const<G>,
+    >>::Convolved;
+    type Error = <(Img, Tensor<Rank4<O, C, K, K>, E, D>) as TryConv2D<
+        Const<S>,
+        Const<P>,
+        Const<L>,
+        Const<G>,
+    >>::Error;
+
+    fn try_forward(&self, x: Img) -> Result<Self::Output, Self::Error> {
+        (x, self.weight.clone()).try_conv2d(Const, Const, Const, Const)
     }
 }
 
-impl<const I: usize, const O: usize, const K: usize, const S: usize, const P: usize, E, D>
-    NonMutableModule for Conv2D<I, O, K, S, P, E, D>
+impl<
+        const I: usize,
+        const O: usize,
+        const K: usize,
+        const S: usize,
+        const P: usize,
+        const L: usize,
+        const G: usize,
+        E,
+        D,
+    > NonMutableModule for Conv2D<I, O, K, S, P, L, G, E, D>
 where
     E: Dtype,
     D: DeviceStorage,
diff --git a/src/nn/pool2d.rs b/src/nn/pool2d.rs
index ab8d5ea51..b36fe5c55 100644
--- a/src/nn/pool2d.rs
+++ b/src/nn/pool2d.rs
@@ -1,5 +1,8 @@
 #[cfg(feature = "nightly")]
-use crate::tensor_ops::{ConstAvgPool2D, ConstMaxPool2D, ConstMinPool2D};
+use crate::{
+    shapes::Const,
+    tensor_ops::{Pool2DKind, TryPool2D},
+};
 
 #[allow(unused)]
 use super::{BuildModule, Module, NonMutableModule, ZeroSizedModule};
@@ -11,8 +14,14 @@ use super::{BuildModule, Module, NonMutableModule, ZeroSizedModule};
 /// - `KERNEL_SIZE`: The size of the kernel applied to both width and height of the images.
 /// - `STRIDE`: How far to move the kernel each step. Defaults to `1`
 /// - `PADDING`: How much zero padding to add around the images. Defaults to `0`.
+/// - `DILATION` How dilated the kernel should be. Defaults to `1`.
 #[derive(Debug, Default, Clone)]
-pub struct AvgPool2D<const KERNEL_SIZE: usize, const STRIDE: usize = 1, const PADDING: usize = 0>;
+pub struct AvgPool2D<
+    const KERNEL_SIZE: usize,
+    const STRIDE: usize = 1,
+    const PADDING: usize = 0,
+    const DILATION: usize = 1,
+>;
 
 /// Max pool with 2d kernel that operates on images (3d) and batches of images (4d).
 /// Each patch reduces to the maximum value in that patch.
@@ -21,8 +30,14 @@ pub struct AvgPool2D<const KERNEL_SIZE: usize, const STRIDE: usize = 1, const PA
 /// - `KERNEL_SIZE`: The size of the kernel applied to both width and height of the images.
 /// - `STRIDE`: How far to move the kernel each step. Defaults to `1`
 /// - `PADDING`: How much zero padding to add around the images. Defaults to `0`.
+/// - `DILATION` How dilated the kernel should be. Defaults to `1`.
 #[derive(Debug, Default, Clone)]
-pub struct MaxPool2D<const KERNEL_SIZE: usize, const STRIDE: usize = 1, const PADDING: usize = 0>;
+pub struct MaxPool2D<
+    const KERNEL_SIZE: usize,
+    const STRIDE: usize = 1,
+    const PADDING: usize = 0,
+    const DILATION: usize = 1,
+>;
 
 /// Minimum pool with 2d kernel that operates on images (3d) and batches of images (4d).
 /// Each patch reduces to the minimum of the values in the patch.
@@ -31,31 +46,48 @@ pub struct MaxPool2D<const KERNEL_SIZE: usize, const STRIDE: usize = 1, const PA
 /// - `KERNEL_SIZE`: The size of the kernel applied to both width and height of the images.
 /// - `STRIDE`: How far to move the kernel each step. Defaults to `1`
 /// - `PADDING`: How much zero padding to add around the images. Defaults to `0`.
+/// - `DILATION` How dilated the kernel should be. Defaults to `1`.
 #[derive(Debug, Default, Clone)]
-pub struct MinPool2D<const KERNEL_SIZE: usize, const STRIDE: usize = 1, const PADDING: usize = 0>;
+pub struct MinPool2D<
+    const KERNEL_SIZE: usize,
+    const STRIDE: usize = 1,
+    const PADDING: usize = 0,
+    const DILATION: usize = 1,
+>;
 
 macro_rules! impl_pools {
-    ($PoolTy:tt, $Trait:ident) => {
-        impl<const K: usize, const S: usize, const P: usize> ZeroSizedModule for $PoolTy<K, S, P> {}
-        impl<const K: usize, const S: usize, const P: usize> NonMutableModule for $PoolTy<K, S, P> {}
+    ($PoolTy:tt, $Op:expr) => {
+        impl<const K: usize, const S: usize, const P: usize, const L: usize> ZeroSizedModule
+            for $PoolTy<K, S, P, L>
+        {
+        }
+        impl<const K: usize, const S: usize, const P: usize, const L: usize> NonMutableModule
+            for $PoolTy<K, S, P, L>
+        {
+        }
 
         #[cfg(feature = "nightly")]
-        impl<const K: usize, const S: usize, const P: usize, Img: $Trait<K, S, P>> Module<Img>
-            for $PoolTy<K, S, P>
+        impl<
+                const K: usize,
+                const S: usize,
+                const P: usize,
+                const L: usize,
+                Img: TryPool2D<Const<K>, Const<S>, Const<P>, Const<L>>,
+            > Module<Img> for $PoolTy<K, S, P, L>
         {
-            type Output = Img::Output;
-            type Error = Img::Err;
+            type Output = Img::Pooled;
+            type Error = Img::Error;
 
-            fn try_forward(&self, x: Img) -> Result<Self::Output, Img::Err> {
-                x.try_pool2d()
+            fn try_forward(&self, x: Img) -> Result<Self::Output, Self::Error> {
+                x.try_pool2d($Op, Const, Const, Const, Const)
             }
         }
     };
 }
 
-impl_pools!(AvgPool2D, ConstAvgPool2D);
-impl_pools!(MaxPool2D, ConstMaxPool2D);
-impl_pools!(MinPool2D, ConstMinPool2D);
+impl_pools!(AvgPool2D, Pool2DKind::Avg);
+impl_pools!(MaxPool2D, Pool2DKind::Max);
+impl_pools!(MinPool2D, Pool2DKind::Min);
 
 #[cfg(feature = "nightly")]
 #[cfg(test)]
diff --git a/src/shapes/shape.rs b/src/shapes/shape.rs
index 2adaf92b2..6df4dd4ff 100644
--- a/src/shapes/shape.rs
+++ b/src/shapes/shape.rs
@@ -164,6 +164,30 @@ where
     }
 }
 
+impl<const N: usize> core::ops::Mul<Const<N>> for usize {
+    type Output = usize;
+    fn mul(self, rhs: Const<N>) -> Self::Output {
+        self.size() * rhs.size()
+    }
+}
+impl<const N: usize> core::ops::Mul<usize> for Const<N> {
+    type Output = usize;
+    fn mul(self, rhs: usize) -> Self::Output {
+        self.size() * rhs.size()
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl<const N: usize, const M: usize> core::ops::Mul<Const<N>> for Const<M>
+where
+    Const<{ N * M }>: Sized,
+{
+    type Output = Const<{ N * M }>;
+    fn mul(self, _: Const<N>) -> Self::Output {
+        Const
+    }
+}
+
 /// Represents either `[T; N]` or `Vec<T>`
 pub trait Array<T>: IntoIterator<Item = T> {
     type Dim: Dim;
diff --git a/src/tensor_ops/conv2d/conv2d.cu b/src/tensor_ops/conv2d/conv2d.cu
index 16e2ce9dc..26ba88d23 100644
--- a/src/tensor_ops/conv2d/conv2d.cu
+++ b/src/tensor_ops/conv2d/conv2d.cu
@@ -1,9 +1,11 @@
 #include "cuda_fp16.h"
 
 struct Conv2DOp {
+    size_t kernel;
     size_t stride;
     size_t padding;
-    size_t kernel;
+    size_t dilation;
+    size_t groups;
     size_t batch;
     size_t chan_in;
     size_t chan_out;
@@ -16,12 +18,12 @@ struct Conv2DOp {
 template<typename T>
 __device__ void unfold_input_into_patches(
     const Conv2DOp op,
-    const T *image, // 4d (Batch, Channels, Height, Width)
+    const T *image, // 4d (Batch, Groups * Channels, Height, Width)
     const size_t *strides, // 4d image strides
-    T *patches // 6d (Batch, Channels, KernelSize, KernelSize, HeightOut, WidthOut)
+    T *patches // 6d (Batch, Groups * Channels, KernelSize, KernelSize, HeightOut, WidthOut)
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    if (i >= op.batch * op.chan_in * op.h_out * op.w_out) {
+    if (i >= op.batch * op.groups * op.chan_in * op.h_out * op.w_out) {
         return;
     }
 
@@ -30,21 +32,21 @@ __device__ void unfold_input_into_patches(
     idx /= op.w_out;
     const size_t oh = idx % op.h_out;
     idx /= op.h_out;
-    const size_t c = idx % op.chan_in;
-    idx /= op.chan_in;
+    const size_t c = idx % (op.chan_in * op.groups);
+    idx /= (op.chan_in * op.groups);
     const size_t b = idx % op.batch;
 
     image += b * strides[0] + c * strides[1];
     patches += oh * op.w_out + ow;
     patches += c * (op.kernel * op.kernel * op.h_out * op.w_out);
-    patches += b * (op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out);
+    patches += b * (op.groups * op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out);
 
     T zero = 0.0;
 
     for (int k1 = 0;k1 < op.kernel;k1++) {
-        const size_t y = oh * op.stride + k1 - op.padding;
+        const size_t y = oh * op.stride + op.dilation * k1 - op.padding;
         for (int k2 = 0;k2 < op.kernel;k2++) {
-            const size_t x = ow * op.stride + k2 - op.padding;
+            const size_t x = ow * op.stride + op.dilation * k2 - op.padding;
             *patches = (y >= op.h_in || x >= op.w_in) ? zero : image[y * strides[2] + x * strides[3]];
             patches += op.h_out * op.w_out;
         }
@@ -80,15 +82,15 @@ __device__ void unfold_output_into_patches(
 
     for (int k1 = 0;k1 < op.kernel;k1++) {
         const size_t oh_ks = y + op.padding;
-        const size_t oh_s = oh_ks - k1;
+        const size_t oh_s = oh_ks - op.dilation * k1;
         const size_t oh = oh_s / op.stride;
-        const bool k1_invalid = (oh_ks < k1 || oh_s % op.stride != 0 || oh >= op.h_out);
+        const bool k1_invalid = (oh_ks < op.dilation * k1 || oh_s % op.stride != 0 || oh >= op.h_out);
         for (int k2 = 0;k2 < op.kernel;k2++) {
             const size_t ow_ks = x + op.padding;
-            const size_t ow_s = ow_ks - k2;
+            const size_t ow_s = ow_ks - op.dilation * k2;
             const size_t ow = ow_s / op.stride;
         
-            const bool invalid = k1_invalid || (ow_ks < k2 || ow_s % op.stride != 0 || ow >= op.w_out);
+            const bool invalid = k1_invalid || (ow_ks < op.dilation * k2 || ow_s % op.stride != 0 || ow >= op.w_out);
             *patches = invalid ? zero : image_out[oh * op.w_out + ow];
             patches += op.h_in * op.w_in;
         }
@@ -100,31 +102,39 @@ __device__ void transpose_filters(
     const Conv2DOp op,
     const T *filters, // 4d (ChanOut, ChanIn, KernelSize, KernelSize)
     const size_t *strides, // 4d filters strides
-    T *filters_tr // 4d (ChanIn, ChanOut, KernelSize, KernelSize)
+    T *filters_tr // 5d (Groups, ChanIn, ChanOut/Groups, KernelSize, KernelSize)
 ) {
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= op.chan_in * op.chan_out * op.kernel * op.kernel) {
         return;
     }
 
+    const size_t o_per_g = op.chan_out / op.groups;
+
     unsigned int idx = i;
     const size_t k2 = idx % op.kernel;
     idx /= op.kernel;
     const size_t k1 = idx % op.kernel;
     idx /= op.kernel;
-    const size_t o = idx % op.chan_out;
-    idx /= op.chan_out;
     const size_t c = idx % op.chan_in;
+    idx /= op.chan_in;
+    const size_t o = idx % op.chan_out;
+    const size_t og = o % o_per_g;
+    const size_t g = o / o_per_g;
 
     auto i_no = o * strides[0] + c * strides[1] + k1 * strides[2] + k2 * strides[3];
-
-    filters_tr[i] = filters[i_no];
+    filters_tr += k2;
+    filters_tr += k1 * op.kernel;
+    filters_tr += og * (op.kernel * op.kernel);
+    filters_tr += c * (o_per_g * op.kernel * op.kernel);
+    filters_tr += g * (op.chan_in * o_per_g * op.kernel * op.kernel);
+    *filters_tr = filters[i_no];
 }
 
 template<typename T>
 __device__ void sum_transposed_filters(
     const Conv2DOp op,
-    const T *filters_tr, // 5d (Batch, ChanIn, ChanOut, KernelSize, KernelSize)
+    const T *filters_tr, // 6d (Batch, Groups, ChanIn, ChanOut/Groups, KernelSize, KernelSize)
     T *filters, // 4d (ChanOut, ChanIn, KernelSize, KernelSize)
     const size_t *strides // 4d filter strides
 ) {
@@ -134,6 +144,8 @@ __device__ void sum_transposed_filters(
         return;
     }
 
+    const size_t o_per_g = op.chan_out / op.groups;
+
     unsigned int idx = i;
     const size_t k2 = idx % op.kernel;
     idx /= op.kernel;
@@ -142,12 +154,17 @@ __device__ void sum_transposed_filters(
     const size_t c = idx % op.chan_in;
     idx /= op.chan_in;
     const size_t o = idx % op.chan_out;
-    idx /= op.chan_out;
+    const size_t og = o % o_per_g;
+    const size_t g = o / o_per_g;
 
     auto i_tr = c * (op.chan_out * op.kernel * op.kernel) + o * (op.kernel * op.kernel) + k1 * (op.kernel) + k2;
     auto i_no = o * strides[0] + c * strides[1] + k1 * strides[2] + k2 * strides[3];
 
-    filters_tr += i_tr;
+    filters_tr += k2;
+    filters_tr += k1 * op.kernel;
+    filters_tr += og * (op.kernel * op.kernel);
+    filters_tr += c * (o_per_g * op.kernel * op.kernel);
+    filters_tr += g * (op.chan_in * o_per_g * op.kernel * op.kernel);
 
     T tmp = 0.0;
     for (int b = 0; b < op.batch; b++) {
diff --git a/src/tensor_ops/conv2d/cpu_kernel.rs b/src/tensor_ops/conv2d/cpu_kernel.rs
index c2237ea78..30bf97ce1 100644
--- a/src/tensor_ops/conv2d/cpu_kernel.rs
+++ b/src/tensor_ops/conv2d/cpu_kernel.rs
@@ -1,6 +1,5 @@
-use crate::prelude::Tensorlike;
 use crate::shapes::{Dtype, Shape};
-use crate::tensor::{cpu::*, Tensor, ZerosTensor};
+use crate::tensor::{cpu::*, *};
 use crate::tensor_ops::matmul::cpu_kernel::MatMulImpl;
 
 use super::{Conv2DKernel, Conv2DOp};
@@ -11,10 +10,10 @@ impl Conv2DOp {
     #[inline(always)]
     fn unfold_idx(&self, [k1, k2, y, x]: [usize; 4]) -> Option<[usize; 2]> {
         let mut oh = y + self.padding;
-        if oh < k1 {
+        if oh < self.dilation * k1 {
             return None;
         }
-        oh -= k1;
+        oh -= self.dilation * k1;
         if oh % self.stride != 0 {
             return None;
         }
@@ -24,10 +23,10 @@ impl Conv2DOp {
         }
 
         let mut ow = x + self.padding;
-        if ow < k2 {
+        if ow < self.dilation * k2 {
             return None;
         }
-        ow -= k2;
+        ow -= self.dilation * k2;
         if ow % self.stride != 0 {
             return None;
         }
@@ -42,7 +41,7 @@ impl Conv2DOp {
 
 impl Cpu {
     #[inline]
-    fn conv2d_forward<E: Dtype>(
+    fn fwd<E: Dtype>(
         &self,
         op: &Conv2DOp,
         img: &[E],
@@ -55,19 +54,16 @@ impl Cpu {
     {
         {
             let mut i = 0;
-            for c in 0..op.chan_in {
+            for c in 0..(op.groups * op.chan_in) {
                 for k1 in 0..op.kernel {
                     for k2 in 0..op.kernel {
                         for oh in 0..op.h_out {
-                            let y = (oh * op.stride + k1).wrapping_sub(op.padding);
+                            let y = (oh * op.stride + op.dilation * k1).wrapping_sub(op.padding);
                             for ow in 0..op.w_out {
-                                let x = (ow * op.stride + k2).wrapping_sub(op.padding);
+                                let x =
+                                    (ow * op.stride + op.dilation * k2).wrapping_sub(op.padding);
                                 if y < op.h_in && x < op.w_in {
-                                    unsafe {
-                                        *buf.get_unchecked_mut(i) = *img.get_unchecked(
-                                            c * (op.w_in * op.h_in) + y * op.w_in + x,
-                                        )
-                                    };
+                                    buf[i] = img[c * (op.w_in * op.h_in) + y * op.w_in + x];
                                 }
                                 i += 1;
                             }
@@ -77,25 +73,27 @@ impl Cpu {
             }
         }
 
-        // (O, C * K * K) * (C * K * K, OH * OW) = (O, OH * OW)
-        let m = op.chan_out;
+        // (G, O / G, C * K * K) * (G, C * K * K, OH * OW) = (G, O / G, OH * OW)
+        let m = op.chan_out / op.groups;
         let k = op.chan_in * op.kernel * op.kernel;
         let n = op.w_out * op.h_out;
-        Self::matmul(
-            (m, k, n),
-            filters.as_ptr(),
-            [k, 1],
-            buf.as_ptr(),
-            [n, 1],
-            out.as_mut_ptr(),
-            [n, 1],
-        );
+        for g in 0..op.groups {
+            Self::matmul(
+                (m, k, n),
+                filters[g * m * k..].as_ptr(),
+                [k, 1],
+                buf[g * k * n..].as_ptr(),
+                [n, 1],
+                out[g * m * n..].as_mut_ptr(),
+                [n, 1],
+            );
+        }
         Ok(())
     }
 
     #[inline]
     #[allow(clippy::too_many_arguments)]
-    fn conv2d_backward<E: Dtype>(
+    fn bwd<E: Dtype>(
         &self,
         op: &Conv2DOp,
         img: &[E],
@@ -116,11 +114,8 @@ impl Cpu {
                         for y in 0..op.h_in {
                             for x in 0..op.w_in {
                                 if let Some([oh, ow]) = op.unfold_idx([k1, k2, y, x]) {
-                                    unsafe {
-                                        *buf.get_unchecked_mut(i) = *grad_out.get_unchecked(
-                                            o * (op.h_out * op.w_out) + oh * op.w_out + ow,
-                                        );
-                                    }
+                                    buf[i] =
+                                        grad_out[o * (op.h_out * op.w_out) + oh * op.w_out + ow];
                                 }
                                 i += 1;
                             }
@@ -132,36 +127,40 @@ impl Cpu {
 
         {
             // img_g += filters^T * unfold(grad_out)
-            // (C, H * W) += (C, O * K * K) * (O * K * K, H * W)
+            // (G, C, H * W) += (G, C, O/G * K * K) * (G, O/G * K * K, H * W)
             let m = op.chan_in;
-            let k = op.chan_out * op.kernel * op.kernel;
+            let k = (op.chan_out / op.groups) * op.kernel * op.kernel;
             let n = op.h_in * op.w_in;
-            Self::matmul(
-                (m, k, n),
-                filters_tr.as_ptr(),
-                [k, 1],
-                buf.as_ptr(),
-                [n, 1],
-                grad_img.as_mut_ptr(),
-                [n, 1],
-            );
+            for g in 0..op.groups {
+                Self::matmul(
+                    (m, k, n),
+                    filters_tr[g * m * k..].as_ptr(),
+                    [k, 1],
+                    buf[g * k * n..].as_ptr(),
+                    [n, 1],
+                    grad_img[g * m * n..].as_mut_ptr(),
+                    [n, 1],
+                );
+            }
         }
 
         {
-            // weight_g^T += img * patches^T
-            // (C, O * K * K) += (C, H * W) * (H * W, O * K * K)
+            // weight_g^T += img * unfold(patches)^T
+            // (G, C, O/G * K * K) += (G, C, H * W) * (G, H * W, O/G * K * K)
             let m = op.chan_in;
             let k = op.h_in * op.w_in;
-            let n = op.chan_out * op.kernel * op.kernel;
-            Self::matmul(
-                (m, k, n),
-                img.as_ptr(),
-                [k, 1],
-                buf.as_ptr(),
-                [1, k],
-                grad_filters_tr.as_mut_ptr(),
-                [n, 1],
-            );
+            let n = (op.chan_out / op.groups) * op.kernel * op.kernel;
+            for g in 0..op.groups {
+                Self::matmul(
+                    (m, k, n),
+                    img[g * m * k..].as_ptr(),
+                    [k, 1],
+                    buf[g * k * n..].as_ptr(),
+                    [1, k],
+                    grad_filters_tr[g * m * n..].as_mut_ptr(),
+                    [n, 1],
+                );
+            }
         }
         Ok(())
     }
@@ -182,7 +181,14 @@ where
         rhs: &Tensor<R, E, Self>,
         out: &mut Tensor<O, E, Self>,
     ) -> Result<(), Self::Err> {
-        let mut patches = self.try_alloc_zeros::<E>(op.inp_patches_shape().num_elements())?;
+        let patches = (
+            op.groups * op.chan_in,
+            op.kernel,
+            op.kernel,
+            op.h_out,
+            op.w_out,
+        );
+        let mut patches = self.try_alloc_zeros::<E>(patches.num_elements())?;
         let [lstride, ostride] = match L::NUM_DIMS {
             3 => [0; 2],
             4 => [lhs.strides[0], out.strides[0]],
@@ -192,7 +198,7 @@ where
         let rhs = rhs.data.as_ref();
         let out = Arc::make_mut(&mut out.data);
         for i_batch in 0..op.batch {
-            self.conv2d_forward(
+            self.fwd(
                 &op,
                 &lhs[i_batch * lstride..],
                 rhs,
@@ -213,8 +219,15 @@ where
         out: &impl Tensorlike<O, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
-        let f_tr_shape = op.filters_tr_shape();
-        let mut patches = self.try_alloc_zeros::<E>(op.out_patches_shape().num_elements())?;
+        let f_tr_shape = [
+            op.groups,
+            op.chan_in,
+            op.chan_out / op.groups,
+            op.kernel,
+            op.kernel,
+        ];
+        let patches_shape = [op.chan_out, op.kernel, op.kernel, op.h_in, op.w_in];
+        let mut patches = self.try_alloc_zeros::<E>(patches_shape.num_elements())?;
         let mut f1023 = self.try_alloc_zeros::<E>(f_tr_shape.num_elements())?;
         let mut grad_f1023 = self.try_alloc_zeros::<E>(f_tr_shape.num_elements())?;
 
@@ -222,8 +235,8 @@ where
             // transpose filters in f1023
             let buf = rhs.data.as_ref();
             let mut f_idx = NdIndex::new(f_tr_shape, f_tr_shape.strides());
-            while let Some((i, [c, o, k1, k2])) = f_idx.next_with_idx() {
-                let idx = o * rhs.strides[0]
+            while let Some((i, [g, c, o, k1, k2])) = f_idx.next_with_idx() {
+                let idx = (g * (op.chan_out / op.groups) + o) * rhs.strides[0]
                     + c * rhs.strides[1]
                     + k1 * rhs.strides[2]
                     + k2 * rhs.strides[3];
@@ -239,7 +252,7 @@ where
         let lhs = lhs.data.as_ref();
 
         for i_batch in 0..op.batch {
-            self.conv2d_backward(
+            self.bwd(
                 &op,
                 &lhs[i_batch * lstride..],
                 &mut grad_lhs[i_batch * lstride..],
@@ -253,8 +266,8 @@ where
         {
             // untranspose filters
             let mut f_idx = NdIndex::new(f_tr_shape, f_tr_shape.strides());
-            while let Some((i, [c, o, k1, k2])) = f_idx.next_with_idx() {
-                let idx = o * rhs.strides[0]
+            while let Some((i, [g, c, o, k1, k2])) = f_idx.next_with_idx() {
+                let idx = (g * (op.chan_out / op.groups) + o) * rhs.strides[0]
                     + c * rhs.strides[1]
                     + k1 * rhs.strides[2]
                     + k2 * rhs.strides[3];
diff --git a/src/tensor_ops/conv2d/cuda_kernel.rs b/src/tensor_ops/conv2d/cuda_kernel.rs
index d773b3e5e..8dbe25f3b 100644
--- a/src/tensor_ops/conv2d/cuda_kernel.rs
+++ b/src/tensor_ops/conv2d/cuda_kernel.rs
@@ -68,42 +68,66 @@ where
     fn forward<L: Shape, R: Shape, O: Shape>(
         &self,
         op: super::Conv2DOp,
-        lhs: &Tensor<L, E, Self>,
-        rhs: &Tensor<R, E, Self>,
+        img: &Tensor<L, E, Self>,
+        fil: &Tensor<R, E, Self>,
         out: &mut Tensor<O, E, Self>,
     ) -> Result<(), Self::Err> {
         if !self.dev.has_func(Self::MOD, Self::FNS[0]) {
             self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?;
         }
 
-        let patches_item_numel = op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out;
+        let patches_item_numel =
+            op.groups * op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out;
         let patches_numel = op.batch * patches_item_numel;
 
         let mut patches = unsafe { self.get_workspace::<E>(patches_numel) }?;
         let mut patches = unsafe { patches.transmute_mut::<E>(patches_numel).unwrap() };
 
-        let img_strides = self.dev.htod_copy(make_4d::<L>(lhs.strides).into())?;
-        let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
-        let cfg = launch_cfg::<128>((op.batch * op.chan_in * op.h_out * op.w_out) as u32);
-        let params = (op, lhs.data.as_ref(), &img_strides, &mut patches);
-        unsafe { unfold_fn.launch(cfg, params) }?;
+        let img_strides = self.dev.htod_copy(make_4d::<L>(img.strides).into())?;
+
+        let out_buf = Arc::get_mut(&mut out.data).unwrap();
 
-        // (O, C * K * K) * (B, C * K * K, OH * OW) = (B, O, OH * OW)
-        let m = op.chan_out;
-        let k = op.chan_in * op.kernel * op.kernel;
-        let n = op.h_out * op.w_out;
         unsafe {
-            self.gemm_batch(
-                (op.batch, m, k, n),
-                rhs.data.as_ref(),
-                [0, k, 1],
-                &patches,
-                [k * n, n, 1],
-                Default::default(),
-                Arc::get_mut(&mut out.data).unwrap(),
-                [m * n, n, 1],
-            )
-            .unwrap();
+            let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[0]).unwrap();
+            let cfg =
+                launch_cfg::<128>((op.batch * op.groups * op.chan_in * op.h_out * op.w_out) as u32);
+            let params = (op, img.data.as_ref(), &img_strides, &mut patches);
+            unfold_fn.launch(cfg, params)?;
+
+            // LHS    (G, O/G, C*K*K)
+            // RHS (B, G, C*K*K, OH*OW)
+            // OUT (B, G, O/G, OH*OW)
+            let m = op.chan_out / op.groups;
+            let k = op.chan_in * op.kernel * op.kernel;
+            let n = op.h_out * op.w_out;
+            if op.groups == 1 {
+                // optimizing here for common case
+                self.gemm_batch(
+                    (op.batch, m, k, n),
+                    fil.data.as_ref(),
+                    [0, k, 1],
+                    &patches,
+                    [k * n, n, 1],
+                    Default::default(),
+                    out_buf,
+                    [m * n, n, 1],
+                )
+                .unwrap();
+            } else {
+                for i_batch in 0..op.batch {
+                    self.gemm_batch(
+                        (op.groups, m, k, n),
+                        fil.data.as_ref(),
+                        [m * k, k, 1],
+                        &patches.slice(i_batch * op.groups * k * n..),
+                        [k * n, n, 1],
+                        Default::default(),
+                        &mut out_buf.slice_mut(i_batch * op.groups * m * n..),
+                        [m * n, n, 1],
+                    )
+                    .unwrap();
+                }
+            }
         }
 
         Ok(())
@@ -126,44 +150,45 @@ where
         let mut patches = unsafe { self.get_workspace::<E>(patches_numel) }?;
         let mut patches = unsafe { patches.transmute_mut::<E>(patches_numel).unwrap() };
 
-        let mut f_b1023 = unsafe { self.alloc_empty::<E>(filters_numel) }?;
-        let mut grad_f_b1023 = unsafe { self.alloc_empty::<E>(op.batch * filters_numel) }?;
+        let mut ftr = unsafe { self.alloc_empty::<E>(filters_numel) }?;
+        let mut grad_ftr = unsafe { self.alloc_empty::<E>(op.batch * filters_numel) }?;
         let f_strides = self.dev.htod_copy(rhs.strides.into())?;
 
         self.par_stream.wait_for_default()?;
 
-        {
+        unsafe {
             // unfold grad_out into patches
             let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
             let cfg = launch_cfg::<128>((op.batch * op.chan_out * op.h_in * op.w_in) as u32);
-            unsafe { unfold_fn.launch(cfg, (op, grad_out, &mut patches)) }?;
+            unfold_fn.launch(cfg, (op, grad_out, &mut patches))?;
         }
 
-        {
+        unsafe {
             // prepare filters for backward operations by
             // swapping dims 0 and 1
             let tr_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
             let cfg = launch_cfg::<128>(rhs.shape.num_elements() as u32);
-            unsafe {
-                tr_fn.launch_on_stream(
-                    self.par_stream.as_ref(),
-                    cfg,
-                    (op, rhs.data.as_ref(), &f_strides, &mut f_b1023),
-                )
-            }?;
+            tr_fn.launch_on_stream(
+                self.par_stream.as_ref(),
+                cfg,
+                (op, rhs.data.as_ref(), &f_strides, &mut ftr),
+            )?;
 
             self.par_stream.wait_for_default()?;
 
             // img_g += filters * patches
-            // (B, C, H * W) += (B, C, O * K * K) * (B, O * K * K, H * W)
+            // LHS =    (G, C, O/G*K*K)
+            // RHS = (B, G, O/G*K*K, H*W)
+            // OUT = (B, G, C, H*W)
             let m = op.chan_in;
-            let k = op.chan_out * op.kernel * op.kernel;
+            let k = (op.chan_out / op.groups) * op.kernel * op.kernel;
             let n = op.h_in * op.w_in;
-            unsafe {
-                self.blas.set_stream(Some(self.par_stream.as_ref()))?;
+            self.blas.set_stream(Some(self.par_stream.as_ref()))?;
+            if op.groups == 1 {
+                // optimizing here for common case
                 self.gemm_batch(
                     (op.batch, m, k, n),
-                    &f_b1023,
+                    &ftr,
                     [0, k, 1],
                     &patches,
                     [k * n, n, 1],
@@ -172,17 +197,34 @@ where
                     [m * n, n, 1],
                 )
                 .unwrap();
-                self.blas.set_stream(None)?;
+            } else {
+                for i_batch in 0..op.batch {
+                    self.gemm_batch(
+                        (op.groups, m, k, n),
+                        &ftr,
+                        [m * k, k, 1],
+                        &patches.slice(i_batch * op.groups * k * n..),
+                        [k * n, n, 1],
+                        <E>::ONE,
+                        &mut grad_lhs.slice_mut(i_batch * op.groups * m * n..),
+                        [m * n, n, 1],
+                    )
+                    .unwrap();
+                }
             }
+            self.blas.set_stream(None)?;
         }
 
-        {
+        unsafe {
             // weight_g += img * patches^T
-            // (B, C, O * K * K) += (B, C, H * W) * (B, H * W, O * K * K)
+            // LHS = (B, G, C, H*W)
+            // RHS = (B, H*W, G, O/G*K*K)
+            // OUT = (B, G, C, O/G*K*K)
             let m = op.chan_in;
             let k = op.h_in * op.w_in;
-            let n = op.chan_out * op.kernel * op.kernel;
-            unsafe {
+            let n = (op.chan_out / op.groups) * op.kernel * op.kernel;
+            if op.groups == 1 {
+                // optimizing here for common case
                 self.gemm_batch(
                     (op.batch, m, k, n),
                     lhs.data.as_ref(),
@@ -190,17 +232,32 @@ where
                     &patches,
                     [k * n, 1, k],
                     Default::default(),
-                    &mut grad_f_b1023,
+                    &mut grad_ftr,
                     [m * n, n, 1],
                 )
                 .unwrap();
+            } else {
+                let lhs_buf = lhs.data.as_ref();
+                for i_batch in 0..op.batch {
+                    self.gemm_batch(
+                        (op.groups, m, k, n),
+                        &lhs_buf.slice(i_batch * op.groups * m * k..),
+                        [m * k, k, 1],
+                        &patches.slice(i_batch * op.groups * k * n..),
+                        [k * n, 1, k],
+                        Default::default(),
+                        &mut grad_ftr.slice_mut(i_batch * op.groups * m * n..),
+                        [m * n, n, 1],
+                    )
+                    .unwrap();
+                }
             }
 
             // sum all the gradients collected in our broadcasted grad_f
             // into grad_rhs
             let sum_fn = self.dev.get_func(Self::MOD, Self::FNS[3]).unwrap();
             let cfg = launch_cfg::<128>(rhs.shape.num_elements() as u32);
-            unsafe { sum_fn.launch(cfg, (op, &grad_f_b1023, grad_rhs, &f_strides)) }?;
+            sum_fn.launch(cfg, (op, &grad_ftr, grad_rhs, &f_strides))?;
         }
 
         self.dev.wait_for(self.par_stream.as_ref())?;
diff --git a/src/tensor_ops/conv2d/cudnn_kernel.rs b/src/tensor_ops/conv2d/cudnn_kernel.rs
index 1f2792eca..2e0a1429d 100644
--- a/src/tensor_ops/conv2d/cudnn_kernel.rs
+++ b/src/tensor_ops/conv2d/cudnn_kernel.rs
@@ -37,12 +37,13 @@ where
         rhs: &Tensor<R, E, Self>,
         out: &mut Tensor<O, E, Self>,
     ) -> Result<(), Self::Err> {
-        let conv = self.cudnn.create_conv2d::<E>(
+        let mut conv = self.cudnn.create_conv2d::<E>(
             [op.padding as i32, op.padding as i32],
             [op.stride as i32, op.stride as i32],
-            [1, 1],
+            [op.dilation as i32, op.dilation as i32],
             cudnn::sys::cudnnConvolutionMode_t::CUDNN_CROSS_CORRELATION,
         )?;
+        conv.set_group_count(op.groups as i32)?;
         let img = self.cudnn.create_4d_tensor_ex::<E>(
             make_4d::<L>(lhs.shape.concrete(), 1).map(|x| x as i32),
             make_4d::<L>(lhs.strides, 0).map(|x| x as i32),
@@ -94,12 +95,13 @@ where
         out: &impl Tensorlike<O, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
-        let conv = self.cudnn.create_conv2d::<E>(
+        let mut conv = self.cudnn.create_conv2d::<E>(
             [op.padding as i32, op.padding as i32],
             [op.stride as i32, op.stride as i32],
-            [1, 1],
+            [op.dilation as i32, op.dilation as i32],
             cudnn::sys::cudnnConvolutionMode_t::CUDNN_CROSS_CORRELATION,
         )?;
+        conv.set_group_count(op.groups as i32)?;
         let img = self.cudnn.create_4d_tensor_ex::<E>(
             make_4d::<L>(lhs.shape.concrete(), 1).map(|x| x as i32),
             make_4d::<L>(lhs.strides, 0).map(|x| x as i32),
diff --git a/src/tensor_ops/conv2d/mod.rs b/src/tensor_ops/conv2d/mod.rs
index 593eccb52..c1d7dd28d 100644
--- a/src/tensor_ops/conv2d/mod.rs
+++ b/src/tensor_ops/conv2d/mod.rs
@@ -1,3 +1,5 @@
+use crate::{shapes::*, tensor::*, tensor_ops::ReshapeTo};
+
 mod cpu_kernel;
 
 #[cfg(all(not(feature = "cudnn"), feature = "cuda"))]
@@ -6,14 +8,17 @@ mod cuda_kernel;
 #[cfg(feature = "cudnn")]
 mod cudnn_kernel;
 
-use crate::{shapes::*, tensor::*};
+#[cfg(test)]
+mod tests;
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub(super) struct Conv2DOp {
+    pub kernel: usize,
     pub stride: usize,
     pub padding: usize,
-    pub kernel: usize,
+    pub dilation: usize,
+    pub groups: usize,
     pub batch: usize,
     pub chan_in: usize,
     pub chan_out: usize,
@@ -23,37 +28,6 @@ pub(super) struct Conv2DOp {
     pub w_out: usize,
 }
 
-impl Conv2DOp {
-    fn new(s: usize, p: usize, k: usize, [b, c, h_in, w_in]: [usize; 4], o: usize) -> Self {
-        Self {
-            stride: s,
-            padding: p,
-            kernel: k,
-            batch: b,
-            chan_in: c,
-            chan_out: o,
-            h_in,
-            h_out: (h_in + 2 * p - k) / s + 1,
-            w_in,
-            w_out: (w_in + 2 * p - k) / s + 1,
-        }
-    }
-
-    #[rustfmt::skip]
-    pub(super) fn inp_patches_shape(&self) -> (usize, usize, usize, usize, usize) {
-        (self.chan_in, self.kernel, self.kernel, self.h_out, self.w_out)
-    }
-
-    #[rustfmt::skip]
-    pub(super) fn out_patches_shape(&self) -> (usize, usize, usize, usize, usize) {
-        (self.chan_out, self.kernel, self.kernel, self.h_in, self.w_in)
-    }
-
-    pub(super) fn filters_tr_shape(&self) -> (usize, usize, usize, usize) {
-        (self.chan_in, self.chan_out, self.kernel, self.kernel)
-    }
-}
-
 pub(super) trait Conv2DKernel<E: Dtype>: DeviceStorage {
     fn alloc<S: Shape>(&self, s: S) -> Result<Tensor<S, E, Self>, Self::Err>;
 
@@ -78,135 +52,204 @@ pub(super) trait Conv2DKernel<E: Dtype>: DeviceStorage {
     ) -> Result<(), Self::Err>;
 }
 
-pub trait ConvAlgebra<const K: usize, const S: usize, const P: usize>: Dim {
-    type Convolved: Dim;
+pub trait TryConv2D<Stride, Padding, Dilation, Groups>: Sized {
+    type Convolved;
+    type Error: std::fmt::Debug;
+
+    fn conv2d(
+        self,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+        groups: Groups,
+    ) -> Self::Convolved {
+        self.try_conv2d(stride, padding, dilation, groups).unwrap()
+    }
 
-    fn convolve_dim(&self) -> Self::Convolved;
+    fn try_conv2d(
+        self,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+        groups: Groups,
+    ) -> Result<Self::Convolved, Self::Error>;
 }
 
-impl<const D: usize, const K: usize, const S: usize, const P: usize> ConvAlgebra<K, S, P>
-    for Const<D>
+impl<
+        const KERNEL: usize,
+        const STRIDE: usize,
+        const PADDING: usize,
+        const DILATION: usize,
+        Groups: Dim,
+        const DIM: usize,
+    > TryConv2D<Const<STRIDE>, Const<PADDING>, Const<DILATION>, Groups>
+    for (Const<DIM>, Const<KERNEL>)
 where
-    Const<{ (D + 2 * P - K) / S + 1 }>: Sized,
+    Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>: Sized,
 {
-    type Convolved = Const<{ (D + 2 * P - K) / S + 1 }>;
-
-    fn convolve_dim(&self) -> Self::Convolved {
-        Self::Convolved::default()
+    type Convolved = Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>;
+    type Error = std::convert::Infallible;
+    fn try_conv2d(
+        self,
+        _: Const<STRIDE>,
+        _: Const<PADDING>,
+        _: Const<DILATION>,
+        _: Groups,
+    ) -> Result<Self::Convolved, Self::Error> {
+        Ok(Const)
     }
 }
 
-impl<const K: usize, const S: usize, const P: usize> ConvAlgebra<K, S, P> for usize {
+impl<Kernel: Dim, Stride: Dim, Padding: Dim, Dilation: Dim, Groups: Dim>
+    TryConv2D<Stride, Padding, Dilation, Groups> for (usize, Kernel)
+{
     type Convolved = usize;
-
-    fn convolve_dim(&self) -> Self::Convolved {
-        (self.size() + 2 * P).checked_sub(K).unwrap() / S + 1
-    }
-}
-
-pub trait TryConv2DTo<F, const S: usize, const P: usize>: HasErr {
-    type Output;
-    fn conv2d_to(self, filters: F) -> Self::Output {
-        self.try_conv2d_to(filters).unwrap()
-    }
-    fn try_conv2d_to(self, filters: F) -> Result<Self::Output, Self::Err>;
-}
-
-pub trait TryConv2D<F> {
-    fn conv2d<const S: usize, const P: usize>(self, filters: F) -> Self::Output
-    where
-        Self: TryConv2DTo<F, S, P>,
-    {
-        self.conv2d_to(filters)
-    }
-    fn try_conv2d<const S: usize, const P: usize>(
+    type Error = std::convert::Infallible;
+    fn try_conv2d(
         self,
-        filters: F,
-    ) -> Result<Self::Output, Self::Err>
-    where
-        Self: TryConv2DTo<F, S, P>,
-    {
-        self.try_conv2d_to(filters)
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+        _: Groups,
+    ) -> Result<Self::Convolved, Self::Error> {
+        let (dim, kernel) = self;
+        Ok((dim + 2 * padding.size() - 1)
+            .checked_sub(dilation.size() * (kernel.size() - 1))
+            .unwrap()
+            / stride.size()
+            + 1)
     }
 }
 
-impl<S: Shape, E: Dtype, D: DeviceStorage, T, F> TryConv2D<F> for Tensor<S, E, D, T> {}
-
-impl<
-        const C: usize,
-        H: Dim + ConvAlgebra<K, S, P>,
-        W: Dim + ConvAlgebra<K, S, P>,
-        const O: usize,
-        const K: usize,
-        const S: usize,
-        const P: usize,
-        E: Dtype,
-        D: Conv2DKernel<E> + ZerosTensor<E>,
-        T: 'static + Tape<E, D>,
-    > TryConv2DTo<Tensor<Rank4<O, C, K, K>, E, D>, S, P> for Tensor<(Const<C>, H, W), E, D, T>
+impl<InpChan, OutChan, Kernel, Stride, Padding, Dilation, Groups, H, W, E, D, T>
+    TryConv2D<Stride, Padding, Dilation, Groups>
+    for (
+        Tensor<(<InpChan as std::ops::Mul<Groups>>::Output, H, W), E, D, T>,
+        Tensor<(OutChan, InpChan, Kernel, Kernel), E, D>,
+    )
+where
+    InpChan: Dim,
+    OutChan: Dim,
+    Kernel: Dim,
+    Stride: Dim,
+    Padding: Dim,
+    Dilation: Dim,
+    Groups: Dim,
+    H: Dim,
+    W: Dim,
+    E: Dtype,
+    D: Conv2DKernel<E> + crate::tensor_ops::reshape_to::ReshapeKernel<E>,
+    T: Tape<E, D>,
+    InpChan: std::ops::Mul<Groups>,
+    <InpChan as std::ops::Mul<Groups>>::Output: Dim,
+    (H, Kernel): TryConv2D<Stride, Padding, Dilation, Groups>,
+    (W, Kernel): TryConv2D<Stride, Padding, Dilation, Groups>,
+    <(H, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved: Dim,
+    <(W, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved: Dim,
 {
-    type Output = Tensor<(Const<O>, H::Convolved, W::Convolved), E, D, T>;
-
-    fn try_conv2d_to(
+    type Convolved = Tensor<
+        (
+            OutChan,
+            <(H, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved,
+            <(W, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved,
+        ),
+        E,
+        D,
+        T,
+    >;
+    type Error = D::Err;
+
+    fn try_conv2d(
         self,
-        filters: Tensor<Rank4<O, C, K, K>, E, D>,
-    ) -> Result<Self::Output, Self::Err> {
-        let h = self.shape.1;
-        let w = self.shape.2;
-
-        let op = Conv2DOp::new(S, P, K, [1, C, h.size(), w.size()], O);
-        let (lhs, ltape) = self.split_tape();
-        let (rhs, rtape) = filters.split_tape();
-        let mut tape = ltape.merge(rtape);
-        let mut out = lhs
-            .device
-            .alloc((Const, h.convolve_dim(), w.convolve_dim()))?;
-        lhs.device.forward(op, &lhs, &rhs, &mut out)?;
-        let lhs_ghost = lhs.ghost();
-        let rhs_ghost = rhs.ghost();
-        let out_ghost = out.ghost();
-        tape.add_backward_op(move |grads| {
-            grads.try_alloc_for(&rhs_ghost)?;
-            grads.try_alloc_for(&lhs_ghost)?;
-            grads.try_alloc_for(&out_ghost)?;
-            let (grad_lhs, grad_rhs, grad_out) =
-                grads.muts_and_ref(&lhs_ghost, &rhs_ghost, &out_ghost);
-            lhs.device
-                .backward(op, &lhs, grad_lhs, &rhs, grad_rhs, &out_ghost, grad_out)
-        });
-        Ok(out.put_tape(tape))
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+        groups: Groups,
+    ) -> Result<Self::Convolved, Self::Error> {
+        let (img, filters) = self;
+        let (inp_chan, h, w) = img.shape;
+        let img = img.try_reshape_like(&(Const::<1>, inp_chan, h, w))?;
+        let out = (img, filters).try_conv2d(stride, padding, dilation, groups)?;
+        let (_, out_chan, out_h, out_w) = out.shape;
+        out.try_reshape_like(&(out_chan, out_h, out_w))
     }
 }
 
-impl<
-        B: Dim,
-        const C: usize,
-        H: Dim + ConvAlgebra<K, S, P>,
-        W: Dim + ConvAlgebra<K, S, P>,
-        const O: usize,
-        const K: usize,
-        const S: usize,
-        const P: usize,
-        E: Dtype,
-        D: Conv2DKernel<E> + ZerosTensor<E>,
-        T: 'static + Tape<E, D>,
-    > TryConv2DTo<Tensor<Rank4<O, C, K, K>, E, D>, S, P> for Tensor<(B, Const<C>, H, W), E, D, T>
+impl<InpChan, OutChan, Kernel, Stride, Padding, Dilation, Groups, Batch, H, W, E, D, T>
+    TryConv2D<Stride, Padding, Dilation, Groups>
+    for (
+        Tensor<(Batch, <InpChan as std::ops::Mul<Groups>>::Output, H, W), E, D, T>,
+        Tensor<(OutChan, InpChan, Kernel, Kernel), E, D>,
+    )
+where
+    InpChan: Dim,
+    OutChan: Dim,
+    Kernel: Dim,
+    Stride: Dim,
+    Padding: Dim,
+    Dilation: Dim,
+    Groups: Dim,
+    Batch: Dim,
+    H: Dim,
+    W: Dim,
+    E: Dtype,
+    D: Conv2DKernel<E>,
+    T: Tape<E, D>,
+    InpChan: std::ops::Mul<Groups>,
+    <InpChan as std::ops::Mul<Groups>>::Output: Dim,
+    (H, Kernel): TryConv2D<Stride, Padding, Dilation, Groups>,
+    (W, Kernel): TryConv2D<Stride, Padding, Dilation, Groups>,
+    <(H, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved: Dim,
+    <(W, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved: Dim,
 {
-    type Output = Tensor<(B, Const<O>, H::Convolved, W::Convolved), E, D, T>;
-    fn try_conv2d_to(
+    type Convolved = Tensor<
+        (
+            Batch,
+            OutChan,
+            <(H, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved,
+            <(W, Kernel) as TryConv2D<Stride, Padding, Dilation, Groups>>::Convolved,
+        ),
+        E,
+        D,
+        T,
+    >;
+    type Error = D::Err;
+
+    fn try_conv2d(
         self,
-        filters: Tensor<Rank4<O, C, K, K>, E, D>,
-    ) -> Result<Self::Output, Self::Err> {
-        let batch = self.shape().0;
-        let h = self.shape().2;
-        let w = self.shape().3;
-
-        let op = Conv2DOp::new(S, P, K, [batch.size(), C, h.size(), w.size()], O);
-        let (lhs, ltape) = self.split_tape();
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+        groups: Groups,
+    ) -> Result<Self::Convolved, Self::Error> {
+        let (img, filters) = self;
+        assert_eq!(img.shape.1.size(), filters.shape.1.size() * groups.size());
+        assert_eq!(filters.shape.2, filters.shape.3);
+        let (batch, _, h, w) = img.shape;
+        let (out_chan, inp_chan, kernel, _) = filters.shape;
+        assert!(out_chan.size() % groups.size() == 0);
+        if img.strides != img.shape.strides() || filters.strides != filters.shape.strides() {
+            panic!("Image & filter inputs to conv2d must be contiguous");
+        }
+        let h_out = (h, kernel).conv2d(stride, padding, dilation, groups);
+        let w_out = (w, kernel).conv2d(stride, padding, dilation, groups);
+        let op = Conv2DOp {
+            stride: stride.size(),
+            padding: padding.size(),
+            kernel: kernel.size(),
+            dilation: dilation.size(),
+            groups: groups.size(),
+            batch: batch.size(),
+            chan_in: inp_chan.size(),
+            chan_out: out_chan.size(),
+            h_in: h.size(),
+            h_out: h_out.size(),
+            w_in: w.size(),
+            w_out: w_out.size(),
+        };
+        let (lhs, ltape) = img.split_tape();
         let (rhs, rtape) = filters.split_tape();
-        let mut out = lhs
-            .device
-            .alloc((batch, Const, h.convolve_dim(), w.convolve_dim()))?;
+        let mut out = lhs.device.alloc((batch, out_chan, h_out, w_out))?;
         let mut tape = ltape.merge(rtape);
         lhs.device.forward(op, &lhs, &rhs, &mut out)?;
         let lhs_ghost = lhs.ghost();
@@ -224,250 +267,3 @@ impl<
         Ok(out.put_tape(tape))
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::{tensor_ops::*, tests::*};
-
-    #[test]
-    /// Produced by
-    /// ```python
-    /// q = torch.nn.Conv2d(1, 2, 2)
-    /// x = torch.sample_normal(1, 2, 3, requires_grad=True)
-    /// q(x).exp().mean().backward()
-    /// ```
-    fn test_conv2d_default_stride_and_padding() {
-        let dev: TestDevice = Default::default();
-        let weight = dev
-            .tensor([
-                [[[-0.04958433, -0.43007267], [0.01935136, 0.09778714]]],
-                [[[0.44083858, -0.20507240], [-0.30017477, -0.10937047]]],
-            ])
-            .to_dtype::<TestDtype>();
-        let bias = dev
-            .tensor([0.36406237, -0.30981010])
-            .to_dtype::<TestDtype>();
-        let x = dev
-            .tensor([[
-                [-0.86713916, 0.52773184, -0.95238322],
-                [-0.64531374, 0.77809018, -0.49099201],
-            ]])
-            .to_dtype::<TestDtype>();
-        let result = x.leaky_trace().conv2d::<1, 0>(weight.clone())
-            + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
-        assert_close_to_literal!(
-            result,
-            [[[0.24369538, 0.71453357]], [[-0.69169492, -0.06172103]]]
-        );
-        let g = result.exp().mean().backward();
-        assert_close_to_literal!(
-            g.get(&x),
-            [[
-                [0.03936806, -0.08457474, -0.26788417],
-                [-0.03140351, -0.04316529, 0.02424446],
-            ]]
-        );
-        assert_close_to_literal!(
-            g.get(&weight),
-            [
-                [[[-0.00703794, -0.31814471], [0.19160703, -0.00260070]]],
-                [[[0.01548620, -0.15778227], [0.10209797, -0.01799832]]],
-            ]
-        );
-        assert_close_to_literal!(g.get(&bias), [0.82979727, 0.36021793]);
-    }
-
-    #[test]
-    /// Produced by
-    /// ```python
-    /// q = torch.nn.Conv2d(1, 2, 2, stride=2)
-    /// x = torch.sample_normal(1, 2, 3, requires_grad=True)
-    /// q(x).exp().mean().backward()
-    /// ```
-    fn test_conv2d_stride_2() {
-        let dev: TestDevice = Default::default();
-        let weight = dev
-            .tensor([
-                [[[0.44704646, -0.29563826], [0.29228759, -0.16575140]]],
-                [[[-0.30488998, 0.25222939], [0.13279295, 0.38153177]]],
-            ])
-            .to_dtype::<TestDtype>();
-        let bias = dev
-            .tensor([-0.44699109, 0.38371694])
-            .to_dtype::<TestDtype>();
-        let x = dev
-            .tensor([[
-                [0.37100124, -0.59504986, -1.19781005],
-                [-0.31547278, 0.58071911, 0.86612970],
-            ]])
-            .to_dtype::<TestDtype>();
-
-        let result = x.leaky_trace().conv2d::<2, 0>(weight.clone())
-            + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
-        assert_close_to_literal!(result, [[[-0.29368058]], [[0.30018353]]]);
-
-        let g = result.exp().mean().backward();
-
-        assert_close_to_literal!(
-            g.get(&x),
-            [[[-0.03917716, 0.06006697, 0.], [0.19859464, 0.19576924, 0.]]]
-        );
-
-        assert_close_to_literal!(
-            g.get(&weight),
-            [
-                [[[0.13829342, -0.22180916], [-0.11759478, 0.21646728]]],
-                [[[0.25044560, -0.40169036], [-0.21296094, 0.39201635]]],
-            ]
-        );
-
-        assert_close_to_literal!(g.get(&bias), [0.37275729, 0.67505330]);
-    }
-
-    #[test]
-    fn test_conv2d_padding_1() {
-        let dev: TestDevice = Default::default();
-        #[rustfmt::skip]
-        let weight = dev.tensor([
-            [[[0.10215953, 0.06263646], [-0.04124039, -0.09729567]], [[-0.32656857, 0.24254093], [-0.27209827, 0.15361503]]],
-            [[[0.03449896, 0.22931078], [-0.17652659, 0.08222872]],[[-0.06016779, 0.29082409], [-0.19154115, 0.13483226]]],
-            [[[-0.14262493, 0.19654515], [0.15921101, 0.01759464]],[[0.16749159, 0.33096817], [0.28376505, -0.05524009]]],
-        ]).to_dtype::<TestDtype>();
-        let bias = dev
-            .tensor([-0.22854491, 0.28763595, 0.20709404])
-            .to_dtype::<TestDtype>();
-        let x = dev
-            .tensor([[[-0.32224107, -0.32800716]], [[-1.13570976, 0.93713200]]])
-            .to_dtype::<TestDtype>();
-
-        let result = x.leaky_trace().conv2d::<1, 1>(weight.clone())
-            + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
-
-        #[rustfmt::skip]
-        assert_close_to_literal!(
-            result,
-            [
-                [[-0.37165433, 0.26964033, -0.47000977],[-0.52418506, 0.3161699, -0.56809187]],
-                [[0.10800815, 0.66143924, 0.16603859],[-0.11654915, 0.5421771, 0.21993488]],
-                [[0.26416105, -0.22402346, 0.420797],[-0.23212466, 0.3085245, 0.41083777]],
-            ]
-        );
-
-        let g = result.exp().mean().backward();
-
-        assert_close_to_literal!(
-            g.get(&x),
-            [[[0.010052743, 0.038219165]], [[0.0013861917, 0.096129306]]]
-        );
-
-        #[rustfmt::skip]
-        assert_close_to_literal!(
-            g.get(&weight),
-            [
-                [[[-0.03488452, -0.035597768], [-0.03483199, -0.036207683]],[[-0.05705857, 0.03406856], [-0.05008337, 0.024666183]]],
-                [[[-0.053492695, -0.04727108], [-0.05620105, -0.055251926]],[[-0.04363727, 0.033381317], [-0.0607851, 0.030584559]]],
-                [[[-0.051853612, -0.03900232], [-0.04206547, -0.037880093]],[[-0.0073834136, 0.0208545], [0.02886929, -0.040557314]]],
-            ]
-        );
-
-        assert_close_to_literal!(g.get(&bias), [0.28636602, 0.44933242, 0.40484178]);
-    }
-
-    #[test]
-    fn test_conv2d_stride_3_padding_4() {
-        let dev: TestDevice = Default::default();
-        #[rustfmt::skip]
-        let weight = dev.tensor([
-            [[[-0.10252278, -0.14387409, -0.14627469],[0.28396228, -0.14590892, 0.29269591],[0.01090384, 0.14785287, 0.29242596]]],
-            [[[-0.31163597, 0.13224581, -0.20954299],[0.27902845, -0.14735751, 0.14001134],[-0.05224654, 0.16499066, -0.13981307]]],
-        ]).to_dtype::<TestDtype>();
-        let bias = dev
-            .tensor([-0.07123789, -0.17244765])
-            .to_dtype::<TestDtype>();
-        #[rustfmt::skip]
-        let x = dev.tensor([[[0.69103152, 0.25624934],[-0.38448590, 0.03110456],[0.83753252, 0.53786588],[1.15540242, -0.54148245]]]).to_dtype::<TestDtype>();
-
-        let result = x.leaky_trace().conv2d::<3, 4>(weight.clone())
-            + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
-
-        #[rustfmt::skip]
-        assert_close_to_literal!(
-            result,
-            [
-                [[-0.07123789, -0.07123789, -0.07123789],[-0.07123789, -0.14481398, -0.07123789],[-0.07123789, -0.59748650, -0.07123789],[-0.07123789, -0.07123789, -0.07123789]],
-                [[-0.17244765, -0.17244765, -0.17244765],[-0.17244765, -0.3061839, -0.17244765],[-0.17244765, -0.42046443, -0.17244765],[-0.17244765, -0.17244765, -0.17244765]],
-            ]
-        );
-
-        let g = result.exp().mean().backward();
-
-        #[rustfmt::skip]
-        assert_close_to_literal!(
-            g.get(&x),
-            [[[-0.009780421, 0.01484663],[0.010391434, 0.0062526874],[0.00032053515, -0.009087289],[-0.0073772445, 0.0105412705]]]
-        );
-
-        #[rustfmt::skip]
-        assert_close_to_literal!(
-            g.get(&weight),
-            [
-                [[[0.0, 0.019200183, 0.012330416],[0.0, 0.051398464, -0.003175714],[0.0, -0.013860448, 0.0011212977]]],
-                [[[0.0, 0.02291844, 0.01471829],[0.0, 0.05281557, -0.0069562597],[0.0, -0.011794927, 0.00095419877]]],
-            ]
-        );
-
-        assert_close_to_literal!(g.get(&bias), [0.44699076, 0.408709]);
-    }
-
-    #[test]
-    fn test_conv2d_s4p3k2() {
-        let dev = TestDevice::seed_from_u64(432);
-
-        let weight: Tensor<Rank4<3, 5, 2, 2>, TestDtype, _> = dev.sample_normal();
-        let bias: Tensor<Rank1<3>, TestDtype, _> = dev.sample_normal();
-        let x: Tensor<Rank3<5, 7, 6>, TestDtype, _> = dev.sample_normal();
-
-        let out = x.conv2d::<4, 3>(weight);
-        let out = out + bias.broadcast::<_, Axes2<1, 2>>();
-
-        #[rustfmt::skip]
-        assert_close_to_literal!(out, [
-            [[-0.57176435, -0.57176435, -0.57176435],[-0.57176435, 1.0759051, 1.4307989],[-0.57176435, -0.86296344, -1.8794353]],
-            [[0.29306656, 0.29306656, 0.29306656],[0.29306656, 0.9771965, 1.467767],[0.29306656, -6.367015, -2.3370528]],
-            [[-0.19717735, -0.19717735, -0.19717735],[-0.19717735, 1.3412137, 2.9476144],[-0.19717735, 4.247249, -2.1779637]],
-        ]);
-    }
-
-    #[test]
-    fn test_batched_conv2d() {
-        let dev: TestDevice = Default::default();
-        let x: Tensor<Rank3<3, 28, 28>, TestDtype, _> = dev.sample_normal();
-        let w: Tensor<Rank4<5, 3, 6, 6>, TestDtype, _> = dev.sample_normal();
-
-        let y: Tensor<Rank3<5, 9, 9>, _, _, _> = x.leaky_trace().conv2d::<3, 2>(w.clone());
-        let y0 = y.retaped::<NoneTape>();
-        let grads0 = y.square().mean().backward();
-        let x0 = grads0.get(&x);
-        let w0 = grads0.get(&w);
-
-        let x = x
-            .broadcast::<Rank4<10, 3, 28, 28>, _>()
-            .reshape::<Rank4<10, 3, 28, 28>>();
-        assert_eq!(x.strides, x.shape.strides());
-
-        let y: Tensor<Rank4<10, 5, 9, 9>, _, _, _> = x.leaky_trace().conv2d::<3, 2>(w.clone());
-        for i in 0..10 {
-            assert_close_to_tensor!(y0, y.retaped::<NoneTape>().select(dev.tensor(i)));
-        }
-
-        let grads = y.square().mean().backward();
-
-        assert_close_to_tensor!(w0, grads.get(&w), 1e-3);
-
-        let x_grad = grads.get(&x) * 10.0;
-        for i in 0..10 {
-            assert_close_to_tensor!(x0, x_grad.clone().select(dev.tensor(i)));
-        }
-    }
-}
diff --git a/src/tensor_ops/matmul/cuda_kernel.rs b/src/tensor_ops/matmul/cuda_kernel.rs
index c6aa739f2..57aa9286e 100644
--- a/src/tensor_ops/matmul/cuda_kernel.rs
+++ b/src/tensor_ops/matmul/cuda_kernel.rs
@@ -245,7 +245,6 @@ where
         lhs: &Tensor<(B, M, K), E, Self>,
         rhs: &Tensor<(K, N), E, Self>,
     ) -> Result<Tensor<(B, M, N), E, Self>, Self::Err> {
-        assert_ne!(lhs.strides[0], 0);
         let (batch, m, _) = lhs.shape;
         let (k, n) = rhs.shape;
         let shape = (batch, m, n);
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
index 9dfd8b45a..f2b742983 100644
--- a/src/tensor_ops/mod.rs
+++ b/src/tensor_ops/mod.rs
@@ -265,7 +265,7 @@ pub(crate) use to_dtype::ToDtypeKernel;
 #[cfg(feature = "nightly")]
 mod conv2d;
 #[cfg(feature = "nightly")]
-pub use conv2d::{TryConv2D, TryConv2DTo};
+pub use conv2d::TryConv2D;
 
 #[cfg(feature = "nightly")]
 mod convtrans2d;
@@ -279,6 +279,4 @@ pub use upscale2d::{Bilinear, GenericUpscale2D, NearestNeighbor, TryUpscale2D, U
 #[cfg(feature = "nightly")]
 mod pool2d;
 #[cfg(feature = "nightly")]
-pub(crate) use pool2d::{ConstAvgPool2D, ConstMaxPool2D, ConstMinPool2D};
-#[cfg(feature = "nightly")]
-pub use pool2d::{TryAvgPool2D, TryMaxPool2D, TryMinPool2D};
+pub use pool2d::{Pool2DKind, TryPool2D};
diff --git a/src/tensor_ops/pool2d/cpu_kernel.rs b/src/tensor_ops/pool2d/cpu_kernel.rs
index 3f4bb1a10..9a64a9de4 100644
--- a/src/tensor_ops/pool2d/cpu_kernel.rs
+++ b/src/tensor_ops/pool2d/cpu_kernel.rs
@@ -1,9 +1,8 @@
-use crate::shapes::*;
-use crate::tensor::{Cpu, Tensor};
+use crate::{shapes::*, tensor::*};
 
 use std::sync::Arc;
 
-use num_traits::Float;
+use num_traits::{Float, FromPrimitive};
 
 fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
     match S::NUM_DIMS {
@@ -13,172 +12,56 @@ fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
     }
 }
 
-impl<E: Float + Unit + std::ops::AddAssign + std::ops::DivAssign> super::AvgPool2DKernel<E>
-    for Cpu
-{
-    fn forward<I: Shape, O: Shape>(
-        &self,
-        op: super::Pool2DOp,
-        inp: &Tensor<I, E, Self>,
-        out: &mut Tensor<O, E, Self>,
-    ) -> Result<(), Self::Err> {
-        let istr = make_4d::<I>(inp.strides);
-        let ostr = make_4d::<O>(out.strides);
-
-        let buf = inp.data.as_ref();
-        let out_buf = Arc::make_mut(&mut out.data);
-        for b in 0..op.batch {
-            for c in 0..op.chan {
-                for oh in 0..op.h_out {
-                    for ow in 0..op.w_out {
-                        let mut tmp = E::zero();
-                        for k1 in 0..op.kernel {
-                            let y = (oh * op.stride + k1).checked_sub(op.padding);
-                            for k2 in 0..op.kernel {
-                                let x = (ow * op.stride + k2).checked_sub(op.padding);
-                                if let Some((y, x)) = y.zip(x) {
-                                    if y < op.h_in && x < op.w_in {
-                                        let inp_idx =
-                                            b * istr[0] + c * istr[1] + y * istr[2] + x * istr[3];
-                                        tmp += buf[inp_idx];
-                                    }
-                                }
-                            }
-                        }
-                        tmp /= E::from(op.kernel * op.kernel).unwrap();
-                        out_buf[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]] = tmp;
-                    }
-                }
-            }
+impl super::Pool2DKind {
+    fn init<E: Float>(&self) -> E {
+        match self {
+            super::Pool2DKind::Avg => E::zero(),
+            super::Pool2DKind::Min => E::infinity(),
+            super::Pool2DKind::Max => E::neg_infinity(),
         }
-        Ok(())
     }
 
-    fn backward<I: Shape, O: Shape>(
-        &self,
-        op: super::Pool2DOp,
-        inp: &Tensor<I, E, Self>,
-        grad_inp: &mut Self::Vec<E>,
-        out: &Tensor<O, E, Self>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err> {
-        let istr = make_4d::<I>(inp.strides);
-        let ostr = make_4d::<O>(out.strides);
-
-        for b in 0..op.batch {
-            for c in 0..op.chan {
-                for oh in 0..op.h_out {
-                    for ow in 0..op.w_out {
-                        let g = grad_out[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]]
-                            / E::from(op.kernel * op.kernel).unwrap();
-
-                        for k1 in 0..op.kernel {
-                            let y = (oh * op.stride + k1).checked_sub(op.padding);
-                            for k2 in 0..op.kernel {
-                                let x = (ow * op.stride + k2).checked_sub(op.padding);
-                                if let Some((y, x)) = y.zip(x) {
-                                    if x < op.w_in && y < op.h_in {
-                                        grad_inp[b * istr[0]
-                                            + c * istr[1]
-                                            + y * istr[2]
-                                            + x * istr[3]] += g;
-                                    }
-                                }
-                            }
-                        }
-                    }
-                }
-            }
+    fn accum<E: Float>(&self, accum: &E, item: &E) -> E {
+        match self {
+            super::Pool2DKind::Avg => *accum + *item,
+            super::Pool2DKind::Min => accum.min(*item),
+            super::Pool2DKind::Max => accum.max(*item),
         }
-        Ok(())
     }
-}
 
-impl<E: Float + Unit + std::ops::AddAssign> super::MaxPool2DKernel<E> for Cpu {
-    fn forward<I: Shape, O: Shape>(
-        &self,
-        op: super::Pool2DOp,
-        inp: &Tensor<I, E, Self>,
-        out: &mut Tensor<O, E, Self>,
-    ) -> Result<(), Self::Err> {
-        let istr = make_4d::<I>(inp.strides);
-        let ostr = make_4d::<O>(out.strides);
-
-        let buf = inp.data.as_ref();
-        let out_buf = Arc::make_mut(&mut out.data);
-        for b in 0..op.batch {
-            for c in 0..op.chan {
-                for oh in 0..op.h_out {
-                    for ow in 0..op.w_out {
-                        let mut tmp = E::neg_infinity();
-                        for k1 in 0..op.kernel {
-                            let y = (oh * op.stride + k1).checked_sub(op.padding);
-                            for k2 in 0..op.kernel {
-                                let x = (ow * op.stride + k2).checked_sub(op.padding);
-                                if let Some((y, x)) = y.zip(x) {
-                                    if y < op.h_in && x < op.w_in {
-                                        tmp = tmp.max(
-                                            buf[b * istr[0]
-                                                + c * istr[1]
-                                                + y * istr[2]
-                                                + x * istr[3]],
-                                        );
-                                    }
-                                }
-                            }
-                        }
-                        out_buf[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]] = tmp;
-                    }
-                }
-            }
+    fn normalize<E: Float + FromPrimitive>(&self, item: E, num_elements: usize) -> E {
+        match self {
+            super::Pool2DKind::Avg => item * E::from_f64(1.0 / num_elements as f64).unwrap(),
+            super::Pool2DKind::Min => item,
+            super::Pool2DKind::Max => item,
         }
-        Ok(())
     }
-    fn backward<I: Shape, O: Shape>(
-        &self,
-        op: super::Pool2DOp,
-        inp: &Tensor<I, E, Self>,
-        grad_inp: &mut Self::Vec<E>,
-        out: &Tensor<O, E, Self>,
-        grad_out: &Self::Vec<E>,
-    ) -> Result<(), Self::Err> {
-        let istr = make_4d::<I>(inp.strides);
-        let ostr = make_4d::<O>(out.strides);
-
-        let inp_buf = inp.data.as_ref();
-        let out_buf = out.data.as_ref();
 
-        for b in 0..op.batch {
-            for c in 0..op.chan {
-                for oh in 0..op.h_out {
-                    for ow in 0..op.w_out {
-                        let out_idx = b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3];
-                        let go = grad_out[out_idx];
-                        let vo = out_buf[out_idx];
-                        for k1 in 0..op.kernel {
-                            let y = (oh * op.stride + k1).checked_sub(op.padding);
-                            for k2 in 0..op.kernel {
-                                let x = (ow * op.stride + k2).checked_sub(op.padding);
-                                if let Some((y, x)) = y.zip(x) {
-                                    if x < op.w_in && y < op.h_in {
-                                        let inp_idx =
-                                            b * istr[0] + c * istr[1] + y * istr[2] + x * istr[3];
-                                        if inp_buf[inp_idx] == vo {
-                                            grad_inp[inp_idx] += go;
-                                        }
-                                    }
-                                }
-                            }
-                        }
-                    }
+    fn filter<E: Float>(&self, item: E, needle: E, haystack: E) -> E {
+        match self {
+            super::Pool2DKind::Avg => item,
+            super::Pool2DKind::Min => {
+                if needle == haystack {
+                    item
+                } else {
+                    E::zero()
+                }
+            }
+            super::Pool2DKind::Max => {
+                if needle == haystack {
+                    item
+                } else {
+                    E::zero()
                 }
             }
         }
-        Ok(())
     }
 }
 
-impl<E: Float + Unit + std::ops::AddAssign> super::MinPool2DKernel<E> for Cpu {
+impl<E: Float + Dtype> super::Pool2DKernel<E> for Cpu {
+    fn alloc<S: Shape>(&self, s: S) -> Result<Tensor<S, E, Self>, Self::Err> {
+        self.try_zeros_like(&s)
+    }
     fn forward<I: Shape, O: Shape>(
         &self,
         op: super::Pool2DOp,
@@ -194,23 +77,21 @@ impl<E: Float + Unit + std::ops::AddAssign> super::MinPool2DKernel<E> for Cpu {
             for c in 0..op.chan {
                 for oh in 0..op.h_out {
                     for ow in 0..op.w_out {
-                        let mut tmp = E::infinity();
+                        let mut tmp = op.kind.init();
                         for k1 in 0..op.kernel {
-                            let y = (oh * op.stride + k1).checked_sub(op.padding);
+                            let y = (oh * op.stride + op.dilation * k1).checked_sub(op.padding);
                             for k2 in 0..op.kernel {
-                                let x = (ow * op.stride + k2).checked_sub(op.padding);
+                                let x = (ow * op.stride + op.dilation * k2).checked_sub(op.padding);
                                 if let Some((y, x)) = y.zip(x) {
                                     if y < op.h_in && x < op.w_in {
-                                        tmp = tmp.min(
-                                            buf[b * istr[0]
-                                                + c * istr[1]
-                                                + y * istr[2]
-                                                + x * istr[3]],
-                                        );
+                                        let inp_idx =
+                                            b * istr[0] + c * istr[1] + y * istr[2] + x * istr[3];
+                                        tmp = op.kind.accum(&tmp, &buf[inp_idx]);
                                     }
                                 }
                             }
                         }
+                        tmp = op.kind.normalize(tmp, op.kernel * op.kernel);
                         out_buf[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]] = tmp;
                     }
                 }
@@ -238,18 +119,19 @@ impl<E: Float + Unit + std::ops::AddAssign> super::MinPool2DKernel<E> for Cpu {
                     for ow in 0..op.w_out {
                         let out_idx = b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3];
                         let go = grad_out[out_idx];
+                        let go = op.kind.normalize(go, op.kernel * op.kernel);
                         let vo = out_buf[out_idx];
+
                         for k1 in 0..op.kernel {
-                            let y = (oh * op.stride + k1).checked_sub(op.padding);
+                            let y = (oh * op.stride + op.dilation * k1).checked_sub(op.padding);
                             for k2 in 0..op.kernel {
-                                let x = (ow * op.stride + k2).checked_sub(op.padding);
+                                let x = (ow * op.stride + op.dilation * k2).checked_sub(op.padding);
                                 if let Some((y, x)) = y.zip(x) {
                                     if x < op.w_in && y < op.h_in {
                                         let inp_idx =
                                             b * istr[0] + c * istr[1] + y * istr[2] + x * istr[3];
-                                        if inp_buf[inp_idx] == vo {
-                                            grad_inp[inp_idx] += go;
-                                        }
+                                        grad_inp[inp_idx] +=
+                                            op.kind.filter(go, inp_buf[inp_idx], vo);
                                     }
                                 }
                             }
diff --git a/src/tensor_ops/pool2d/cuda_kernel.rs b/src/tensor_ops/pool2d/cuda_kernel.rs
index d750fdfb1..b24952493 100644
--- a/src/tensor_ops/pool2d/cuda_kernel.rs
+++ b/src/tensor_ops/pool2d/cuda_kernel.rs
@@ -19,108 +19,82 @@ fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
     }
 }
 
-macro_rules! pool_impl {
-    ($Trait:tt<$TypeName:ty>, $Fwd:tt, $Bwd:tt) => {
-        impl super::$Trait<$TypeName> for Cuda {
-            fn forward<I: Shape, O: Shape>(
-                &self,
-                op: super::Pool2DOp,
-                inp: &Tensor<I, $TypeName, Self>,
-                out: &mut Tensor<O, $TypeName, Self>,
-            ) -> Result<(), Self::Err> {
-                if !self.dev.has_func($Fwd, $Fwd) {
-                    self.dev.load_ptx(PTX_SRC.into(), $Fwd, &[$Fwd, $Bwd])?;
-                }
-
-                let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
-                let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
-                let fwd_fn = self.dev.get_func($Fwd, $Fwd).unwrap();
-                let cfg = launch_cfg::<128>(out.shape().num_elements() as u32);
-                let params = (
-                    op,                           // const Pool2dOp op,
-                    &inp_strides,                 // const size_t *inp_strides,
-                    &out_strides,                 // const size_t *out_strides,
-                    inp.data.as_ref(),            // const float *inp,
-                    Arc::make_mut(&mut out.data), // float *out
-                );
-                unsafe { fwd_fn.launch(cfg, params) }?;
-                Ok(())
-            }
-            fn backward<I: Shape, O: Shape>(
-                &self,
-                op: super::Pool2DOp,
-                inp: &Tensor<I, $TypeName, Self>,
-                grad_inp: &mut Self::Vec<$TypeName>,
-                out: &Tensor<O, $TypeName, Self>,
-                grad_out: &Self::Vec<$TypeName>,
-            ) -> Result<(), Self::Err> {
-                let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
-                let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
-                let bwd_fn = self.dev.get_func($Fwd, $Bwd).unwrap();
-                let cfg = launch_cfg::<128>(inp.shape().num_elements() as u32);
-                let params = (
-                    op,                // const Pool2dOp op,
-                    &inp_strides,      // const size_t *inp_strides,
-                    &out_strides,      // const size_t *out_strides,
-                    inp.data.as_ref(), // const float *inp,
-                    grad_inp,          // float *grad_inp,
-                    out.data.as_ref(), // const float *out,
-                    grad_out,          // const float *grad_out
-                );
-                unsafe { bwd_fn.launch(cfg, params) }?;
-                Ok(())
-            }
-        }
-    };
+trait HasCudaKernel<E> {
+    const FWD: &'static str;
+    const BWD: &'static str;
 }
 
 #[cfg(feature = "f16")]
-pool_impl!(
-    AvgPool2DKernel<half::f16>,
-    "avg_pool2d_fwd_f16",
-    "avg_pool2d_bwd_f16"
-);
-#[cfg(feature = "f16")]
-pool_impl!(
-    MaxPool2DKernel<half::f16>,
-    "max_pool2d_fwd_f16",
-    "max_pool2d_bwd_f16"
-);
-#[cfg(feature = "f16")]
-pool_impl!(
-    MinPool2DKernel<half::f16>,
-    "min_pool2d_fwd_f16",
-    "min_pool2d_bwd_f16"
-);
+impl HasCudaKernel<half::f16> for Cuda {
+    const FWD: &'static str = "pool2d_fwd_f16";
+    const BWD: &'static str = "pool2d_bwd_f16";
+}
+
+impl HasCudaKernel<f32> for Cuda {
+    const FWD: &'static str = "pool2d_fwd_f32";
+    const BWD: &'static str = "pool2d_bwd_f32";
+}
 
-pool_impl!(
-    AvgPool2DKernel<f32>,
-    "avg_pool2d_fwd_f32",
-    "avg_pool2d_bwd_f32"
-);
-pool_impl!(
-    MaxPool2DKernel<f32>,
-    "max_pool2d_fwd_f32",
-    "max_pool2d_bwd_f32"
-);
-pool_impl!(
-    MinPool2DKernel<f32>,
-    "min_pool2d_fwd_f32",
-    "min_pool2d_bwd_f32"
-);
+impl HasCudaKernel<f64> for Cuda {
+    const FWD: &'static str = "pool2d_fwd_f64";
+    const BWD: &'static str = "pool2d_bwd_f64";
+}
+
+impl<E: Dtype> super::Pool2DKernel<E> for Cuda
+where
+    Self: HasCudaKernel<E>,
+{
+    fn alloc<S: Shape>(&self, s: S) -> Result<Tensor<S, E, Self>, Self::Err> {
+        let data = unsafe { self.alloc_empty::<E>(s.num_elements()) }?;
+        Ok(self.build_tensor(s, s.strides(), data))
+    }
+    fn forward<I: Shape, O: Shape>(
+        &self,
+        op: super::Pool2DOp,
+        inp: &Tensor<I, E, Self>,
+        out: &mut Tensor<O, E, Self>,
+    ) -> Result<(), Self::Err> {
+        if !self.dev.has_func(Self::FWD, Self::FWD) {
+            self.dev
+                .load_ptx(PTX_SRC.into(), Self::FWD, &[Self::FWD, Self::BWD])?;
+        }
 
-pool_impl!(
-    AvgPool2DKernel<f64>,
-    "avg_pool2d_fwd_f64",
-    "avg_pool2d_bwd_f64"
-);
-pool_impl!(
-    MaxPool2DKernel<f64>,
-    "max_pool2d_fwd_f64",
-    "max_pool2d_bwd_f64"
-);
-pool_impl!(
-    MinPool2DKernel<f64>,
-    "min_pool2d_fwd_f64",
-    "min_pool2d_bwd_f64"
-);
+        let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
+        let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
+        let fwd_fn = self.dev.get_func(Self::FWD, Self::FWD).unwrap();
+        let cfg = launch_cfg::<128>(out.shape().num_elements() as u32);
+        let params = (
+            op,                           // const Pool2dOp op,
+            &inp_strides,                 // const size_t *inp_strides,
+            &out_strides,                 // const size_t *out_strides,
+            inp.data.as_ref(),            // const float *inp,
+            Arc::make_mut(&mut out.data), // float *out
+        );
+        unsafe { fwd_fn.launch(cfg, params) }?;
+        Ok(())
+    }
+    fn backward<I: Shape, O: Shape>(
+        &self,
+        op: super::Pool2DOp,
+        inp: &Tensor<I, E, Self>,
+        grad_inp: &mut Self::Vec<E>,
+        out: &Tensor<O, E, Self>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err> {
+        let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
+        let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
+        let bwd_fn = self.dev.get_func(Self::FWD, Self::BWD).unwrap();
+        let cfg = launch_cfg::<128>(inp.shape().num_elements() as u32);
+        let params = (
+            op,                // const Pool2dOp op,
+            &inp_strides,      // const size_t *inp_strides,
+            &out_strides,      // const size_t *out_strides,
+            inp.data.as_ref(), // const float *inp,
+            grad_inp,          // float *grad_inp,
+            out.data.as_ref(), // const float *out,
+            grad_out,          // const float *grad_out
+        );
+        unsafe { bwd_fn.launch(cfg, params) }?;
+        Ok(())
+    }
+}
diff --git a/src/tensor_ops/pool2d/mod.rs b/src/tensor_ops/pool2d/mod.rs
index 666f23e65..44629beec 100644
--- a/src/tensor_ops/pool2d/mod.rs
+++ b/src/tensor_ops/pool2d/mod.rs
@@ -3,19 +3,26 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
-use crate::{
-    shapes::*,
-    tensor::{DeviceStorage, HasErr, PutTape, SplitTape, Tape, Tensor, ZerosTensor},
-};
+use crate::{shapes::*, tensor::*};
 
-use super::conv2d::ConvAlgebra;
+use super::ReshapeTo;
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub enum Pool2DKind {
+    Avg,
+    Min,
+    Max,
+}
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
 pub struct Pool2DOp {
+    pub kind: Pool2DKind,
     pub kernel: usize,
     pub stride: usize,
     pub padding: usize,
+    pub dilation: usize,
     pub batch: usize,
     pub chan: usize,
     pub h_in: usize,
@@ -24,179 +31,201 @@ pub struct Pool2DOp {
     pub w_out: usize,
 }
 
-impl Pool2DOp {
-    fn new(k: usize, s: usize, p: usize, [b, c, h_in, w_in]: [usize; 4]) -> Self {
-        Self {
-            kernel: k,
-            stride: s,
-            padding: p,
-            batch: b,
-            chan: c,
-            h_in,
-            h_out: (h_in + 2 * p - k) / s + 1,
-            w_in,
-            w_out: (w_in + 2 * p - k) / s + 1,
-        }
-    }
-}
-
-macro_rules! pool2d {
-    (Kernel=$Kernel:ident, ConstTrait=$ConstTrait:ident, TryTrait=$TryTrait:ident, Meth=$Meth:ident, TryMeth=$TryMeth:ident) => {
-        pub trait $Kernel<E: Unit>: DeviceStorage {
-            fn forward<I: Shape, O: Shape>(
-                &self,
-                op: Pool2DOp,
-                inp: &Tensor<I, E, Self>,
-                out: &mut Tensor<O, E, Self>,
-            ) -> Result<(), Self::Err>;
+pub(super) trait Pool2DKernel<E: Dtype>: DeviceStorage {
+    fn alloc<S: Shape>(&self, s: S) -> Result<Tensor<S, E, Self>, Self::Err>;
 
-            fn backward<I: Shape, O: Shape>(
-                &self,
-                op: Pool2DOp,
-                inp: &Tensor<I, E, Self>,
-                grad_inp: &mut Self::Vec<E>,
-                out: &Tensor<O, E, Self>,
-                grad_out: &Self::Vec<E>,
-            ) -> Result<(), Self::Err>;
-        }
+    fn forward<I: Shape, O: Shape>(
+        &self,
+        op: Pool2DOp,
+        inp: &Tensor<I, E, Self>,
+        out: &mut Tensor<O, E, Self>,
+    ) -> Result<(), Self::Err>;
 
-        pub trait $ConstTrait<const K: usize, const S: usize, const P: usize>: HasErr {
-            type Output;
-            fn try_pool2d(self) -> Result<Self::Output, Self::Err>;
-        }
+    #[allow(clippy::too_many_arguments)]
+    fn backward<I: Shape, O: Shape>(
+        &self,
+        op: Pool2DOp,
+        inp: &Tensor<I, E, Self>,
+        grad_inp: &mut Self::Vec<E>,
+        out: &Tensor<O, E, Self>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err>;
+}
 
-        pub trait $TryTrait {
-            fn $Meth<const K: usize, const S: usize, const P: usize>(self) -> Self::Output
-            where
-                Self: $ConstTrait<K, S, P>,
-            {
-                self.try_pool2d().unwrap()
-            }
-            fn $TryMeth<const K: usize, const S: usize, const P: usize>(
-                self,
-            ) -> Result<Self::Output, Self::Err>
-            where
-                Self: $ConstTrait<K, S, P>,
-            {
-                self.try_pool2d()
-            }
-        }
-        impl<S: Shape, E: Dtype, D: DeviceStorage, T> $TryTrait for Tensor<S, E, D, T> {}
+pub trait TryPool2D<Kernel, Stride, Padding, Dilation>: Sized {
+    type Pooled;
+    type Error: std::fmt::Debug;
 
-        impl<
-                C: Dim,
-                H: Dim + ConvAlgebra<K, S, P>,
-                W: Dim + ConvAlgebra<K, S, P>,
-                E: Dtype,
-                D: $Kernel<E> + ZerosTensor<E>,
-                T: 'static + Tape<E, D>,
-                const K: usize,
-                const S: usize,
-                const P: usize,
-            > $ConstTrait<K, S, P> for Tensor<(C, H, W), E, D, T>
-        {
-            type Output = Tensor<(C, H::Convolved, W::Convolved), E, D, T>;
+    fn pool2d(
+        self,
+        kind: Pool2DKind,
+        kernel: Kernel,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+    ) -> Self::Pooled {
+        self.try_pool2d(kind, kernel, stride, padding, dilation)
+            .unwrap()
+    }
 
-            fn try_pool2d(self) -> Result<Self::Output, Self::Err> {
-                let h = self.shape.1;
-                let w = self.shape.2;
+    fn try_pool2d(
+        self,
+        kind: Pool2DKind,
+        kernel: Kernel,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+    ) -> Result<Self::Pooled, Self::Error>;
+}
 
-                let &(chan, _, _) = self.shape();
-                let op = Pool2DOp::new(K, S, P, [1, chan.size(), h.size(), w.size()]);
-                let (inp, mut tape) = self.split_tape();
-                let mut out =
-                    inp.device
-                        .try_zeros_like(&(chan, h.convolve_dim(), w.convolve_dim()))?;
-                inp.device.forward(op, &inp, &mut out)?;
-                let inp_ghost = inp.ghost();
-                let out_ghost = out.ghost();
-                let out_clone = out.clone();
-                tape.add_backward_op(move |grads| {
-                    grads.try_alloc_for(&inp_ghost)?;
-                    grads.try_alloc_for(&out_ghost)?;
-                    let (grad_inp, grad_out) = grads.mut_and_ref(&inp_ghost, &out_ghost);
-                    inp_ghost
-                        .dev
-                        .backward(op, &inp, grad_inp, &out_clone, grad_out)
-                });
-                Ok(out.put_tape(tape))
-            }
-        }
+impl<
+        const KERNEL: usize,
+        const STRIDE: usize,
+        const PADDING: usize,
+        const DILATION: usize,
+        const DIM: usize,
+    > TryPool2D<Const<KERNEL>, Const<STRIDE>, Const<PADDING>, Const<DILATION>> for Const<DIM>
+where
+    Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>: Sized,
+{
+    type Pooled = Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>;
+    type Error = std::convert::Infallible;
+    fn try_pool2d(
+        self,
+        _: Pool2DKind,
+        _: Const<KERNEL>,
+        _: Const<STRIDE>,
+        _: Const<PADDING>,
+        _: Const<DILATION>,
+    ) -> Result<Self::Pooled, Self::Error> {
+        Ok(Const)
+    }
+}
 
-        impl<
-                B: Dim,
-                C: Dim,
-                H: Dim + ConvAlgebra<K, S, P>,
-                W: Dim + ConvAlgebra<K, S, P>,
-                E: Dtype,
-                D: $Kernel<E> + ZerosTensor<E>,
-                T: 'static + Tape<E, D>,
-                const K: usize,
-                const S: usize,
-                const P: usize,
-            > $ConstTrait<K, S, P> for Tensor<(B, C, H, W), E, D, T>
-        {
-            type Output = Tensor<(B, C, H::Convolved, W::Convolved), E, D, T>;
+impl<Kernel: Dim, Stride: Dim, Padding: Dim, Dilation: Dim>
+    TryPool2D<Kernel, Stride, Padding, Dilation> for usize
+{
+    type Pooled = usize;
+    type Error = std::convert::Infallible;
+    fn try_pool2d(
+        self,
+        _: Pool2DKind,
+        kernel: Kernel,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+    ) -> Result<Self::Pooled, Self::Error> {
+        Ok((self + 2 * padding.size() - 1)
+            .checked_sub(dilation.size() * (kernel.size() - 1))
+            .unwrap()
+            / stride.size()
+            + 1)
+    }
+}
 
-            fn try_pool2d(self) -> Result<Self::Output, Self::Err> {
-                let h = self.shape.2;
-                let w = self.shape.3;
+impl<Chan, Kernel, Stride, Padding, Dilation, H, W, E, D, T>
+    TryPool2D<Kernel, Stride, Padding, Dilation> for Tensor<(Chan, H, W), E, D, T>
+where
+    Chan: Dim,
+    Kernel: Dim,
+    Stride: Dim,
+    Padding: Dim,
+    Dilation: Dim,
+    H: Dim + TryPool2D<Kernel, Stride, Padding, Dilation>,
+    H::Pooled: Dim,
+    W: Dim + TryPool2D<Kernel, Stride, Padding, Dilation>,
+    W::Pooled: Dim,
+    E: Dtype,
+    D: Pool2DKernel<E> + crate::tensor_ops::reshape_to::ReshapeKernel<E>,
+    T: Tape<E, D>,
+{
+    type Pooled = Tensor<(Chan, H::Pooled, W::Pooled), E, D, T>;
+    type Error = D::Err;
 
-                let &(batch, chan, _, _) = self.shape();
-                let op = Pool2DOp::new(K, S, P, [batch.size(), chan.size(), h.size(), w.size()]);
-                let (inp, mut tape) = self.split_tape();
-                let mut out = inp.device.try_zeros_like(&(
-                    batch,
-                    chan,
-                    h.convolve_dim(),
-                    w.convolve_dim(),
-                ))?;
-                inp.device.forward(op, &inp, &mut out)?;
-                let inp_ghost = inp.ghost();
-                let out_ghost = out.ghost();
-                let out_clone = out.clone();
-                tape.add_backward_op(move |grads| {
-                    grads.try_alloc_for(&inp_ghost)?;
-                    grads.try_alloc_for(&out_ghost)?;
-                    let (grad_inp, grad_out) = grads.mut_and_ref(&inp_ghost, &out_ghost);
-                    inp_ghost
-                        .dev
-                        .backward(op, &inp, grad_inp, &out_clone, grad_out)
-                });
-                Ok(out.put_tape(tape))
-            }
-        }
-    };
+    fn try_pool2d(
+        self,
+        kind: Pool2DKind,
+        kernel: Kernel,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+    ) -> Result<Self::Pooled, Self::Error> {
+        let (chan, h, w) = self.shape;
+        let img = self.try_reshape_like(&(Const::<1>, chan, h, w))?;
+        let out = img.try_pool2d(kind, kernel, stride, padding, dilation)?;
+        let (_, _, out_h, out_w) = out.shape;
+        out.try_reshape_like(&(chan, out_h, out_w))
+    }
 }
 
-pool2d!(
-    Kernel = AvgPool2DKernel,
-    ConstTrait = ConstAvgPool2D,
-    TryTrait = TryAvgPool2D,
-    Meth = avg_pool2d,
-    TryMeth = try_avg_pool2d
-);
+impl<Chan, Kernel, Stride, Padding, Dilation, Batch, H, W, E, D, T>
+    TryPool2D<Kernel, Stride, Padding, Dilation> for Tensor<(Batch, Chan, H, W), E, D, T>
+where
+    Chan: Dim,
+    Kernel: Dim,
+    Stride: Dim,
+    Padding: Dim,
+    Dilation: Dim,
+    Batch: Dim,
+    H: Dim + TryPool2D<Kernel, Stride, Padding, Dilation>,
+    H::Pooled: Dim,
+    W: Dim + TryPool2D<Kernel, Stride, Padding, Dilation>,
+    W::Pooled: Dim,
+    E: Dtype,
+    D: Pool2DKernel<E>,
+    T: Tape<E, D>,
+{
+    type Pooled = Tensor<(Batch, Chan, H::Pooled, W::Pooled), E, D, T>;
+    type Error = D::Err;
 
-pool2d!(
-    Kernel = MaxPool2DKernel,
-    ConstTrait = ConstMaxPool2D,
-    TryTrait = TryMaxPool2D,
-    Meth = max_pool2d,
-    TryMeth = try_max_pool2d
-);
-
-pool2d!(
-    Kernel = MinPool2DKernel,
-    ConstTrait = ConstMinPool2D,
-    TryTrait = TryMinPool2D,
-    Meth = min_pool2d,
-    TryMeth = try_min_pool2d
-);
+    fn try_pool2d(
+        self,
+        kind: Pool2DKind,
+        kernel: Kernel,
+        stride: Stride,
+        padding: Padding,
+        dilation: Dilation,
+    ) -> Result<Self::Pooled, Self::Error> {
+        let (batch, chan, h, w) = self.shape;
+        if self.strides != self.shape.strides() {
+            panic!("Image input to pool2d must be contiguous");
+        }
+        let h_out = h.pool2d(kind, kernel, stride, padding, dilation);
+        let w_out = w.pool2d(kind, kernel, stride, padding, dilation);
+        let op = Pool2DOp {
+            kind,
+            stride: stride.size(),
+            padding: padding.size(),
+            kernel: kernel.size(),
+            dilation: dilation.size(),
+            batch: batch.size(),
+            chan: chan.size(),
+            h_in: h.size(),
+            h_out: h_out.size(),
+            w_in: w.size(),
+            w_out: w_out.size(),
+        };
+        let (img, mut tape) = self.split_tape();
+        let mut out = img.device.alloc((batch, chan, h_out, w_out))?;
+        img.device.forward(op, &img, &mut out)?;
+        let img_ghost = img.ghost();
+        let out_ghost = out.ghost();
+        let out_clone = out.clone();
+        tape.add_backward_op(move |grads| {
+            grads.try_alloc_for(&img_ghost)?;
+            grads.try_alloc_for(&out_ghost)?;
+            let (grad_img, grad_out) = grads.mut_and_ref(&img_ghost, &out_ghost);
+            img.device
+                .backward(op, &img, grad_img, &out_clone, grad_out)
+        });
+        Ok(out.put_tape(tape))
+    }
+}
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{tensor::*, tensor_ops::*, tests::*};
+    use crate::{tensor_ops::*, tests::*};
 
     #[test]
     fn test_pool2d_3d_max2d_eq_grads() {
@@ -204,7 +233,13 @@ mod tests {
         let x = dev
             .tensor([[[1.0, 1., 0.5, 0.2], [0.2, 0.2, 0.5, 1.2]]])
             .to_dtype::<TestDtype>();
-        let r = x.leaky_trace().max_pool2d::<2, 1, 0>();
+        let r = x.leaky_trace().pool2d(
+            Pool2DKind::Max,
+            Const::<2>,
+            Const::<1>,
+            Const::<0>,
+            Const::<1>,
+        );
         assert_close_to_literal!(r, [[[1., 1., 1.2]]]);
         let g = r.sum().backward();
         assert_close_to_literal!(g.get(&x), [[[1., 2., 0., 0.], [0., 0., 0., 1.]]]);
@@ -216,7 +251,13 @@ mod tests {
         let x = dev
             .tensor([[[1., 1., 0.5, 0.2], [0.2, 0.2, 0.5, 1.2]]])
             .to_dtype::<TestDtype>();
-        let r = x.leaky_trace().min_pool2d::<2, 1, 0>();
+        let r = x.leaky_trace().pool2d(
+            Pool2DKind::Min,
+            Const::<2>,
+            Const::<1>,
+            Const::<0>,
+            Const::<1>,
+        );
         assert_close_to_literal!(r, [[[0.2, 0.2, 0.2]]]);
         let g = r.sum().backward();
         assert_close_to_literal!(g.get(&x), [[[0., 0., 0., 1.], [1., 2., 0., 0.]]]);
@@ -226,7 +267,13 @@ mod tests {
     fn test_pool2d_3d_max2d() {
         let dev = TestDevice::seed_from_u64(234);
         let x: Tensor<Rank3<2, 3, 4>, TestDtype, _> = dev.sample_normal();
-        let r = x.leaky_trace().max_pool2d::<2, 2, 0>();
+        let r = x.leaky_trace().pool2d(
+            Pool2DKind::Max,
+            Const::<2>,
+            Const::<2>,
+            Const::<0>,
+            Const::<1>,
+        );
         assert_close_to_literal!(r, [[[1.79155397, 1.10126066]], [[1.14464748, 2.26301837]]]);
         let g = r.exp().mean().backward();
         #[rustfmt::skip]
@@ -243,7 +290,13 @@ mod tests {
     fn test_pool2d_3d_min2d() {
         let dev = TestDevice::seed_from_u64(234);
         let x: Tensor<Rank3<2, 3, 4>, TestDtype, _> = dev.sample_normal();
-        let r = x.leaky_trace().min_pool2d::<2, 2, 0>();
+        let r = x.leaky_trace().pool2d(
+            Pool2DKind::Min,
+            Const::<2>,
+            Const::<2>,
+            Const::<0>,
+            Const::<1>,
+        );
         assert_close_to_literal!(
             r,
             [[[-1.09635627, -1.07717276]], [[-0.01996479, -1.82562149]]]
@@ -259,28 +312,17 @@ mod tests {
         );
     }
 
-    #[test]
-    fn test_pool2d_3d_avg2d() {
-        let dev = TestDevice::seed_from_u64(234);
-        let x: Tensor<Rank3<2, 3, 4>, TestDtype, _> = dev.sample_normal();
-        let r = x.leaky_trace().avg_pool2d::<2, 2, 0>();
-        assert_close_to_literal!(r, [[[0.03031558, -0.25052455]], [[0.39499030, 0.04878314]]]);
-        let g = r.exp().mean().backward();
-        #[rustfmt::skip]
-        assert_close_to_literal!(
-            g.get(&x),
-            [
-                [[0.06442373, 0.06442373, 0.048649523, 0.048649523],[0.06442373, 0.06442373, 0.048649523, 0.048649523],[0.0, 0.0, 0.0, 0.0]],
-                [[0.09277311, 0.09277311, 0.06562454, 0.06562454],[0.09277311, 0.09277311, 0.06562454, 0.06562454],[0.0, 0.0, 0.0, 0.0]]
-            ]
-        );
-    }
-
     #[test]
     fn test_pool2d_4d_avg2d() {
         let dev = TestDevice::seed_from_u64(234);
         let x: Tensor<Rank4<2, 4, 2, 2>, TestDtype, _> = dev.sample_normal();
-        let r = x.leaky_trace().avg_pool2d::<1, 2, 0>();
+        let r = x.leaky_trace().pool2d(
+            Pool2DKind::Avg,
+            Const::<1>,
+            Const::<2>,
+            Const::<0>,
+            Const::<1>,
+        );
         assert_close_to_literal!(
             r,
             [
@@ -298,4 +340,45 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn test_pool2d_dilated() {
+        let dev: TestDevice = Default::default();
+        let x = dev
+            .tensor([[
+                [0., 1., 2., 4., 5.],
+                [6., 7., 8., 9., 10.],
+                [11., 12., 13., 14., 15.],
+                [16., 17., 18., 19., 20.],
+            ]])
+            .to_dtype::<TestDtype>();
+        let y_max = x.leaky_trace().pool2d(
+            Pool2DKind::Max,
+            Const::<2>,
+            Const::<1>,
+            Const::<0>,
+            Const::<2>,
+        );
+        assert_close_to_literal!(y_max, [[[13., 14., 15.], [18., 19., 20.]]]);
+        let y_min = x.clone().pool2d(
+            Pool2DKind::Min,
+            Const::<2>,
+            Const::<1>,
+            Const::<0>,
+            Const::<2>,
+        );
+        assert_close_to_literal!(y_min, [[[0., 1., 2.], [6., 7., 8.]]]);
+
+        let grads = y_max.mean().backward();
+        let v = 1.0 / 6.0;
+        assert_close_to_literal!(
+            grads.get(&x),
+            [[
+                [0., 0., 0., 0., 0.],
+                [0., 0., 0., 0., 0.],
+                [0., 0., v, v, v],
+                [0., 0., v, v, v]
+            ]]
+        );
+    }
 }
diff --git a/src/tensor_ops/pool2d/pool2d.cu b/src/tensor_ops/pool2d/pool2d.cu
index a39d86d11..a84f0f320 100644
--- a/src/tensor_ops/pool2d/pool2d.cu
+++ b/src/tensor_ops/pool2d/pool2d.cu
@@ -1,9 +1,17 @@
 #include "cuda_utils.cuh"
 
+enum Pool2dKind {
+    AVG,
+    MIN,
+    MAX,
+};
+
 struct Pool2dOp {
+    Pool2dKind kind;
     size_t kernel;
     size_t stride;
     size_t padding;
+    size_t dilation;
     size_t batch;
     size_t chan;
     size_t h_in;
@@ -12,208 +20,59 @@ struct Pool2dOp {
     size_t w_out;
 };
 
-template<typename T>
-__device__ void avg_pool2d_fwd(
-    const Pool2dOp op,
-    const size_t *inp_strides,
-    const size_t *out_strides,
-    const T *inp, // 4d (Batch, Channels, Height, Width)
-    T *out // 4d (Batch, Channels, HeightOut, WidthOut)
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const size_t numel = op.batch * op.chan * op.h_out * op.w_out;
-    if (i >= numel) {
-        return;
+__device__ double init(const Pool2dOp op) {
+    switch(op.kind) {
+        case AVG:
+            return 0.0;
+        case MIN:
+            return INFINITY;
+        case MAX:
+            return -INFINITY;
     }
-
-    unsigned int idx = i;
-    const size_t ow = idx % op.w_out;
-    idx /= op.w_out;
-    const size_t oh = idx % op.h_out;
-    idx /= op.h_out;
-    const size_t c = idx % op.chan;
-    idx /= op.chan;
-    const size_t b = idx % op.batch;
-    idx /= op.batch;
-    
-    T tmp = 0.0;
-    for(size_t k1 = 0; k1 < op.kernel; k1++) {
-        for (size_t k2 = 0; k2 < op.kernel; k2++) {
-            const size_t y_plus_p = oh * op.stride + k1;
-            if (y_plus_p < op.padding) { continue; }
-            const size_t y = y_plus_p - op.padding;
-            if (y >= op.h_in) { continue; }
-            const size_t x_plus_p = ow * op.stride + k2;
-            if (x_plus_p < op.padding) { continue; }
-            const size_t x = x_plus_p - op.padding;
-            if (x >= op.w_in) { continue; }
-
-            auto inp_i = b * inp_strides[0] + c * inp_strides[1] + y * inp_strides[2] + x * inp_strides[3];
-            tmp += inp[inp_i];
-        }
-    }
-
-    double num_f64 = op.kernel * op.kernel;
-    T num = num_f64;
-    tmp /= num;
-    out[i] = tmp;
 }
 
 template<typename T>
-__device__ void avg_pool2d_bwd(
-    const Pool2dOp op,
-    const size_t *inp_strides,
-    const size_t *out_strides,
-    const T *inp, // 4d (Batch, Channels, Height, Width)
-    T *grad_inp,
-    const T *out, // 4d (Batch, Channels, HeightOut, WidthOut)
-    const T *grad_out
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const size_t numel = op.batch * op.chan * op.h_in * op.w_in;
-    if (i >= numel) {
-        return;
-    }
-
-    unsigned int idx = i;
-    const size_t x = idx % op.w_in;
-    idx /= op.w_in;
-    const size_t y = idx % op.h_in;
-    idx /= op.h_in;
-    const size_t c = idx % op.chan;
-    idx /= op.chan;
-    const size_t b = idx % op.batch;
-    idx /= op.batch;
-
-    T tmp = 0.0;
-    for(size_t k1 = 0; k1 < op.kernel; k1++) {
-        for (size_t k2 = 0; k2 < op.kernel; k2++) {
-            size_t oh = y + op.padding;
-            if (oh < k1) { continue; }
-            oh -= k1;
-            if (oh % op.stride != 0) { continue; }
-            oh /= op.stride;
-            if (oh >= op.h_out) { continue; }
-
-            size_t ow = x + op.padding;
-            if (ow < k2) { continue; }
-            ow -= k2;
-            if (ow % op.stride != 0) { continue; }
-            ow /= op.stride;
-            if (ow >= op.w_out) { continue; }
-
-            auto out_i = b * out_strides[0] + c * out_strides[1] + oh * out_strides[2] + ow * out_strides[3];
-            tmp += grad_out[out_i];
-        }
+__device__ T accum(const Pool2dOp op, const T accum, const T item) {
+    switch(op.kind) {
+        case AVG:
+            return accum + item;
+        case MIN:
+            return ming(accum, item);
+        case MAX:
+            return maxg(accum, item);
     }
-
-    double num_f64 = op.kernel * op.kernel;
-    T num = num_f64;
-    grad_inp[i] += tmp / num;
 }
 
 template<typename T>
-__device__ void max_pool2d_fwd(
-    const Pool2dOp op,
-    const size_t *inp_strides,
-    const size_t *out_strides,
-    const T *inp, // 4d (Batch, Channels, Height, Width)
-    T *out // 4d (Batch, Channels, HeightOut, WidthOut)
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const size_t numel = op.batch * op.chan * op.h_out * op.w_out;
-    if (i >= numel) {
-        return;
-    }
-
-    unsigned int idx = i;
-    const size_t ow = idx % op.w_out;
-    idx /= op.w_out;
-    const size_t oh = idx % op.h_out;
-    idx /= op.h_out;
-    const size_t c = idx % op.chan;
-    idx /= op.chan;
-    const size_t b = idx % op.batch;
-    idx /= op.batch;
-
-    T tmp = -INFINITY;
-    for(size_t k1 = 0; k1 < op.kernel; k1++) {
-        for (size_t k2 = 0; k2 < op.kernel; k2++) {
-            const size_t y_plus_p = oh * op.stride + k1;
-            if (y_plus_p < op.padding) { continue; }
-            const size_t y = y_plus_p - op.padding;
-            if (y >= op.h_in) { continue; }
-            const size_t x_plus_p = ow * op.stride + k2;
-            if (x_plus_p < op.padding) { continue; }
-            const size_t x = x_plus_p - op.padding;
-            if (x >= op.w_in) { continue; }
-
-            auto inp_i = b * inp_strides[0] + c * inp_strides[1] + y * inp_strides[2] + x * inp_strides[3];
-            tmp = maxg(tmp, inp[inp_i]);
-        }
+__device__ T normalize(const Pool2dOp op, const T item, const size_t num_elements) {
+    double num_f64 = num_elements;
+    double scale_f64 = 1.0 / num_f64;
+    T scale = scale_f64;
+    switch(op.kind) {
+        case AVG:
+            return item * scale;
+        case MIN:
+            return item;
+        case MAX:
+            return item;
     }
-
-    out[i] = tmp;
 }
 
 template<typename T>
-__device__ void max_pool2d_bwd(
-    const Pool2dOp op,
-    const size_t *inp_strides,
-    const size_t *out_strides,
-    const T *inp, // 4d (Batch, Channels, Height, Width)
-    T *grad_inp,
-    const T *out, // 4d (Batch, Channels, HeightOut, WidthOut)
-    const T *grad_out
-) {
-    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
-    const size_t numel = op.batch * op.chan * op.h_in * op.w_in;
-    if (i >= numel) {
-        return;
-    }
-
-    unsigned int idx = i;
-    const size_t x = idx % op.w_in;
-    idx /= op.w_in;
-    const size_t y = idx % op.h_in;
-    idx /= op.h_in;
-    const size_t c = idx % op.chan;
-    idx /= op.chan;
-    const size_t b = idx % op.batch;
-    idx /= op.batch;
-
-    const T inp_v = inp[i];
-
-    T tmp = 0.0;
-    for(size_t k1 = 0; k1 < op.kernel; k1++) {
-        for (size_t k2 = 0; k2 < op.kernel; k2++) {
-            size_t oh = y + op.padding;
-            if (oh < k1) { continue; }
-            oh -= k1;
-            if (oh % op.stride != 0) { continue; }
-            oh /= op.stride;
-            if (oh >= op.h_out) { continue; }
-
-            size_t ow = x + op.padding;
-            if (ow < k2) { continue; }
-            ow -= k2;
-            if (ow % op.stride != 0) { continue; }
-            ow /= op.stride;
-            if (ow >= op.w_out) { continue; }
-
-            auto out_i = b * out_strides[0] + c * out_strides[1] + oh * out_strides[2] + ow * out_strides[3];
-
-            if (out[out_i] == inp_v) {
-                tmp += grad_out[out_i];
-            }
-        }
+__device__ T filter(const Pool2dOp op, const T item, const T needle, const T haystack) {
+    T zero = 0.0;
+    switch(op.kind){
+        case AVG:
+            return item;
+        case MIN:
+            return (needle == haystack) ? item : zero;
+        case MAX:
+            return (needle == haystack) ? item : zero;
     }
-
-    grad_inp[i] += tmp;
 }
 
 template<typename T>
-__device__ void min_pool2d_fwd(
+__device__ void pool2d_fwd(
     const Pool2dOp op,
     const size_t *inp_strides,
     const size_t *out_strides,
@@ -235,29 +94,29 @@ __device__ void min_pool2d_fwd(
     idx /= op.chan;
     const size_t b = idx % op.batch;
     idx /= op.batch;
-
-    T tmp = INFINITY;
+    
+    T tmp = init(op);
     for(size_t k1 = 0; k1 < op.kernel; k1++) {
         for (size_t k2 = 0; k2 < op.kernel; k2++) {
-            const size_t y_plus_p = oh * op.stride + k1;
+            const size_t y_plus_p = oh * op.stride + op.dilation * k1;
             if (y_plus_p < op.padding) { continue; }
             const size_t y = y_plus_p - op.padding;
             if (y >= op.h_in) { continue; }
-            const size_t x_plus_p = ow * op.stride + k2;
+            const size_t x_plus_p = ow * op.stride + op.dilation * k2;
             if (x_plus_p < op.padding) { continue; }
             const size_t x = x_plus_p - op.padding;
             if (x >= op.w_in) { continue; }
 
             auto inp_i = b * inp_strides[0] + c * inp_strides[1] + y * inp_strides[2] + x * inp_strides[3];
-            tmp = ming(tmp, inp[inp_i]);
+            tmp = accum(op, tmp, inp[inp_i]);
         }
     }
 
-    out[i] = tmp;
+    out[i] = normalize(op, tmp, op.kernel * op.kernel);
 }
 
 template<typename T>
-__device__ void min_pool2d_bwd(
+__device__ void pool2d_bwd(
     const Pool2dOp op,
     const size_t *inp_strides,
     const size_t *out_strides,
@@ -288,31 +147,27 @@ __device__ void min_pool2d_bwd(
     for(size_t k1 = 0; k1 < op.kernel; k1++) {
         for (size_t k2 = 0; k2 < op.kernel; k2++) {
             size_t oh = y + op.padding;
-            if (oh < k1) { continue; }
-            oh -= k1;
+            if (oh < op.dilation * k1) { continue; }
+            oh -= op.dilation * k1;
             if (oh % op.stride != 0) { continue; }
             oh /= op.stride;
             if (oh >= op.h_out) { continue; }
 
             size_t ow = x + op.padding;
-            if (ow < k2) { continue; }
-            ow -= k2;
+            if (ow < op.dilation * k2) { continue; }
+            ow -= op.dilation * k2;
             if (ow % op.stride != 0) { continue; }
             ow /= op.stride;
             if (ow >= op.w_out) { continue; }
 
             auto out_i = b * out_strides[0] + c * out_strides[1] + oh * out_strides[2] + ow * out_strides[3];
-
-            if (out[out_i] == inp_v) {
-                tmp += grad_out[out_i];
-            }
+            tmp += filter(op, grad_out[out_i], out[out_i], inp_v);
         }
     }
-
-    grad_inp[i] += tmp;
+    grad_inp[i] += normalize(op, tmp, op.kernel * op.kernel);
 }
 
-#define POOL_OP(TYPENAME, fwd, bwd, fwd_FN, bwd_FN) \
+#define POOL_OP(TYPENAME, fwd, bwd) \
 extern "C" __global__ void fwd( \
     const Pool2dOp op, \
     const size_t *inp_strides, \
@@ -320,7 +175,7 @@ extern "C" __global__ void fwd( \
     const TYPENAME *inp, \
     TYPENAME *out \
 ) { \
-    fwd_FN(op, inp_strides, out_strides, inp, out); \
+    pool2d_fwd(op, inp_strides, out_strides, inp, out); \
 } \
 extern "C" __global__ void bwd( \
     const Pool2dOp op, \
@@ -331,53 +186,9 @@ extern "C" __global__ void bwd( \
     const TYPENAME *out, \
     const TYPENAME *grad_out \
 ) { \
-    bwd_FN(op, inp_strides, out_strides, inp, grad_inp, out, grad_out); \
+    pool2d_bwd(op, inp_strides, out_strides, inp, grad_inp, out, grad_out); \
 }
 
-POOL_OP(
-    __half,
-    avg_pool2d_fwd_f16, avg_pool2d_bwd_f16,
-    avg_pool2d_fwd, avg_pool2d_bwd
-);
-POOL_OP(
-    __half,
-    min_pool2d_fwd_f16, min_pool2d_bwd_f16,
-    min_pool2d_fwd, min_pool2d_bwd
-);
-POOL_OP(
-    __half,
-    max_pool2d_fwd_f16, max_pool2d_bwd_f16,
-    max_pool2d_fwd, max_pool2d_bwd
-);
-
-POOL_OP(
-    float,
-    avg_pool2d_fwd_f32, avg_pool2d_bwd_f32,
-    avg_pool2d_fwd, avg_pool2d_bwd
-);
-POOL_OP(
-    float,
-    min_pool2d_fwd_f32, min_pool2d_bwd_f32,
-    min_pool2d_fwd, min_pool2d_bwd
-);
-POOL_OP(
-    float,
-    max_pool2d_fwd_f32, max_pool2d_bwd_f32,
-    max_pool2d_fwd, max_pool2d_bwd
-);
-
-POOL_OP(
-    double,
-    avg_pool2d_fwd_f64, avg_pool2d_bwd_f64,
-    avg_pool2d_fwd, avg_pool2d_bwd
-);
-POOL_OP(
-    double,
-    min_pool2d_fwd_f64, min_pool2d_bwd_f64,
-    min_pool2d_fwd, min_pool2d_bwd
-);
-POOL_OP(
-    double,
-    max_pool2d_fwd_f64, max_pool2d_bwd_f64,
-    max_pool2d_fwd, max_pool2d_bwd
-);
+POOL_OP(__half, pool2d_fwd_f16, pool2d_bwd_f16);
+POOL_OP(float, pool2d_fwd_f32, pool2d_bwd_f32);
+POOL_OP(double, pool2d_fwd_f64, pool2d_bwd_f64);


Test patch:
diff --git a/src/tensor_ops/conv2d/tests.rs b/src/tensor_ops/conv2d/tests.rs
new file mode 100644
index 000000000..375705bec
--- /dev/null
+++ b/src/tensor_ops/conv2d/tests.rs
@@ -0,0 +1,341 @@
+use super::*;
+use crate::{tensor_ops::*, tests::*};
+
+#[test]
+/// Produced by
+/// ```python
+/// q = torch.nn.Conv2d(1, 2, 2)
+/// x = torch.sample_normal(1, 2, 3, requires_grad=True)
+/// q(x).exp().mean().backward()
+/// ```
+fn test_conv2d_default_stride_and_padding() {
+    let dev: TestDevice = Default::default();
+    let weight = dev
+        .tensor([
+            [[[-0.04958433, -0.43007267], [0.01935136, 0.09778714]]],
+            [[[0.44083858, -0.20507240], [-0.30017477, -0.10937047]]],
+        ])
+        .to_dtype::<TestDtype>();
+    let bias = dev
+        .tensor([0.36406237, -0.30981010])
+        .to_dtype::<TestDtype>();
+    let x = dev
+        .tensor([[
+            [-0.86713916, 0.52773184, -0.95238322],
+            [-0.64531374, 0.77809018, -0.49099201],
+        ]])
+        .to_dtype::<TestDtype>();
+    let result = (x.leaky_trace(), weight.clone())
+        .conv2d(Const::<1>, Const::<0>, Const::<1>, Const::<1>)
+        + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
+    assert_close_to_literal!(
+        result,
+        [[[0.24369538, 0.71453357]], [[-0.69169492, -0.06172103]]]
+    );
+    let g = result.exp().mean().backward();
+    assert_close_to_literal!(
+        g.get(&x),
+        [[
+            [0.03936806, -0.08457474, -0.26788417],
+            [-0.03140351, -0.04316529, 0.02424446],
+        ]]
+    );
+    assert_close_to_literal!(
+        g.get(&weight),
+        [
+            [[[-0.00703794, -0.31814471], [0.19160703, -0.00260070]]],
+            [[[0.01548620, -0.15778227], [0.10209797, -0.01799832]]],
+        ]
+    );
+    assert_close_to_literal!(g.get(&bias), [0.82979727, 0.36021793]);
+}
+
+#[test]
+/// Produced by
+/// ```python
+/// q = torch.nn.Conv2d(1, 2, 2, stride=2)
+/// x = torch.sample_normal(1, 2, 3, requires_grad=True)
+/// q(x).exp().mean().backward()
+/// ```
+fn test_conv2d_stride_2() {
+    let dev: TestDevice = Default::default();
+    let weight = dev
+        .tensor([
+            [[[0.44704646, -0.29563826], [0.29228759, -0.16575140]]],
+            [[[-0.30488998, 0.25222939], [0.13279295, 0.38153177]]],
+        ])
+        .to_dtype::<TestDtype>();
+    let bias = dev
+        .tensor([-0.44699109, 0.38371694])
+        .to_dtype::<TestDtype>();
+    let x = dev
+        .tensor([[
+            [0.37100124, -0.59504986, -1.19781005],
+            [-0.31547278, 0.58071911, 0.86612970],
+        ]])
+        .to_dtype::<TestDtype>();
+
+    let result = (x.leaky_trace(), weight.clone())
+        .conv2d(Const::<2>, Const::<0>, Const::<1>, Const::<1>)
+        + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
+    assert_close_to_literal!(result, [[[-0.29368058]], [[0.30018353]]]);
+
+    let g = result.exp().mean().backward();
+
+    assert_close_to_literal!(
+        g.get(&x),
+        [[[-0.03917716, 0.06006697, 0.], [0.19859464, 0.19576924, 0.]]]
+    );
+
+    assert_close_to_literal!(
+        g.get(&weight),
+        [
+            [[[0.13829342, -0.22180916], [-0.11759478, 0.21646728]]],
+            [[[0.25044560, -0.40169036], [-0.21296094, 0.39201635]]],
+        ]
+    );
+
+    assert_close_to_literal!(g.get(&bias), [0.37275729, 0.67505330]);
+}
+
+#[test]
+fn test_conv2d_padding_1() {
+    let dev: TestDevice = Default::default();
+    #[rustfmt::skip]
+        let weight = dev.tensor([
+            [[[0.10215953, 0.06263646], [-0.04124039, -0.09729567]], [[-0.32656857, 0.24254093], [-0.27209827, 0.15361503]]],
+            [[[0.03449896, 0.22931078], [-0.17652659, 0.08222872]],[[-0.06016779, 0.29082409], [-0.19154115, 0.13483226]]],
+            [[[-0.14262493, 0.19654515], [0.15921101, 0.01759464]],[[0.16749159, 0.33096817], [0.28376505, -0.05524009]]],
+        ]).to_dtype::<TestDtype>();
+    let bias = dev
+        .tensor([-0.22854491, 0.28763595, 0.20709404])
+        .to_dtype::<TestDtype>();
+    let x = dev
+        .tensor([[[-0.32224107, -0.32800716]], [[-1.13570976, 0.93713200]]])
+        .to_dtype::<TestDtype>();
+
+    let result = (x.leaky_trace(), weight.clone())
+        .conv2d(Const::<1>, Const::<1>, Const::<1>, Const::<1>)
+        + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
+
+    #[rustfmt::skip]
+        assert_close_to_literal!(
+            result,
+            [
+                [[-0.37165433, 0.26964033, -0.47000977],[-0.52418506, 0.3161699, -0.56809187]],
+                [[0.10800815, 0.66143924, 0.16603859],[-0.11654915, 0.5421771, 0.21993488]],
+                [[0.26416105, -0.22402346, 0.420797],[-0.23212466, 0.3085245, 0.41083777]],
+            ]
+        );
+
+    let g = result.exp().mean().backward();
+
+    assert_close_to_literal!(
+        g.get(&x),
+        [[[0.010052743, 0.038219165]], [[0.0013861917, 0.096129306]]]
+    );
+
+    #[rustfmt::skip]
+        assert_close_to_literal!(
+            g.get(&weight),
+            [
+                [[[-0.03488452, -0.035597768], [-0.03483199, -0.036207683]],[[-0.05705857, 0.03406856], [-0.05008337, 0.024666183]]],
+                [[[-0.053492695, -0.04727108], [-0.05620105, -0.055251926]],[[-0.04363727, 0.033381317], [-0.0607851, 0.030584559]]],
+                [[[-0.051853612, -0.03900232], [-0.04206547, -0.037880093]],[[-0.0073834136, 0.0208545], [0.02886929, -0.040557314]]],
+            ]
+        );
+
+    assert_close_to_literal!(g.get(&bias), [0.28636602, 0.44933242, 0.40484178]);
+}
+
+#[test]
+fn test_conv2d_stride_3_padding_4() {
+    let dev: TestDevice = Default::default();
+    #[rustfmt::skip]
+        let weight = dev.tensor([
+            [[[-0.10252278, -0.14387409, -0.14627469],[0.28396228, -0.14590892, 0.29269591],[0.01090384, 0.14785287, 0.29242596]]],
+            [[[-0.31163597, 0.13224581, -0.20954299],[0.27902845, -0.14735751, 0.14001134],[-0.05224654, 0.16499066, -0.13981307]]],
+        ]).to_dtype::<TestDtype>();
+    let bias = dev
+        .tensor([-0.07123789, -0.17244765])
+        .to_dtype::<TestDtype>();
+    #[rustfmt::skip]
+        let x = dev.tensor([[[0.69103152, 0.25624934],[-0.38448590, 0.03110456],[0.83753252, 0.53786588],[1.15540242, -0.54148245]]]).to_dtype::<TestDtype>();
+
+    let result = (x.leaky_trace(), weight.clone())
+        .conv2d(Const::<3>, Const::<4>, Const::<1>, Const::<1>)
+        + bias.leaky_trace().broadcast::<_, Axes2<1, 2>>();
+
+    #[rustfmt::skip]
+        assert_close_to_literal!(
+            result,
+            [
+                [[-0.07123789, -0.07123789, -0.07123789],[-0.07123789, -0.14481398, -0.07123789],[-0.07123789, -0.59748650, -0.07123789],[-0.07123789, -0.07123789, -0.07123789]],
+                [[-0.17244765, -0.17244765, -0.17244765],[-0.17244765, -0.3061839, -0.17244765],[-0.17244765, -0.42046443, -0.17244765],[-0.17244765, -0.17244765, -0.17244765]],
+            ]
+        );
+
+    let g = result.exp().mean().backward();
+
+    #[rustfmt::skip]
+        assert_close_to_literal!(
+            g.get(&x),
+            [[[-0.009780421, 0.01484663],[0.010391434, 0.0062526874],[0.00032053515, -0.009087289],[-0.0073772445, 0.0105412705]]]
+        );
+
+    #[rustfmt::skip]
+        assert_close_to_literal!(
+            g.get(&weight),
+            [
+                [[[0.0, 0.019200183, 0.012330416],[0.0, 0.051398464, -0.003175714],[0.0, -0.013860448, 0.0011212977]]],
+                [[[0.0, 0.02291844, 0.01471829],[0.0, 0.05281557, -0.0069562597],[0.0, -0.011794927, 0.00095419877]]],
+            ]
+        );
+
+    assert_close_to_literal!(g.get(&bias), [0.44699076, 0.408709]);
+}
+
+#[test]
+fn test_conv2d_s4p3k2() {
+    let dev = TestDevice::seed_from_u64(432);
+
+    let weight: Tensor<Rank4<3, 5, 2, 2>, TestDtype, _> = dev.sample_normal();
+    let bias: Tensor<Rank1<3>, TestDtype, _> = dev.sample_normal();
+    let x: Tensor<Rank3<5, 7, 6>, TestDtype, _> = dev.sample_normal();
+
+    let out =
+        (x.leaky_trace(), weight.clone()).conv2d(Const::<4>, Const::<3>, Const::<1>, Const::<1>);
+    let out = out + bias.broadcast::<_, Axes2<1, 2>>();
+
+    #[rustfmt::skip]
+        assert_close_to_literal!(out, [
+            [[-0.57176435, -0.57176435, -0.57176435],[-0.57176435, 1.0759051, 1.4307989],[-0.57176435, -0.86296344, -1.8794353]],
+            [[0.29306656, 0.29306656, 0.29306656],[0.29306656, 0.9771965, 1.467767],[0.29306656, -6.367015, -2.3370528]],
+            [[-0.19717735, -0.19717735, -0.19717735],[-0.19717735, 1.3412137, 2.9476144],[-0.19717735, 4.247249, -2.1779637]],
+        ]);
+}
+
+#[test]
+fn test_batched_conv2d() {
+    let dev: TestDevice = Default::default();
+    let x: Tensor<Rank3<3, 28, 28>, TestDtype, _> = dev.sample_normal();
+    let w: Tensor<Rank4<5, 3, 6, 6>, TestDtype, _> = dev.sample_normal();
+
+    let y: Tensor<Rank3<5, 9, 9>, _, _, _> =
+        (x.leaky_trace(), w.clone()).conv2d(Const::<3>, Const::<2>, Const::<1>, Const::<1>);
+    let y0 = y.retaped::<NoneTape>();
+    let grads0 = y.square().mean().backward();
+    let x0 = grads0.get(&x);
+    let w0 = grads0.get(&w);
+
+    let x = x
+        .broadcast::<Rank4<10, 3, 28, 28>, _>()
+        .reshape::<Rank4<10, 3, 28, 28>>();
+    assert_eq!(x.strides, x.shape.strides());
+
+    let y: Tensor<Rank4<10, 5, 9, 9>, _, _, _> =
+        (x.leaky_trace(), w.clone()).conv2d(Const::<3>, Const::<2>, Const::<1>, Const::<1>);
+    for i in 0..10 {
+        assert_close_to_tensor!(y0, y.retaped::<NoneTape>().select(dev.tensor(i)));
+    }
+
+    let grads = y.square().mean().backward();
+
+    assert_close_to_tensor!(w0, grads.get(&w), 1e-3);
+
+    let x_grad = grads.get(&x) * 10.0;
+    for i in 0..10 {
+        assert_close_to_tensor!(x0, x_grad.clone().select(dev.tensor(i)));
+    }
+}
+
+#[test]
+fn test_conv2d_dilated() {
+    let dev: TestDevice = Default::default();
+    let x = dev
+        .tensor([[
+            [0., 1., 2., 4., 5.],
+            [6., 7., 8., 9., 10.],
+            [11., 12., 13., 14., 15.],
+            [16., 17., 18., 19., 20.],
+        ]])
+        .to_dtype::<TestDtype>();
+    let w = dev
+        .tensor([[[[0.1, 0.5], [1.0, 2.0]]]])
+        .to_dtype::<TestDtype>();
+    let y = (x.leaky_trace(), w.clone()).conv2d(Const::<1>, Const::<0>, Const::<2>, Const::<1>);
+    assert_close_to_literal!(y, [[[38.0, 42.1, 45.7], [56.6, 60.2, 63.8]]]);
+    let grads = y.mean().backward();
+    assert_close_to_literal!(
+        grads.get(&x),
+        [[
+            [0.016666668, 0.016666668, 0.1, 0.083333336, 0.083333336],
+            [0.016666668, 0.016666668, 0.1, 0.083333336, 0.083333336],
+            [0.16666667, 0.16666667, 0.5, 0.33333334, 0.33333334],
+            [0.16666667, 0.16666667, 0.5, 0.33333334, 0.33333334]
+        ]]
+    );
+    assert_close_to_literal!(grads.get(&w), [[[[4.0, 6.3333335], [14.5, 16.5]]]]);
+}
+
+#[test]
+fn test_conv2d_grouped_forward() {
+    const NUM_GROUPS: usize = 3;
+    let dev: TestDevice = Default::default();
+    let x: Tensor<Rank4<2, 9, 14, 14>, TestDtype, _> = dev.sample_normal();
+    let w: Tensor<Rank4<15, 3, 3, 3>, TestDtype, _> = dev.sample_normal();
+
+    let y = (x.leaky_trace(), w.clone()).conv2d(
+        Const::<1>,
+        Const::<0>,
+        Const::<1>,
+        Const::<NUM_GROUPS>,
+    );
+
+    for i in 0..NUM_GROUPS {
+        let x_group = x
+            .clone()
+            .slice((.., 3 * i..3 * (i + 1), .., ..))
+            .realize::<(Const<2>, Const<3>, Const<14>, Const<14>)>();
+        let w_group = w
+            .clone()
+            .slice((5 * i..5 * (i + 1), .., .., ..))
+            .realize::<(Const<5>, Const<3>, Const<3>, Const<3>)>();
+        let y_group = (x_group, w_group).conv2d(Const::<1>, Const::<0>, Const::<1>, Const::<1>);
+        let y_group_true = y
+            .retaped::<NoneTape>()
+            .slice((.., 5 * i..5 * (i + 1), .., ..))
+            .realize::<(Const<2>, Const<5>, Const<12>, Const<12>)>();
+        assert_close_to_tensor!(y_group, y_group_true);
+    }
+
+    let grads = y.exp().sum().backward();
+    let x_grad = grads.get(&x);
+    let w_grad = grads.get(&w);
+
+    for i in 0..NUM_GROUPS {
+        let x_group = x
+            .clone()
+            .slice((.., 3 * i..3 * (i + 1), .., ..))
+            .realize::<(Const<2>, Const<3>, Const<14>, Const<14>)>();
+        let w_group = w
+            .clone()
+            .slice((5 * i..5 * (i + 1), .., .., ..))
+            .realize::<(Const<5>, Const<3>, Const<3>, Const<3>)>();
+        let y_group = (x_group.leaky_trace(), w_group.clone())
+            .conv2d(Const::<1>, Const::<0>, Const::<1>, Const::<1>);
+        let grads = y_group.exp().sum().backward();
+
+        let x_grad_group_true = x_grad
+            .clone()
+            .slice((.., 3 * i..3 * (i + 1), .., ..))
+            .realize::<(Const<2>, Const<3>, Const<14>, Const<14>)>();
+        let w_grad_group_true = w_grad
+            .clone()
+            .slice((5 * i..5 * (i + 1), .., .., ..))
+            .realize::<(Const<5>, Const<3>, Const<3>, Const<3>)>();
+
+        assert_close_to_tensor!(grads.get(&x_group), x_grad_group_true);
+        assert_close_to_tensor!(grads.get(&w_group), w_grad_group_true);
+    }
+}
diff --git a/tests/test_resnet18_f32.rs b/tests/test_resnet18_f32.rs
index 117b00f1b..3dbfb120b 100644
--- a/tests/test_resnet18_f32.rs
+++ b/tests/test_resnet18_f32.rs
@@ -1,4 +1,7 @@
-#![cfg_attr(feature = "nightly", feature(generic_const_exprs))]
+#![cfg_attr(
+    all(feature = "test-integrations", feature = "nightly"),
+    feature(generic_const_exprs)
+)]
 
 #[cfg(all(feature = "test-integrations", feature = "nightly"))]
 mod model {