Repository: coreylowman/dfdx (language: rust)

Issue:
Add upsample and conv-transpose operations
In order to support a UNet architecture, used my many generative models for images, there needs to be a way to increase the tensor resolution. The common ways to do this are the convolution transpose operation, or the upsample operation. Sadly, dfdx supports neither of these right now.

It's possible, with some work, to construct a simple nearest-neighbor upsample with an integer increase with the select operation, but doing so is not intuitive or easy, and doesn't seem like it'd be very efficient.

While not needed for UNets specifically, PyTorch also supports upsamples by fractional values, such as increasing the resolution by 1.6x. This would pose a significant challenge for the dfdx architecture with it's hardcoded tensor dimensions, as I don't think it's possible to do the necessary math with const generics to statically determine the dimensions of the output tensor.


Discussion / hints:
Luckily we have runtime dimensions now, so I think this would be pretty straightforward
Oh, excellent. I see you just pressed a new release today. I was still running on v0.10.
Is there an example on using runtime dimensions and a guideline about when to use it?
I would be interested in working on conv-transpose. I don't know much about cuda but I think I could get a cpu forward implementation working, but I'm not familiar with how the backward is working (I'm familiar with calculus but not as much with matrix calculus). I assume it should have the layer struct as `src/nn/conv-transpose.rs` and then the actual ops at `src/tensor-ops/conv-transpose/[mod|cpu_kernel].rs`.

Fix patch:
diff --git a/src/nn/convtrans.rs b/src/nn/convtrans.rs
new file mode 100644
index 000000000..12f708e79
--- /dev/null
+++ b/src/nn/convtrans.rs
@@ -0,0 +1,194 @@
+use num_traits::Float;
+use rand_distr::uniform::SampleUniform;
+
+use crate::{shapes::*, tensor::*, tensor_ops::*};
+
+use super::*;
+
+pub mod builder {
+    #[derive(Debug)]
+    pub struct ConvTrans2D<
+        const IN_CHAN: usize,
+        const OUT_CHAN: usize,
+        const KERNEL_SIZE: usize,
+        const STRIDE: usize = 1,
+        const PADDING: usize = 0,
+    >;
+}
+
+impl<const I: usize, const O: usize, const K: usize, const S: usize, const P: usize, E, D>
+    BuildOnDevice<D, E> for builder::ConvTrans2D<I, O, K, S, P>
+where
+    E: Dtype,
+    D: Device<E>,
+    ConvTrans2D<I, O, K, S, P, E, D>: BuildModule<D, E>,
+{
+    type Built = ConvTrans2D<I, O, K, S, P, E, D>;
+    fn try_build_on_device(device: &D) -> Result<Self::Built, <D>::Err> {
+        Self::Built::try_build(device)
+    }
+}
+
+/// **Requires Nightly** Performs *unbiased* 2d deconvolutions on 3d and 4d images.
+///
+/// **Pytorch Equivalent**: `torch.nn.ConvTranspose2d(..., bias=False)`
+///
+/// To create a biased conv, combine with [crate::nn::modules::Bias2D]:
+/// ```ignore
+/// # use dfdx::prelude::*;
+/// type BiasedConv = (ConvTrans2D<3, 5, 4>, Bias2D<5>);
+/// ```
+///
+/// Generics:
+/// - `IN_CHAN`: The number of input channels in an image.
+/// - `OUT_CHAN`: The number of channels in the output of the layer.
+/// - `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`.
+#[derive(Debug, Clone)]
+pub struct ConvTrans2D<
+    const IN_CHAN: usize,
+    const OUT_CHAN: usize,
+    const KERNEL_SIZE: usize,
+    const STRIDE: usize,
+    const PADDING: 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 ConvTrans2D<I, O, K, S, P, E, D>
+where
+    E: Dtype + Float + SampleUniform,
+    D: Device<E>,
+{
+    type To<E2: Dtype, D2: Device<E2>> = ConvTrans2D<I, O, K, S, P, E2, D2>;
+
+    fn iter_tensors<V: ModuleVisitor<Self, E, D>>(
+        visitor: &mut V,
+    ) -> Result<Option<Self::To<V::E2, V::D2>>, V::Err> {
+        visitor.visit_fields(
+            Self::tensor(
+                "weight",
+                |s| &s.weight,
+                |s| &mut s.weight,
+                TensorOptions::reset_with(|t| {
+                    let b = E::ONE / E::from_usize(I * K * K).unwrap().sqrt();
+                    t.try_fill_with_distr(rand_distr::Uniform::new(-b, b))
+                }),
+            ),
+            |weight| ConvTrans2D { weight },
+        )
+    }
+}
+
+#[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 ConvTrans2D<C, O, K, S, P, E, D>
+where
+    E: Dtype,
+    D: Device<E>,
+    Img: TryConvTrans2DTo<Tensor<Rank4<O, C, K, K>, E, D>, S, P> + HasErr<Err = D::Err>,
+{
+    type Output = Img::Output;
+    type Error = D::Err;
+
+    fn try_forward(&self, x: Img) -> Result<Self::Output, D::Err> {
+        x.try_convtrans2d_to(self.weight.clone())
+    }
+}
+
+impl<const I: usize, const O: usize, const K: usize, const S: usize, const P: usize, E, D>
+    NonMutableModule for ConvTrans2D<I, O, K, S, P, E, D>
+where
+    E: Dtype,
+    D: DeviceStorage,
+{
+}
+
+#[cfg(feature = "nightly")]
+#[cfg(test)]
+mod tests {
+    use crate::{
+        optim::*,
+        tensor::{AsArray, SampleTensor, ZerosTensor},
+        tests::*,
+    };
+
+    use super::{builder::ConvTrans2D, *};
+
+    #[rustfmt::skip]
+    #[test]
+    fn test_forward_3d_sizes() {
+        let dev: TestDevice = Default::default();
+        let x = dev.zeros::<Rank3<3, 8, 8>>();
+        let _: Tensor<Rank3<2, 10, 10>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<4, 10, 10>, _, _, _> = dev.build_module::<ConvTrans2D<3, 4, 3>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<4, 9, 9>, _, _, _> = dev.build_module::<ConvTrans2D<3, 4, 2>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<4, 11, 11>, _, _, _> = dev.build_module::<ConvTrans2D<3, 4, 4>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<2, 17, 17>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 2>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<2, 24, 24>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 3>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<2, 8, 8>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 1, 1>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<2, 6, 6>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 1, 2>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank3<2, 13, 13>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 2, 2>, TestDtype>().forward(x.clone());
+    }
+
+    #[rustfmt::skip]
+    #[test]
+    fn test_forward_4d_sizes() {
+        let dev: TestDevice = Default::default();
+        let x = dev.zeros::<Rank4<5, 3, 8, 8>>();
+        let _: Tensor<Rank4<5, 2, 10, 10>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 4, 10, 10>, _, _, _> = dev.build_module::<ConvTrans2D<3, 4, 3>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 4, 9, 9>, _, _, _> = dev.build_module::<ConvTrans2D<3, 4, 2>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 4, 11, 11>, _, _, _> = dev.build_module::<ConvTrans2D<3, 4, 4>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 2, 17, 17>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 2>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 2, 24, 24>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 3>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 2, 8, 8>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 1, 1>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 2, 6, 6>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 1, 2>, TestDtype>().forward(x.clone());
+        let _: Tensor<Rank4<5, 2, 13, 13>, _, _, _> = dev.build_module::<ConvTrans2D<3, 2, 3, 2, 2>, TestDtype>().forward(x.clone());
+    }
+
+    #[test]
+    fn test_2_conv_sizes() {
+        let dev = Cpu::default();
+        type A = ConvTrans2D<4, 2, 3>;
+        type B = ConvTrans2D<2, 1, 3>;
+        let _: Tensor<Rank3<1, 10, 10>, _, _> = dev
+            .build_module::<(A, B), TestDtype>()
+            .forward(dev.zeros::<Rank3<4, 6, 6>>());
+    }
+
+    #[test]
+    fn test_3_conv_sizes() {
+        type A = ConvTrans2D<2, 1, 3>;
+        type B = ConvTrans2D<4, 2, 3>;
+        type C = ConvTrans2D<1, 4, 1, 1, 1>;
+
+        let dev = Cpu::default();
+        let _: Tensor<Rank3<1, 10, 10>, _, _> = dev
+            .build_module::<(C, B, A), TestDtype>()
+            .forward_mut(dev.zeros::<Rank3<1, 8, 8>>());
+    }
+
+    #[test]
+    fn test_conv_with_optimizer() {
+        let dev: TestDevice = Default::default();
+
+        let mut m = dev.build_module::<ConvTrans2D<2, 4, 3>, TestDtype>();
+
+        let weight_init = m.weight.clone();
+
+        let mut opt = Sgd::new(&m, Default::default());
+        let out = m.forward(dev.sample_normal::<Rank4<8, 2, 28, 28>>().leaky_trace());
+        let g = out.square().mean().backward();
+
+        assert_ne!(g.get(&m.weight).array(), [[[[0.0; 3]; 3]; 2]; 4]);
+
+        opt.update(&mut m, &g).expect("unused params");
+
+        assert_ne!(weight_init.array(), m.weight.array());
+    }
+}
diff --git a/src/nn/mod.rs b/src/nn/mod.rs
index ae483ae28..95c564764 100644
--- a/src/nn/mod.rs
+++ b/src/nn/mod.rs
@@ -189,6 +189,7 @@ mod batchnorm1d;
 mod batchnorm2d;
 mod bias2d;
 mod conv;
+mod convtrans;
 mod dropout;
 mod ema;
 mod embedding;
@@ -208,6 +209,7 @@ mod safetensors;
 mod split_into;
 mod transformer;
 mod unbiased_linear;
+mod upscale;
 
 pub use module::{
     BuildModule, BuildOnDevice, DeviceBuildExt, Module, ModuleMut, NonMutableModule,
@@ -238,6 +240,7 @@ pub mod modules {
     pub use super::bias2d::Bias2D;
     #[cfg(feature = "nightly")]
     pub use super::conv::Conv2D;
+    pub use super::convtrans::ConvTrans2D;
     pub use super::dropout::{Dropout, DropoutOneIn};
     pub use super::embedding::Embedding;
     #[cfg(feature = "nightly")]
@@ -256,6 +259,7 @@ pub mod modules {
         TransformerEncoder, TransformerEncoderBlock,
     };
     pub use super::unbiased_linear::UnbiasedLinear;
+    pub use super::upscale::Upscale2D;
     pub use super::*;
 }
 
@@ -269,6 +273,7 @@ pub mod builders {
     pub use super::bias2d::builder::Bias2D;
     #[cfg(feature = "nightly")]
     pub use super::conv::builder::Conv2D;
+    pub use super::convtrans::builder::ConvTrans2D;
     pub use super::dropout::{Dropout, DropoutOneIn};
     pub use super::embedding::builder::Embedding;
     #[cfg(feature = "nightly")]
diff --git a/src/nn/upscale.rs b/src/nn/upscale.rs
new file mode 100644
index 000000000..3f85d71e4
--- /dev/null
+++ b/src/nn/upscale.rs
@@ -0,0 +1,112 @@
+#[cfg(feature = "nightly")]
+use crate::prelude::{Const, Dim, Dtype, HasErr, Tape, Tensor, Upscale2DKernel, ZerosTensor};
+use crate::prelude::{ConstUpscale2D, NearestNeighbor, UpscaleMethod};
+
+#[allow(unused)]
+use super::{BuildModule, Module, NonMutableModule, ZeroSizedModule};
+
+#[derive(Debug, Default, Clone)]
+pub struct Upscale2D<const OH: usize, const OW: usize = OH, M: UpscaleMethod = NearestNeighbor>(M);
+
+impl<const OH: usize, const OW: usize, M: UpscaleMethod> ZeroSizedModule for Upscale2D<OH, OW, M> {}
+impl<const OH: usize, const OW: usize, M: UpscaleMethod> NonMutableModule for Upscale2D<OH, OW, M> {}
+
+impl<const OH: usize, const OW: usize, M: UpscaleMethod, Img: ConstUpscale2D<M>> Module<Img>
+    for Upscale2D<OH, OW, M>
+{
+    type Output = Img::Output<OH, OW>;
+    type Error = Img::Err;
+
+    fn try_forward(&self, x: Img) -> Result<Self::Output, Img::Err> {
+        x.try_upscale2d()
+    }
+}
+
+#[cfg(feature = "nightly")]
+#[derive(Debug, Default, Clone)]
+pub struct Upscale2DBy<const H: usize, const W: usize = H, M: UpscaleMethod = NearestNeighbor>(M);
+
+#[cfg(feature = "nightly")]
+impl<const H: usize, const W: usize, M: UpscaleMethod> ZeroSizedModule for Upscale2DBy<H, W, M> {}
+#[cfg(feature = "nightly")]
+impl<const H: usize, const W: usize, M: UpscaleMethod> NonMutableModule for Upscale2DBy<H, W, M> {}
+
+#[cfg(feature = "nightly")]
+impl<
+        const H: usize,
+        const W: usize,
+        const IH: usize,
+        const IW: usize,
+        C: Dim,
+        E: Dtype,
+        M: UpscaleMethod,
+        D: Upscale2DKernel<E, M> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > Module<Tensor<(C, Const<IH>, Const<IW>), E, D, T>> for Upscale2DBy<H, W, M>
+where
+    Tensor<(C, Const<{ IH * H }>, Const<{ IW * W }>), E, D, T>: Sized,
+{
+    type Output = Tensor<(C, Const<{ IH * H }>, Const<{ IW * W }>), E, D, T>;
+    type Error = <Self::Output as HasErr>::Err;
+
+    fn try_forward(
+        &self,
+        x: Tensor<(C, Const<IH>, Const<IW>), E, D, T>,
+    ) -> Result<Self::Output, Self::Error> {
+        x.try_upscale2d()
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl<
+        const H: usize,
+        const W: usize,
+        const IH: usize,
+        const IW: usize,
+        B: Dim,
+        C: Dim,
+        E: Dtype,
+        M: UpscaleMethod,
+        D: Upscale2DKernel<E, M> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > Module<Tensor<(B, C, Const<IH>, Const<IW>), E, D, T>> for Upscale2DBy<H, W, M>
+where
+    Tensor<(B, C, Const<{ IH * H }>, Const<{ IW * W }>), E, D, T>: Sized,
+{
+    type Output = Tensor<(B, C, Const<{ IH * H }>, Const<{ IW * W }>), E, D, T>;
+    type Error = <Self::Output as HasErr>::Err;
+
+    fn try_forward(
+        &self,
+        x: Tensor<(B, C, Const<IH>, Const<IW>), E, D, T>,
+    ) -> Result<Self::Output, Self::Error> {
+        x.try_upscale2d()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{prelude::Bilinear, shapes::*, tensor::*, tests::*};
+
+    #[test]
+    fn test_upscale2d() {
+        let dev: TestDevice = Default::default();
+        let x: Tensor<Rank3<3, 4, 4>, TestDtype, _> = dev.zeros();
+        let _: Tensor<Rank3<3, 8, 8>, _, _> = Upscale2D::<8>::default().forward(x.clone());
+        let _: Tensor<Rank3<3, 8, 12>, _, _> = Upscale2D::<8, 12>::default().forward(x.clone());
+        let _: Tensor<Rank3<3, 9, 9>, _, _> =
+            Upscale2D::<9, 9, NearestNeighbor>::default().forward(x.clone());
+    }
+
+    #[cfg(feature = "nightly")]
+    #[test]
+    fn test_upscale2dby() {
+        let dev: TestDevice = Default::default();
+        let x: Tensor<Rank3<3, 4, 4>, TestDtype, _> = dev.zeros();
+        let _: Tensor<Rank3<3, 8, 8>, _, _> = Upscale2DBy::<2>::default().forward(x.clone());
+        let _: Tensor<Rank3<3, 8, 12>, _, _> = Upscale2DBy::<2, 3>::default().forward(x.clone());
+        let _: Tensor<Rank3<3, 12, 12>, _, _> =
+            Upscale2DBy::<3, 3, Bilinear>::default().forward(x.clone());
+    }
+}
diff --git a/src/tensor_ops/convtrans2d/convtrans2d.cu b/src/tensor_ops/convtrans2d/convtrans2d.cu
new file mode 100644
index 000000000..d2d98abd1
--- /dev/null
+++ b/src/tensor_ops/convtrans2d/convtrans2d.cu
@@ -0,0 +1,232 @@
+struct Conv2DOp {
+    size_t stride;
+    size_t padding;
+    size_t kernel;
+    size_t batch;
+    size_t chan_in;
+    size_t chan_out;
+    size_t h_in;
+    size_t h_out;
+    size_t w_in;
+    size_t w_out;
+};
+
+template<typename T>
+__device__ void unfold_input_into_patches(
+    const Conv2DOp op,
+    const T *image, // 4d (Batch, Channels, Height, Width)
+    const size_t *strides, // 4d image strides
+    T *patches // 6d (Batch, Channels, KernelSize, KernelSize, HeightOut, WidthOut)
+) {
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    const auto patches_numel = op.batch * op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out;
+    if (i >= patches_numel) {
+        return;
+    }
+
+    // patches shape is (B, C, K, K, h_out, w_out)
+    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 k2 = idx % op.kernel;
+    idx /= op.kernel;
+    const size_t k1 = idx % op.kernel;
+    idx /= op.kernel;
+    const size_t c = idx % op.chan_in;
+    idx /= op.chan_in;
+    const size_t b = idx % op.batch;
+    idx /= op.batch;
+
+    size_t y = oh + op.padding;
+    if (y < k1) {
+        return;
+    }
+    y -= k1;
+    if (y % op.stride != 0) {
+        return;
+    }
+    y /= op.stride;
+    if (y >= op.h_in) {
+        return;
+    }
+    
+    size_t x = ow + op.padding;
+    if (x < k2) {
+        return;
+    }
+    x -= k2;
+    if (x % op.stride != 0) {
+        return;
+    }
+    x /= op.stride;
+    if (x >= op.w_in) {
+        return;
+    }
+
+    const size_t i_image = b * strides[0] + c * strides[1] + y * strides[2] + x * strides[3];
+    patches[i] = image[i_image];
+}
+
+template<typename T>
+__device__ void unfold_output_into_patches(
+    const Conv2DOp op,
+    const T *image_out, // 4d (Batch, ChanOut, HeightOut, WidthOut)
+    T *patches // 6d (Batch, ChanOut, KernelSize, KernelSize, HeightIn, WidthIn)
+) {
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    const auto patches_numel = op.batch * op.chan_out * op.kernel * op.kernel * op.h_in * op.w_in;
+    if (i >= patches_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 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 b = idx % op.batch;
+    idx /= op.batch;
+
+    const size_t y_plus_p = y * op.stride + k1;
+    if (y_plus_p < op.padding) {
+        return;
+    }
+    const size_t oh = y_plus_p - op.padding;
+    if (oh >= op.h_out) {
+        return;
+    }
+
+    const size_t x_plus_p = x * op.stride + k2;
+    if (x_plus_p < op.padding) {
+        return;
+    }
+    const size_t ow = x_plus_p - op.padding;
+    if (ow >= op.w_out) {
+        return;
+    }
+
+    size_t image_i = b * (op.chan_out * op.h_out * op.w_out) + o * (op.h_out * op.w_out) + oh * (op.w_out)  + ow;
+    patches[i] = image_out[image_i];
+}
+
+template<typename T>
+__device__ void transpose_and_broadcast_filters(
+    const Conv2DOp op,
+    const T *filters, // 4d (ChanOut, ChanIn, KernelSize, KernelSize)
+    const size_t *strides, // 4d filters strides
+    T *filters_tr // 5d (Batch, ChanIn, ChanOut, KernelSize, KernelSize)
+) {
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    auto numel = op.chan_in * op.chan_out * op.kernel * op.kernel;
+    if (i >= numel) {
+        return;
+    }
+
+    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 c = idx % op.chan_in;
+    idx /= op.chan_in;
+    const size_t o = idx % op.chan_out;
+    idx /= op.chan_out;
+
+    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];
+
+    const T f = filters[i_no];
+    for (auto b = 0; b < op.batch; b++) {
+        filters_tr[b * numel + i_tr] = f;
+    }
+}
+
+template<typename T>
+__device__ void sum_transposed_filters(
+    const Conv2DOp op,
+    const T *filters_tr, // 5d (Batch, ChanIn, ChanOut, KernelSize, KernelSize)
+    T *filters, // 4d (ChanOut, ChanIn, KernelSize, KernelSize)
+    const size_t *strides // 4d filter strides
+) {
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    auto numel = op.chan_out * op.chan_in * op.kernel * op.kernel;
+    if (i >= numel) {
+        return;
+    }
+
+    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 c = idx % op.chan_in;
+    idx /= op.chan_in;
+    const size_t o = idx % op.chan_out;
+    idx /= op.chan_out;
+
+    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];
+
+    T tmp = 0.0;
+    for (auto b = 0; b < op.batch; b++) {
+        tmp += filters_tr[b * numel + i_tr];
+    }
+
+    filters[i_no] += tmp;
+}
+
+#define CONV_OP(TYPENAME, UNFOLD_INPUT, UNFOLD_OUTPUT, TR_FILTERS, SUM_TR_FILTERS) \
+extern "C" __global__ void UNFOLD_INPUT( \
+    const Conv2DOp op, \
+    const TYPENAME *image, \
+    const size_t *strides, \
+    TYPENAME *patches \
+) { \
+    unfold_input_into_patches(op, image, strides, patches); \
+} \
+extern "C" __global__ void UNFOLD_OUTPUT( \
+    const Conv2DOp op, \
+    const TYPENAME *image_out, \
+    TYPENAME *patches \
+) { \
+    unfold_output_into_patches(op, image_out, patches); \
+} \
+extern "C" __global__ void TR_FILTERS( \
+    const Conv2DOp op, \
+    const TYPENAME *filters, \
+    const size_t *strides, \
+    TYPENAME *filters_tr \
+) { \
+    transpose_and_broadcast_filters(op, filters, strides, filters_tr); \
+} \
+extern "C" __global__ void SUM_TR_FILTERS( \
+    const Conv2DOp op, \
+    const TYPENAME *filters_tr, \
+    TYPENAME *filters, \
+    const size_t *strides \
+) { \
+    sum_transposed_filters(op, filters_tr, filters, strides); \
+}
+
+CONV_OP(
+    float,
+    unfold_input_into_patches_f32,
+    unfold_output_into_patches_f32,
+    transpose_and_broadcast_filters_f32,
+    sum_transposed_filters_f32
+);
+CONV_OP(
+    double,
+    unfold_input_into_patches_f64,
+    unfold_output_into_patches_f64,
+    transpose_and_broadcast_filters_f64,
+    sum_transposed_filters_f64
+);
\ No newline at end of file
diff --git a/src/tensor_ops/convtrans2d/cpu_kernel.rs b/src/tensor_ops/convtrans2d/cpu_kernel.rs
new file mode 100644
index 000000000..36cdb4256
--- /dev/null
+++ b/src/tensor_ops/convtrans2d/cpu_kernel.rs
@@ -0,0 +1,261 @@
+use crate::shapes::{Dtype, Shape};
+use crate::tensor::{cpu::*, Tensor};
+use crate::tensor_ops::matmul::cpu_kernel::MatMulImpl;
+
+use std::sync::Arc;
+
+use super::{ConvTrans2DKernel, ConvTrans2DOp};
+
+impl ConvTrans2DOp {
+    #[inline(always)]
+    fn unfold_idx(&self, [k1, k2, y, x]: [usize; 4]) -> Option<[usize; 2]> {
+        let mut oh = y * self.stride;
+        oh += k1;
+        oh -= self.padding;
+
+        let mut ow = x * self.stride;
+        ow += k2;
+        ow -= self.padding;
+
+        Some([oh, ow])
+    }
+}
+
+impl Cpu {
+    #[inline]
+    fn convtrans2d_forward<E: Dtype>(
+        &self,
+        op: &ConvTrans2DOp,
+        img: &[E],
+        filters: &[E],
+        out: &mut [E],
+        buf: &mut [E],
+    ) -> Result<(), CpuError>
+    where
+        Self: MatMulImpl<E>,
+    {
+        {
+            let mut i = 0;
+            for c in 0..op.chan_in {
+                for k1 in 0..op.kernel {
+                    for k2 in 0..op.kernel {
+                        for oh in 0..op.h_out {
+                            for ow in 0..op.w_out {
+                                i += 1;
+                                let mut y = oh + op.padding;
+                                if y < k1 {
+                                    continue;
+                                }
+                                y -= k1;
+                                if y % op.stride != 0 {
+                                    continue;
+                                }
+                                y /= op.stride;
+                                if y >= op.h_in {
+                                    continue;
+                                }
+
+                                let mut x = ow + op.padding;
+                                if x < k2 {
+                                    continue;
+                                }
+                                x -= k2;
+                                if x % op.stride != 0 {
+                                    continue;
+                                }
+                                x /= op.stride;
+                                if x >= op.w_in {
+                                    continue;
+                                }
+
+                                if y < op.h_in && x < op.w_in {
+                                    buf[i - 1] = img[c * (op.w_in * op.h_in) + y * op.w_in + x];
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        // (O, C * K * K) * (C * K * K, OH * OW) = (O, OH * OW)
+        let m = op.chan_out;
+        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],
+        );
+        Ok(())
+    }
+
+    #[inline]
+    #[allow(clippy::too_many_arguments)]
+    fn convtrans2d_backward<E: Dtype>(
+        &self,
+        op: &ConvTrans2DOp,
+        img: &[E],
+        grad_img: &mut [E],
+        filters_tr: &[E],
+        grad_filters_tr: &mut [E],
+        grad_out: &[E],
+        buf: &mut [E],
+    ) -> Result<(), CpuError>
+    where
+        Self: MatMulImpl<E>,
+    {
+        {
+            let mut i = 0;
+            for o in 0..op.chan_out {
+                for k1 in 0..op.kernel {
+                    for k2 in 0..op.kernel {
+                        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]) {
+                                    buf[i] =
+                                        grad_out[o * (op.h_out * op.w_out) + oh * op.w_out + ow];
+                                }
+                                i += 1;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        {
+            // img_g += filters^T * unfold(grad_out)
+            // (C, H * W) += (C, O * K * K) * (O * K * K, H * W)
+            let m = op.chan_in;
+            let k = op.chan_out * 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],
+            );
+        }
+
+        {
+            // weight_g^T += img * patches^T
+            // (C, O * K * K) += (C, H * W) * (H * W, O * 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],
+            );
+        }
+        Ok(())
+    }
+}
+
+impl<E: Dtype> ConvTrans2DKernel<E> for Cpu
+where
+    Self: MatMulImpl<E>,
+{
+    fn forward<L: Shape, R: Shape, O: Shape>(
+        &self,
+        op: ConvTrans2DOp,
+        lhs: &Tensor<L, E, Self>,
+        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 [lstride, ostride] = match L::NUM_DIMS {
+            3 => [0; 2],
+            4 => [lhs.strides[0], out.strides[0]],
+            _ => unreachable!(),
+        };
+        let lhs = lhs.data.as_ref();
+        let rhs = rhs.data.as_ref();
+        let out = Arc::make_mut(&mut out.data);
+        for i_batch in 0..op.batch {
+            self.convtrans2d_forward(
+                &op,
+                &lhs[i_batch * lstride..],
+                rhs,
+                &mut out[i_batch * ostride..],
+                &mut patches,
+            )?;
+        }
+        Ok(())
+    }
+
+    fn backward<L: Shape, R: Shape, O: Shape>(
+        &self,
+        op: ConvTrans2DOp,
+        lhs: &Tensor<L, E, Self>,
+        grad_lhs: &mut Self::Vec<E>,
+        rhs: &Tensor<R, E, Self>,
+        grad_rhs: &mut Self::Vec<E>,
+        out: &Tensor<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 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())?;
+
+        {
+            // 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]
+                    + c * rhs.strides[1]
+                    + k1 * rhs.strides[2]
+                    + k2 * rhs.strides[3];
+                f1023[i] = buf[idx];
+            }
+        }
+
+        let [lstride, ostride] = match L::NUM_DIMS {
+            3 => [0; 2],
+            4 => [lhs.strides[0], out.strides[0]],
+            _ => unreachable!(),
+        };
+        let lhs = lhs.data.as_ref();
+
+        for i_batch in 0..op.batch {
+            self.convtrans2d_backward(
+                &op,
+                &lhs[i_batch * lstride..],
+                &mut grad_lhs[i_batch * lstride..],
+                &f1023,
+                &mut grad_f1023,
+                &grad_out[i_batch * ostride..],
+                &mut patches,
+            )?;
+        }
+
+        {
+            // 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]
+                    + c * rhs.strides[1]
+                    + k1 * rhs.strides[2]
+                    + k2 * rhs.strides[3];
+                grad_rhs[idx] += grad_f1023[i];
+            }
+        }
+
+        Ok(())
+    }
+}
diff --git a/src/tensor_ops/convtrans2d/cuda_kernel.rs b/src/tensor_ops/convtrans2d/cuda_kernel.rs
new file mode 100644
index 000000000..a7bf44a95
--- /dev/null
+++ b/src/tensor_ops/convtrans2d/cuda_kernel.rs
@@ -0,0 +1,176 @@
+use cudarc::cublas::{CudaBlas, Gemm};
+use cudarc::driver::{DeviceRepr, DeviceSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits};
+
+use crate::{
+    shapes::*,
+    tensor::{Cuda, Tensor},
+};
+
+use std::sync::Arc;
+
+unsafe impl DeviceRepr for super::ConvTrans2DOp {}
+
+const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/convtrans2d.ptx"));
+
+trait HasCudaKernel<E> {
+    const MOD: &'static str;
+    const FNS: &'static [&'static str];
+}
+
+impl HasCudaKernel<f32> for Cuda {
+    const MOD: &'static str = "convtrans2d_f32";
+    const FNS: &'static [&'static str] = &[
+        "unfold_input_into_patches_f32",
+        "unfold_output_into_patches_f32",
+        "transpose_and_broadcast_filters_f32",
+        "sum_transposed_filters_f32",
+    ];
+}
+
+impl HasCudaKernel<f64> for Cuda {
+    const MOD: &'static str = "convtrans2d_f64";
+    const FNS: &'static [&'static str] = &[
+        "unfold_input_into_patches_f64",
+        "unfold_output_into_patches_f64",
+        "transpose_and_broadcast_filters_f64",
+        "sum_transposed_filters_f64",
+    ];
+}
+
+fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
+    match S::NUM_DIMS {
+        3 => [0, strides[0], strides[1], strides[2]],
+        4 => [strides[0], strides[1], strides[2], strides[3]],
+        _ => unreachable!("Only implemented for 3d & 4d arrays"),
+    }
+}
+
+impl<E: Dtype + ValidAsZeroBits> super::ConvTrans2DKernel<E> for Cuda
+where
+    Self: HasCudaKernel<E>,
+    CudaBlas: Gemm<E>,
+{
+    fn forward<L: Shape, R: Shape, O: Shape>(
+        &self,
+        op: super::ConvTrans2DOp,
+        lhs: &Tensor<L, E, Self>,
+        rhs: &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_numel = op.batch * op.chan_in * op.kernel * op.kernel * op.h_out * op.w_out;
+        let mut patches = self.dev.alloc_zeros::<E>(patches_numel)?;
+        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 = LaunchConfig::for_num_elems(patches.len() as u32);
+        let params = (op, lhs.data.as_ref(), &img_strides, &mut patches);
+        unsafe { unfold_fn.launch(cfg, params) }?;
+
+        // (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::make_mut(&mut out.data),
+                [m * n, n, 1],
+            )
+            .unwrap();
+        }
+
+        Ok(())
+    }
+
+    fn backward<L: Shape, R: Shape, O: Shape>(
+        &self,
+        op: super::ConvTrans2DOp,
+        lhs: &Tensor<L, E, Self>,
+        grad_lhs: &mut Self::Vec<E>,
+        rhs: &Tensor<R, E, Self>,
+        grad_rhs: &mut Self::Vec<E>,
+        _: &Tensor<O, E, Self>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err> {
+        let patches_numel = op.batch * op.chan_out * op.kernel * op.kernel * op.h_in * op.w_in;
+        let mut patches = self.dev.alloc_zeros::<E>(patches_numel)?;
+
+        {
+            // unfold grad_out into patches
+            let unfold_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
+            let cfg = LaunchConfig::for_num_elems(patches_numel as u32);
+            unsafe { unfold_fn.launch(cfg, (op, grad_out, &mut patches)) }?;
+        }
+
+        let filters_numel = op.batch * op.chan_in * op.chan_out * op.kernel * op.kernel;
+        let mut f_b1023 = self.dev.alloc_zeros::<E>(filters_numel)?;
+        let mut grad_f_b1023 = self.dev.alloc_zeros::<E>(filters_numel)?;
+        let f_strides = self.dev.htod_copy(rhs.strides.into())?;
+
+        {
+            // prepare filters for backward operations by
+            // swapping dims 0 and 1 and adding a batch dimension
+            let tr_fn = self.dev.get_func(Self::MOD, Self::FNS[2]).unwrap();
+            let cfg = LaunchConfig::for_num_elems(rhs.shape.num_elements() as u32);
+            unsafe { tr_fn.launch(cfg, (op, rhs.data.as_ref(), &f_strides, &mut f_b1023)) }?;
+        }
+
+        {
+            // img_g += filters * patches
+            // (B, C, H * W) += (B, C, O * K * K) * (B, O * K * K, H * W)
+            let m = op.chan_in;
+            let k = op.chan_out * op.kernel * op.kernel;
+            let n = op.h_in * op.w_in;
+            unsafe {
+                self.gemm_batch(
+                    (op.batch, m, k, n),
+                    &f_b1023,
+                    [m * k, k, 1],
+                    &patches,
+                    [k * n, n, 1],
+                    <E>::ONE,
+                    grad_lhs,
+                    [m * n, n, 1],
+                )
+                .unwrap();
+            }
+        }
+
+        {
+            // weight_g += img * patches^T
+            // (B, C, O * K * K) += (B, C, H * W) * (B, H * W, O * 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 {
+                self.gemm_batch(
+                    (op.batch, m, k, n),
+                    lhs.data.as_ref(),
+                    [m * k, k, 1],
+                    &patches,
+                    [k * n, 1, k],
+                    <E>::ONE,
+                    &mut grad_f_b1023,
+                    [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 = LaunchConfig::for_num_elems(rhs.shape.num_elements() as u32);
+            unsafe { sum_fn.launch(cfg, (op, &grad_f_b1023, grad_rhs, &f_strides)) }?;
+        }
+
+        Ok(())
+    }
+}
diff --git a/src/tensor_ops/convtrans2d/mod.rs b/src/tensor_ops/convtrans2d/mod.rs
new file mode 100644
index 000000000..6ea55d052
--- /dev/null
+++ b/src/tensor_ops/convtrans2d/mod.rs
@@ -0,0 +1,584 @@
+mod cpu_kernel;
+
+#[cfg(feature = "cuda")]
+mod cuda_kernel;
+
+use crate::{
+    shapes::*,
+    tensor::{DeviceStorage, HasErr, PutTape, SplitTape, Tape, Tensor, ZerosTensor},
+};
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub(super) struct ConvTrans2DOp {
+    pub stride: usize,
+    pub padding: usize,
+    pub kernel: usize,
+    pub batch: usize,
+    pub chan_in: usize,
+    pub chan_out: usize,
+    pub h_in: usize,
+    pub h_out: usize,
+    pub w_in: usize,
+    pub w_out: usize,
+}
+
+impl ConvTrans2DOp {
+    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 - 1) * s - 2 * p + k,
+            w_in,
+            w_out: (w_in - 1) * s - 2 * p + k,
+        }
+    }
+
+    #[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 ConvTrans2DKernel<E: Dtype>: DeviceStorage {
+    fn forward<L: Shape, R: Shape, O: Shape>(
+        &self,
+        op: ConvTrans2DOp,
+        lhs: &Tensor<L, E, Self>,
+        rhs: &Tensor<R, E, Self>,
+        out: &mut Tensor<O, E, Self>,
+    ) -> Result<(), Self::Err>;
+
+    #[allow(clippy::too_many_arguments)]
+    fn backward<L: Shape, R: Shape, O: Shape>(
+        &self,
+        op: ConvTrans2DOp,
+        lhs: &Tensor<L, E, Self>,
+        grad_lhs: &mut Self::Vec<E>,
+        rhs: &Tensor<R, E, Self>,
+        grad_rhs: &mut Self::Vec<E>,
+        out: &Tensor<O, E, Self>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err>;
+}
+
+pub trait ConvTransAlgebra<const K: usize, const S: usize, const P: usize>: ConstDim {
+    type Convolved: ConstDim;
+}
+
+impl<const D: usize, const K: usize, const S: usize, const P: usize> ConvTransAlgebra<K, S, P>
+    for Const<D>
+where
+    Const<{ D * S + K - S - 2 * P }>: Sized,
+{
+    type Convolved = Const<{ D * S + K - S - 2 * P }>;
+}
+
+pub trait TryConvTrans2DTo<F, const S: usize, const P: usize>: HasErr {
+    type Output;
+    fn convtrans2d_to(self, filters: F) -> Self::Output {
+        self.try_convtrans2d_to(filters).unwrap()
+    }
+    fn try_convtrans2d_to(self, filters: F) -> Result<Self::Output, Self::Err>;
+}
+
+pub trait TryConvTrans2D<F> {
+    fn convtrans2d<const S: usize, const P: usize>(self, filters: F) -> Self::Output
+    where
+        Self: TryConvTrans2DTo<F, S, P>,
+    {
+        self.convtrans2d_to(filters)
+    }
+    fn try_convtrans2d<const S: usize, const P: usize>(
+        self,
+        filters: F,
+    ) -> Result<Self::Output, Self::Err>
+    where
+        Self: TryConvTrans2DTo<F, S, P>,
+    {
+        self.try_convtrans2d_to(filters)
+    }
+}
+
+impl<T, F> TryConvTrans2D<F> for T {}
+
+impl<
+        const C: usize,
+        const H: usize,
+        const W: usize,
+        const O: usize,
+        const K: usize,
+        const S: usize,
+        const P: usize,
+        E: Dtype,
+        D: ConvTrans2DKernel<E> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > TryConvTrans2DTo<Tensor<Rank4<O, C, K, K>, E, D>, S, P> for Tensor<Rank3<C, H, W>, E, D, T>
+where
+    Const<H>: ConvTransAlgebra<K, S, P>,
+    Const<W>: ConvTransAlgebra<K, S, P>,
+{
+    type Output = Tensor<
+        (
+            Const<O>,
+            <Const<H> as ConvTransAlgebra<K, S, P>>::Convolved,
+            <Const<W> as ConvTransAlgebra<K, S, P>>::Convolved,
+        ),
+        E,
+        D,
+        T,
+    >;
+
+    fn try_convtrans2d_to(
+        self,
+        filters: Tensor<Rank4<O, C, K, K>, E, D>,
+    ) -> Result<Self::Output, Self::Err> {
+        let op = ConvTrans2DOp::new(S, P, K, [1, C, H, W], O);
+        let (lhs, ltape) = self.split_tape();
+        let (rhs, rtape) = filters.split_tape();
+        let mut tape = ltape.merge(rtape);
+        let mut out = lhs.device.try_zeros()?;
+        lhs.device.forward(op, &lhs, &rhs, &mut out)?;
+        let phantom_out = out.clone();
+        tape.try_alloc_grad(&lhs)?;
+        tape.try_alloc_grad(&rhs)?;
+        tape.try_alloc_grad(&out)?;
+        tape.add_backward_op(move |grads| {
+            let (grad_lhs, grad_rhs, grad_out) = grads.muts_and_ref(&lhs, &rhs, &phantom_out);
+            lhs.device
+                .backward(op, &lhs, grad_lhs, &rhs, grad_rhs, &phantom_out, grad_out)
+        });
+        Ok(out.put_tape(tape))
+    }
+}
+
+impl<
+        B: Dim,
+        const C: usize,
+        const H: usize,
+        const W: usize,
+        const O: usize,
+        const K: usize,
+        const S: usize,
+        const P: usize,
+        E: Dtype,
+        D: ConvTrans2DKernel<E> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > TryConvTrans2DTo<Tensor<Rank4<O, C, K, K>, E, D>, S, P>
+    for Tensor<(B, Const<C>, Const<H>, Const<W>), E, D, T>
+where
+    Const<H>: ConvTransAlgebra<K, S, P>,
+    Const<W>: ConvTransAlgebra<K, S, P>,
+{
+    type Output = Tensor<
+        (
+            B,
+            Const<O>,
+            <Const<H> as ConvTransAlgebra<K, S, P>>::Convolved,
+            <Const<W> as ConvTransAlgebra<K, S, P>>::Convolved,
+        ),
+        E,
+        D,
+        T,
+    >;
+    fn try_convtrans2d_to(
+        self,
+        filters: Tensor<Rank4<O, C, K, K>, E, D>,
+    ) -> Result<Self::Output, Self::Err> {
+        let batch = self.shape().0;
+        let op = ConvTrans2DOp::new(S, P, K, [batch.size(), C, H, W], O);
+        let (lhs, ltape) = self.split_tape();
+        let (rhs, rtape) = filters.split_tape();
+        let mut out =
+            lhs.device
+                .try_zeros_like(&(batch, Const, Default::default(), Default::default()))?;
+        let mut tape = ltape.merge(rtape);
+        lhs.device.forward(op, &lhs, &rhs, &mut out)?;
+        let phantom_out = out.clone();
+        tape.try_alloc_grad(&lhs)?;
+        tape.try_alloc_grad(&rhs)?;
+        tape.try_alloc_grad(&out)?;
+        tape.add_backward_op(move |grads| {
+            let (grad_lhs, grad_rhs, grad_out) = grads.muts_and_ref(&lhs, &rhs, &phantom_out);
+            lhs.device
+                .backward(op, &lhs, grad_lhs, &rhs, grad_rhs, &phantom_out, grad_out)?;
+            Ok(())
+        });
+        Ok(out.put_tape(tape))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{tensor::*, tensor_ops::*, tests::*};
+
+    #[test]
+    /// TODO
+    /// Produced by
+    /// ```python
+    /// x = torch.rand((2, 3, 3), requires_grad=True)
+    /// w = torch.rand((2, 3, 3, 3), requires_grad=True)
+    /// print(x)
+    /// print(torch.swapaxes(w, 0, 1))
+    /// y = torch.conv_transpose2d(x, w)
+    /// print(y)
+    /// y.exp().mean().backward()
+    /// print(x.grad)
+    /// print(torch.swapaxes(w.grad, 0, 1))
+    /// ```
+    fn convtrans2d_test() {
+        let device = TestDevice::default();
+
+        let x = device.tensor([
+            [
+                [0.0907329, 0.5784497, 0.1818193],
+                [0.9867508, 0.0566732, 0.9057426],
+                [0.3095418, 0.7836370, 0.3519793],
+            ],
+            [
+                [0.0969202, 0.6929486, 0.7536632],
+                [0.8163304, 0.4960053, 0.4525285],
+                [0.2100813, 0.2163504, 0.3710884],
+            ],
+        ]);
+        let w = device.tensor([
+            [
+                [
+                    [0.3401043, 0.4955706, 0.2441649],
+                    [0.4701799, 0.6495003, 0.2529446],
+                    [0.4017784, 0.1768293, 0.2353096],
+                ],
+                [
+                    [0.6621207, 0.7709801, 0.5986264],
+                    [0.2237560, 0.6466616, 0.7321741],
+                    [0.1578425, 0.9478191, 0.9861543],
+                ],
+            ],
+            [
+                [
+                    [0.8291070, 0.0848221, 0.3680936],
+                    [0.4642293, 0.1073243, 0.1073309],
+                    [0.7863810, 0.3699800, 0.4956312],
+                ],
+                [
+                    [0.0681600, 0.5616951, 0.4053129],
+                    [0.1850831, 0.8223089, 0.0667553],
+                    [0.8905262, 0.6328429, 0.8180532],
+                ],
+            ],
+            [
+                [
+                    [0.7582999, 0.9763424, 0.5727801],
+                    [0.3743349, 0.4793805, 0.6885015],
+                    [0.8183323, 0.1882774, 0.9794642],
+                ],
+                [
+                    [0.5606869, 0.7552301, 0.6572021],
+                    [0.8761331, 0.2401637, 0.1778120],
+                    [0.2065960, 0.4133974, 0.8821540],
+                ],
+            ],
+        ]);
+        let y = x.leaky_trace().convtrans2d::<1, 0>(w.clone());
+        assert_close(
+            &y.array(),
+            &[
+                [
+                    [0.0950315, 0.7752370, 1.4619386, 1.2272182, 0.4955567],
+                    [0.9404547, 2.0147018, 2.9196219, 2.3676410, 1.0898490],
+                    [0.9427375, 2.4812443, 3.9218845, 3.6057489, 1.6545289],
+                    [0.7178541, 1.8030399, 3.0822182, 1.9167527, 1.0201255],
+                    [0.1575270, 0.6028528, 0.8236330, 0.8117172, 0.4487746],
+                ],
+                [
+                    [0.0818333, 0.5889638, 0.7130897, 0.9325359, 0.3723960],
+                    [0.9338222, 1.1092455, 2.6313026, 1.3005096, 0.5866395],
+                    [1.0377907, 2.8708880, 3.3737209, 2.5207422, 1.1140431],
+                    [1.6855066, 1.9777625, 3.1483138, 1.4968101, 0.8816571],
+                    [0.4305007, 1.0563757, 1.3593760, 0.9304471, 0.4780220],
+                ],
+                [
+                    [0.1231446, 0.9889490, 1.7642095, 1.5334388, 0.5994515],
+                    [1.3248383, 2.7914243, 3.7239599, 2.3741565, 1.0753872],
+                    [1.5313777, 2.9749527, 4.2994099, 3.4086916, 1.9924896],
+                    [1.2760720, 1.3538387, 3.8719988, 1.6865263, 1.5946647],
+                    [0.2967100, 0.8310994, 1.0901904, 1.1780756, 0.6721083],
+                ],
+            ],
+        );
+
+        let g = y.exp().mean().backward();
+
+        assert_close(
+            &g.get(&x).array(),
+            &[
+                [
+                    [2.4066830, 2.4581399, 2.5645943],
+                    [3.0028410, 3.2547507, 3.4216807],
+                    [2.1464431, 3.0581608, 2.8662176],
+                ],
+                [
+                    [2.9864695, 3.5932014, 2.5797451],
+                    [3.3677268, 3.8909531, 3.2242548],
+                    [2.4629762, 3.3527191, 2.9590628],
+                ],
+            ],
+        );
+        assert_close(
+            &g.get(&w).array(),
+            &[
+                [
+                    [
+                        [0.6642238, 1.0354118, 0.9408946],
+                        [0.9326871, 1.1026800, 1.0413336],
+                        [0.5472590, 0.7134937, 0.7234858],
+                    ],
+                    [
+                        [0.5456561, 0.7068147, 0.6173539],
+                        [0.8016681, 1.0878984, 1.0714644],
+                        [0.8350498, 1.1260045, 0.7775879],
+                    ],
+                ],
+                [
+                    [
+                        [0.5567597, 0.5549879, 0.4975571],
+                        [0.6702054, 0.8184303, 0.6338357],
+                        [0.6227797, 0.5077031, 0.5278049],
+                    ],
+                    [
+                        [0.3733614, 0.3889205, 0.3597363],
+                        [0.6457027, 0.7389204, 0.5783513],
+                        [0.7389930, 0.7089815, 0.5071381],
+                    ],
+                ],
+                [
+                    [
+                        [1.1678052, 1.4273405, 1.2900156],
+                        [1.4765850, 1.5869446, 1.4983673],
+                        [1.0089380, 0.8733283, 1.0910161],
+                    ],
+                    [
+                        [0.9175905, 1.0371233, 0.9381008],
+                        [1.4550014, 1.5706275, 1.4026034],
+                        [1.3066854, 1.4330946, 1.0638479],
+                    ],
+                ],
+            ],
+        );
+    }
+
+    #[test]
+    /// torch.set_printoptions(precision=7)
+    /// x = torch.rand((2, 3, 3), requires_grad=True)
+    /// w = torch.rand((2, 3, 3, 3), requires_grad=True)
+    /// print(x)
+    /// print(torch.swapaxes(w, 0, 1))
+    /// y = torch.conv_transpose2d(x, w, stride=2)
+    /// print(y)
+    /// y.exp().mean().backward()
+    /// print(x.grad)
+    /// print(torch.swapaxes(w.grad, 0, 1))
+    fn convtrans2d_s2() {
+        let device = TestDevice::default();
+
+        let x = device.tensor([
+            [
+                [0.0357635, 0.0225288, 0.3642959],
+                [0.6850907, 0.2586224, 0.2234361],
+                [0.9315249, 0.7850553, 0.7588840],
+            ],
+            [
+                [0.1240514, 0.4712945, 0.8732865],
+                [0.1023245, 0.1211519, 0.0407664],
+                [0.5106173, 0.9263544, 0.3101138],
+            ],
+        ]);
+        let w = device.tensor([
+            [
+                [
+                    [0.7007528, 0.1896583, 0.9991148],
+                    [0.6587640, 0.9383754, 0.3999129],
+                    [0.0035173, 0.4376699, 0.3985791],
+                ],
+                [
+                    [0.2180834, 0.4829719, 0.0272914],
+                    [0.2712103, 0.8577049, 0.8002768],
+                    [0.7074867, 0.4011419, 0.8835942],
+                ],
+            ],
+            [
+                [
+                    [0.0067961, 0.4006048, 0.3549793],
+                    [0.5392876, 0.3803764, 0.6090584],
+                    [0.4874769, 0.5006863, 0.8963661],
+                ],
+                [
+                    [0.3751084, 0.5425243, 0.5102475],
+                    [0.6024926, 0.2719866, 0.9794098],
+                    [0.2236674, 0.1083973, 0.4948432],
+                ],
+            ],
+            [
+                [
+                    [0.9486710, 0.9823384, 0.5994584],
+                    [0.2740490, 0.2620903, 0.2716798],
+                    [0.3620688, 0.9108542, 0.9017550],
+                ],
+                [
+                    [0.6089512, 0.4252676, 0.2729263],
+                    [0.8855131, 0.3937372, 0.3419960],
+                    [0.8216078, 0.6664743, 0.5395248],
+                ],
+            ],
+        ]);
+        let y = x.leaky_trace().convtrans2d::<2, 0>(w.clone());
+        assert_close(
+            &y.array(),
+            &[
+                [
+                    [
+                        0.0521149, 0.0666962, 0.1576860, 0.2318948, 0.4811018, 0.4908646, 0.3878066,
+                    ],
+                    [
+                        0.0572037, 0.1399591, 0.2562388, 0.4253721, 0.8630048, 1.0908685, 0.8445575,
+                    ],
+                    [
+                        0.5902850, 0.2447678, 1.3523079, 0.3064790, 1.4716963, 0.5718186, 1.1411824,
+                    ],
+                    [
+                        0.4790645, 0.7306364, 0.5590932, 0.3465975, 0.3586293, 0.2446325, 0.1219794,
+                    ],
+                    [
+                        0.8389287, 0.7641755, 2.1468871, 0.7580858, 1.6488208, 0.4078493, 0.8917535,
+                    ],
+                    [
+                        0.7521397, 1.3120791, 1.5495670, 1.5312154, 1.6393251, 0.9781042, 0.5516644,
+                    ],
+                    [
+                        0.3645314, 0.6125304, 1.4806095, 0.7151946, 1.3534986, 0.4565403, 0.5764900,
+                    ],
+                ],
+                [
+                    [
+                        0.0467758, 0.0816279, 0.2529318, 0.2647139, 0.5785270, 0.6197178, 0.5749097,
+                    ],
+                    [
+                        0.0940269, 0.0473439, 0.4393802, 0.1367552, 1.1979207, 0.3760918, 1.0771828,
+                    ],
+                    [
+                        0.0882189, 0.3613173, 0.5524452, 0.2317001, 0.7967559, 0.3886862, 0.8587984,
+                    ],
+                    [
+                        0.4311106, 0.2884232, 0.7299427, 0.1313255, 0.4212312, 0.0960777, 0.1760126,
+                    ],
+                    [
+                        0.5547202, 1.0043029, 1.7619288, 0.9596879, 1.2826418, 0.5885472, 0.6480764,
+                    ],
+                    [
+                        0.8100029, 0.4932112, 2.0489488, 0.5505725, 1.9815230, 0.3730083, 0.7659331,
+                    ],
+                    [
+                        0.5683053, 0.5217513, 1.6775545, 0.4934808, 1.6013979, 0.4135783, 0.8336955,
+                    ],
+                ],
+                [
+                    [
+                        0.1094690, 0.0878869, 0.3636634, 0.2225572, 1.0195196, 0.7292423, 0.4567231,
+                    ],
+                    [
+                        0.1196501, 0.0582169, 0.4756527, 0.1914707, 1.0404429, 0.4393238, 0.3976323,
+                    ],
+                    [
+                        0.8271067, 0.8317585, 1.2522885, 0.6402028, 1.5488806, 1.1506698, 0.9447322,
+                    ],
+                    [
+                        0.2783581, 0.2198445, 0.3992766, 0.1154844, 0.2090276, 0.0746117, 0.0746450,
+                    ],
+                    [
+                        1.5267723, 1.8244361, 2.8728042, 1.4814504, 2.0451818, 1.1080496, 0.7630367,
+                    ],
+                    [
+                        0.7074417, 0.4451926, 1.4631481, 0.5704955, 1.0126743, 0.3209994, 0.3122311,
+                    ],
+                    [
+                        0.7568032, 1.1887968, 2.1608419, 1.3324623, 1.7372788, 0.8979155, 0.8516415,
+                    ],
+                ],
+            ],
+        );
+
+        let g = y.exp().mean().backward();
+
+        assert_close(
+            &g.get(&x).array(),
+            &[
+                [
+                    [0.1513395, 0.1986136, 0.2298895],
+                    [0.3779295, 0.3469064, 0.2452929],
+                    [0.4825282, 0.5639746, 0.3148936],
+                ],
+                [
+                    [0.1527605, 0.2144486, 0.2480491],
+                    [0.3177541, 0.3597765, 0.2327209],
+                    [0.3772507, 0.5048490, 0.2865718],
+                ],
+            ],
+        );
+        assert_close(
+            &g.get(&w).array(),
+            &[
+                [
+                    [
+                        [0.1134962, 0.0483045, 0.1292279],
+                        [0.0842974, 0.0839551, 0.0851499],
+                        [0.0981565, 0.0517171, 0.1198249],
+                    ],
+                    [
+                        [0.0928453, 0.0412302, 0.0897282],
+                        [0.0699945, 0.0741750, 0.0777097],
+                        [0.0907740, 0.0422520, 0.0901646],
+                    ],
+                ],
+                [
+                    [
+                        [0.0771443, 0.0567608, 0.0866109],
+                        [0.1149883, 0.0410739, 0.1213422],
+                        [0.0952494, 0.0514846, 0.1152932],
+                    ],
+                    [
+                        [0.0687711, 0.0483170, 0.0680406],
+                        [0.1007998, 0.0350681, 0.1096105],
+                        [0.0770078, 0.0379331, 0.0848609],
+                    ],
+                ],
+                [
+                    [
+                        [0.1948610, 0.1028565, 0.1976164],
+                        [0.0683245, 0.0401628, 0.0643963],
+                        [0.1662624, 0.1036348, 0.2046718],
+                    ],
+                    [
+                        [0.1715550, 0.0769297, 0.1411840],
+                        [0.0654903, 0.0355645, 0.0568618],
+                        [0.1351082, 0.0760437, 0.1234931],
+                    ],
+                ],
+            ],
+        );
+    }
+}
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
index 8b99e110a..4c16d56fe 100644
--- a/src/tensor_ops/mod.rs
+++ b/src/tensor_ops/mod.rs
@@ -258,9 +258,17 @@ pub(crate) use to_dtype::ToDtypeKernel;
 #[cfg(feature = "nightly")]
 mod conv2d;
 #[cfg(feature = "nightly")]
-pub use conv2d::TryConv2D;
+pub use conv2d::{TryConv2D, TryConv2DTo};
+
+#[cfg(feature = "nightly")]
+mod convtrans2d;
+#[cfg(feature = "nightly")]
+pub use convtrans2d::{ConvTransAlgebra, TryConvTrans2D, TryConvTrans2DTo};
+
+mod upscale2d;
 #[cfg(feature = "nightly")]
-pub(crate) use conv2d::TryConv2DTo;
+pub(crate) use upscale2d::Upscale2DKernel;
+pub use upscale2d::{Bilinear, ConstUpscale2D, NearestNeighbor, TryUpscale2D, UpscaleMethod};
 
 #[cfg(feature = "nightly")]
 mod pool2d;
diff --git a/src/tensor_ops/upscale2d/cpu_kernel.rs b/src/tensor_ops/upscale2d/cpu_kernel.rs
new file mode 100644
index 000000000..4e13423d9
--- /dev/null
+++ b/src/tensor_ops/upscale2d/cpu_kernel.rs
@@ -0,0 +1,197 @@
+use crate::shapes::*;
+use crate::tensor::{Cpu, Tensor};
+
+use std::sync::Arc;
+
+use num_traits::Float;
+
+use super::{Bilinear, NearestNeighbor};
+
+fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
+    match S::NUM_DIMS {
+        3 => [0, strides[0], strides[1], strides[2]],
+        4 => [strides[0], strides[1], strides[2], strides[3]],
+        _ => panic!("Only implemented for 3d & 4d arrays"),
+    }
+}
+
+impl<E: Float + Unit + std::ops::AddAssign + std::ops::DivAssign>
+    super::Upscale2DKernel<E, NearestNeighbor> for Cpu
+{
+    fn forward<I: Shape, O: Shape>(
+        &self,
+        op: super::Upscale2DOp,
+        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 h_scale = ((istr[1] / istr[2]) as f32) / ((ostr[1] / ostr[2]) as f32);
+        let w_scale = ((istr[2] / istr[3]) as f32) / ((ostr[2] / ostr[3]) as f32);
+
+        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 ih = (h_scale * oh as f32) as usize;
+                        let iw = (w_scale * ow as f32) as usize;
+                        out_buf[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]] =
+                            buf[b * istr[0] + c * istr[1] + ih * istr[2] + iw * istr[3]];
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn backward<I: Shape, O: Shape>(
+        &self,
+        op: super::Upscale2DOp,
+        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 h_scale = ((istr[1] / istr[2]) as f32) / ((ostr[1] / ostr[2]) as f32);
+        let w_scale = ((istr[2] / istr[3]) as f32) / ((ostr[2] / ostr[3]) as f32);
+
+        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 ih = (h_scale * oh as f32) as usize;
+                        let iw = (w_scale * ow as f32) as usize;
+                        grad_inp[b * istr[0] + c * istr[1] + ih * istr[2] + iw * istr[3]] +=
+                            grad_out[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]];
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+impl<E: Float + Unit + std::ops::AddAssign + std::ops::DivAssign>
+    super::Upscale2DKernel<E, Bilinear> for Cpu
+{
+    fn forward<I: Shape, O: Shape>(
+        &self,
+        op: super::Upscale2DOp,
+        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 h_scale = ((istr[1] / istr[2] - 1) as f32) / ((ostr[1] / ostr[2] - 1) as f32);
+        let w_scale = ((istr[2] / istr[3] - 1) as f32) / ((ostr[2] / ostr[3] - 1) as f32);
+
+        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 ih = (h_scale * oh as f32) as usize;
+                        let iw = (w_scale * ow as f32) as usize;
+
+                        let hs = E::from(h_scale * oh as f32 - (ih as f32)).unwrap();
+                        let ws = E::from(w_scale * ow as f32 - (iw as f32)).unwrap();
+
+                        let one = E::from(1.0).unwrap();
+                        let zero = E::from(0.0).unwrap();
+
+                        let ll = buf[b * istr[0] + c * istr[1] + ih * istr[2] + iw * istr[3]]
+                            * (one - hs)
+                            * (one - ws);
+                        let lh = if ws != zero {
+                            buf[b * istr[0] + c * istr[1] + ih * istr[2] + (iw + 1) * istr[3]]
+                                * (one - hs)
+                                * ws
+                        } else {
+                            zero
+                        };
+                        let hl = if hs != zero {
+                            buf[b * istr[0] + c * istr[1] + (ih + 1) * istr[2] + iw * istr[3]]
+                                * hs
+                                * (one - ws)
+                        } else {
+                            zero
+                        };
+                        let hh = if hs != zero && ws != zero {
+                            buf[b * istr[0] + c * istr[1] + (ih + 1) * istr[2] + (iw + 1) * istr[3]]
+                                * hs
+                                * ws
+                        } else {
+                            zero
+                        };
+
+                        out_buf[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]] =
+                            ll + lh + hl + hh;
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn backward<I: Shape, O: Shape>(
+        &self,
+        op: super::Upscale2DOp,
+        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 h_scale = ((istr[1] / istr[2] - 1) as f32) / ((ostr[1] / ostr[2] - 1) as f32);
+        let w_scale = ((istr[2] / istr[3] - 1) as f32) / ((ostr[2] / ostr[3] - 1) as f32);
+
+        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 ih = (h_scale * oh as f32) as usize;
+                        let iw = (w_scale * ow as f32) as usize;
+
+                        let hs = E::from(h_scale * oh as f32 - (ih as f32)).unwrap();
+                        let ws = E::from(w_scale * ow as f32 - (iw as f32)).unwrap();
+
+                        let one = E::from(1.0).unwrap();
+                        let zero = E::from(0.0).unwrap();
+
+                        let g = grad_out[b * ostr[0] + c * ostr[1] + oh * ostr[2] + ow * ostr[3]];
+
+                        grad_inp[b * istr[0] + c * istr[1] + ih * istr[2] + iw * istr[3]] +=
+                            g * (one - hs) * (one - ws);
+                        if ws != zero {
+                            grad_inp
+                                [b * istr[0] + c * istr[1] + ih * istr[2] + (iw + 1) * istr[3]] +=
+                                g * (one - hs) * ws;
+                        }
+                        if hs != zero {
+                            grad_inp
+                                [b * istr[0] + c * istr[1] + (ih + 1) * istr[2] + iw * istr[3]] +=
+                                g * hs * (one - ws);
+                        }
+                        if ws != zero && hs != zero {
+                            grad_inp[b * istr[0]
+                                + c * istr[1]
+                                + (ih + 1) * istr[2]
+                                + (iw + 1) * istr[3]] += g * hs * ws;
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/src/tensor_ops/upscale2d/cuda_kernel.rs b/src/tensor_ops/upscale2d/cuda_kernel.rs
new file mode 100644
index 000000000..562fddd26
--- /dev/null
+++ b/src/tensor_ops/upscale2d/cuda_kernel.rs
@@ -0,0 +1,113 @@
+use crate::{
+    shapes::*,
+    tensor::{Cuda, Tensor},
+};
+
+use std::sync::Arc;
+
+use cudarc::driver::{DeviceRepr, LaunchAsync, LaunchConfig};
+
+use super::{Bilinear, NearestNeighbor};
+
+const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/upscale2d.ptx"));
+
+unsafe impl DeviceRepr for super::Upscale2DOp {}
+
+fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
+    match S::NUM_DIMS {
+        3 => [0, strides[0], strides[1], strides[2]],
+        4 => [strides[0], strides[1], strides[2], strides[3]],
+        _ => panic!("Only implemented for 3d & 4d arrays"),
+    }
+}
+
+fn make_4d_shape<S: Shape>(sizes: S::Concrete) -> [usize; 4] {
+    match S::NUM_DIMS {
+        3 => [1, sizes[0], sizes[1], sizes[2]],
+        4 => [sizes[0], sizes[1], sizes[2], sizes[3]],
+        _ => panic!("Only implemented for 3d & 4d arrays"),
+    }
+}
+
+macro_rules! pool_impl {
+    ($Trait:tt<$TypeName:ty, $UpscaleType:ty>, $Fwd:tt, $Bwd:tt) => {
+        impl super::$Trait<$TypeName, $UpscaleType> for Cuda {
+            fn forward<I: Shape, O: Shape>(
+                &self,
+                op: super::Upscale2DOp,
+                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 inp_sizes = self
+                    .dev
+                    .htod_copy(make_4d_shape::<I>(inp.shape.concrete()).into())?;
+                let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
+                let out_sizes = self
+                    .dev
+                    .htod_copy(make_4d_shape::<O>(out.shape.concrete()).into())?;
+                let fwd_fn = self.dev.get_func($Fwd, $Fwd).unwrap();
+                let cfg = LaunchConfig::for_num_elems(out.shape().num_elements() as u32);
+                let params = (
+                    op,                           // const Pool2dOp op,
+                    &inp_strides,                 // const size_t *inp_strides,
+                    &inp_sizes,                   // const size_t *inp_sizes,
+                    &out_strides,                 // const size_t *out_strides,
+                    &out_sizes,                   // const size_t *out_sizes,
+                    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::Upscale2DOp,
+                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 inp_sizes = self
+                    .dev
+                    .htod_copy(make_4d_shape::<I>(inp.shape.concrete()).into())?;
+                let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides).into())?;
+                let out_sizes = self
+                    .dev
+                    .htod_copy(make_4d_shape::<O>(out.shape.concrete()).into())?;
+                let bwd_fn = self.dev.get_func($Fwd, $Bwd).unwrap();
+                let cfg = LaunchConfig::for_num_elems(inp.shape().num_elements() as u32);
+                let params = (
+                    op,                // const Pool2dOp op,
+                    &inp_strides,      // const size_t *inp_strides,
+                    &inp_sizes,        // const size_t *inp_sizes,
+                    &out_strides,      // const size_t *out_strides,
+                    &out_sizes,        // const size_t *out_sizes,
+                    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(())
+            }
+        }
+    };
+}
+
+pool_impl!(
+    Upscale2DKernel<f32, NearestNeighbor>,
+    "nearest_upscale2d_fwd_f32",
+    "nearest_upscale2d_bwd_f32"
+);
+
+pool_impl!(
+    Upscale2DKernel<f32, Bilinear>,
+    "bilinear_upscale2d_fwd_f32",
+    "bilinear_upscale2d_bwd_f32"
+);
diff --git a/src/tensor_ops/upscale2d/mod.rs b/src/tensor_ops/upscale2d/mod.rs
new file mode 100644
index 000000000..52d9f4a04
--- /dev/null
+++ b/src/tensor_ops/upscale2d/mod.rs
@@ -0,0 +1,273 @@
+mod cpu_kernel;
+
+#[cfg(feature = "cuda")]
+mod cuda_kernel;
+
+use crate::{
+    shapes::*,
+    tensor::{DeviceStorage, HasErr, PutTape, SplitTape, Tape, Tensor, ZerosTensor},
+};
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct Upscale2DOp {
+    pub batch: usize,
+    pub chan: usize,
+    pub h_in: usize,
+    pub h_out: usize,
+    pub w_in: usize,
+    pub w_out: usize,
+}
+
+impl Upscale2DOp {
+    fn new([b, c, h_in, w_in]: [usize; 4], [h_out, w_out]: [usize; 2]) -> Self {
+        assert!(
+            h_out >= h_in,
+            "Output height must be larger than input height"
+        );
+        assert!(
+            w_out >= w_in,
+            "Output width must be larger than input width"
+        );
+        Self {
+            batch: b,
+            chan: c,
+            h_in,
+            h_out,
+            w_in,
+            w_out,
+        }
+    }
+}
+
+pub trait UpscaleMethod: Default {}
+
+#[derive(Clone, Copy, Default)]
+pub struct NearestNeighbor;
+
+impl UpscaleMethod for NearestNeighbor {}
+
+#[derive(Clone, Copy, Default)]
+pub struct Bilinear;
+
+impl UpscaleMethod for Bilinear {}
+
+pub trait Upscale2DKernel<E: Unit, M: UpscaleMethod>: DeviceStorage {
+    fn forward<I: Shape, O: Shape>(
+        &self,
+        op: Upscale2DOp,
+        inp: &Tensor<I, E, Self>,
+        out: &mut Tensor<O, E, Self>,
+    ) -> Result<(), Self::Err>;
+
+    fn backward<I: Shape, O: Shape>(
+        &self,
+        op: Upscale2DOp,
+        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 ConstUpscale2D<M: UpscaleMethod>: HasErr {
+    type Output<const OH: usize, const OW: usize>;
+    fn try_upscale2d<const OH: usize, const OW: usize>(
+        self,
+    ) -> Result<Self::Output<OH, OW>, Self::Err>;
+}
+
+pub trait TryUpscale2D {
+    fn upscale_2d<const OH: usize, const OW: usize, M: UpscaleMethod>(
+        self,
+    ) -> <Self as ConstUpscale2D<M>>::Output<OH, OW>
+    where
+        Self: ConstUpscale2D<M>,
+    {
+        self.try_upscale2d().unwrap()
+    }
+    fn try_upscale_2d<const OH: usize, const OW: usize, M: UpscaleMethod>(
+        self,
+    ) -> Result<<Self as ConstUpscale2D<M>>::Output<OH, OW>, Self::Err>
+    where
+        Self: ConstUpscale2D<M>,
+    {
+        ConstUpscale2D::try_upscale2d(self)
+    }
+}
+impl<T> TryUpscale2D for T {}
+
+impl<
+        C: Dim,
+        const H: usize,
+        const W: usize,
+        E: Dtype,
+        M: UpscaleMethod,
+        D: Upscale2DKernel<E, M> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > ConstUpscale2D<M> for Tensor<(C, Const<H>, Const<W>), E, D, T>
+{
+    type Output<const OH: usize, const OW: usize> = Tensor<(C, Const<OH>, Const<OW>), E, D, T>;
+
+    fn try_upscale2d<const OH: usize, const OW: usize>(
+        self,
+    ) -> Result<Self::Output<OH, OW>, Self::Err> {
+        let &(chan, _, _) = self.shape();
+        let op = Upscale2DOp::new([1, chan.size(), H, W], [OH, OW]);
+        let (inp, mut tape) = self.split_tape();
+        let mut out = inp
+            .device
+            .try_zeros_like(&(chan, Default::default(), Default::default()))?;
+        inp.device.forward(op, &inp, &mut out)?;
+        let phantom_out = out.clone();
+        tape.try_alloc_grad(&inp)?;
+        tape.try_alloc_grad(&out)?;
+        tape.add_backward_op(move |grads| {
+            let (grad_inp, grad_out) = grads.mut_and_ref(&inp, &phantom_out);
+            inp.device
+                .backward(op, &inp, grad_inp, &phantom_out, grad_out)
+        });
+        Ok(out.put_tape(tape))
+    }
+}
+
+impl<
+        B: Dim,
+        C: Dim,
+        const H: usize,
+        const W: usize,
+        E: Dtype,
+        M: UpscaleMethod,
+        D: Upscale2DKernel<E, M> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > ConstUpscale2D<M> for Tensor<(B, C, Const<H>, Const<W>), E, D, T>
+{
+    type Output<const OH: usize, const OW: usize> = Tensor<(B, C, Const<OH>, Const<OW>), E, D, T>;
+
+    fn try_upscale2d<const OH: usize, const OW: usize>(
+        self,
+    ) -> Result<Self::Output<OH, OW>, Self::Err> {
+        let &(batch, chan, _, _) = self.shape();
+        let op = Upscale2DOp::new([batch.size(), chan.size(), H, W], [OH, OW]);
+        let (inp, mut tape) = self.split_tape();
+        let mut out =
+            inp.device
+                .try_zeros_like(&(batch, chan, Default::default(), Default::default()))?;
+        inp.device.forward(op, &inp, &mut out)?;
+        let phantom_out = out.clone();
+        tape.try_alloc_grad(&inp)?;
+        tape.try_alloc_grad(&out)?;
+        tape.add_backward_op(move |grads| {
+            let (grad_inp, grad_out) = grads.mut_and_ref(&inp, &phantom_out);
+            inp.device
+                .backward(op, &inp, grad_inp, &phantom_out, grad_out)
+        });
+        Ok(out.put_tape(tape))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{prelude::*, tests::*};
+
+    use super::{Bilinear, NearestNeighbor, TryUpscale2D};
+
+    #[test]
+    fn nearest_upscale2d_even() {
+        let dev = TestDevice::default();
+
+        let x = dev.tensor([[[1.0, 0.0], [2.0, 3.0]]]);
+        let y = x.leaky_trace().upscale_2d::<4, 4, NearestNeighbor>();
+        assert_close(
+            &y.array(),
+            &[[
+                [1.0, 1.0, 0.0, 0.0],
+                [1.0, 1.0, 0.0, 0.0],
+                [2.0, 2.0, 3.0, 3.0],
+                [2.0, 2.0, 3.0, 3.0],
+            ]],
+        );
+
+        let g = y.exp().mean().backward();
+        assert_close(
+            &g.get(&x).array(),
+            &[[[0.679570457, 0.25], [1.847264025, 5.021384231]]],
+        );
+    }
+
+    #[test]
+    fn nearest_upscale2d_uneven() {
+        let dev = TestDevice::default();
+
+        let x = dev.tensor([[[1.0, 0.0, 2.0], [2.0, 3.0, 4.0]]]);
+        let y = x.leaky_trace().upscale_2d::<2, 7, NearestNeighbor>();
+        assert_close(
+            &y.array(),
+            &[[
+                [1.0, 1.0, 1.0, 0.0, 0.0, 2.0, 2.0],
+                [2.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0],
+            ]],
+        );
+
+        let g = y.exp().mean().backward();
+        assert_close(
+            &g.get(&x).array(),
+            &[[
+                [0.582488963, 0.142857143, 1.055579443],
+                [1.583369164, 2.869362418, 7.799735719],
+            ]],
+        );
+    }
+
+    // Use align_corners when comparing these
+    #[test]
+    fn bilinear_upscale2d_even() {
+        let dev = TestDevice::default();
+
+        let x = dev.tensor([[[1.0, 0.0], [2.0, 3.0]]]);
+        let y = x.leaky_trace().upscale_2d::<4, 4, Bilinear>();
+        assert_close(
+            &y.array(),
+            &[[
+                [1.0000000, 0.6666666, 0.3333333, 0.0000000],
+                [1.3333333, 1.2222222, 1.1111112, 1.0000000],
+                [1.6666667, 1.7777778, 1.8888890, 2.0000000],
+                [2.0000000, 2.3333333, 2.6666665, 3.0000000],
+            ]],
+        );
+
+        let g = y.exp().mean().backward();
+        assert_close(
+            &g.get(&x).array(),
+            &[[[0.8130764, 0.6928807], [1.8153939, 2.7659647]]],
+        );
+    }
+
+    #[test]
+    fn bilinear_upscale2d_uneven() {
+        let dev = TestDevice::default();
+
+        let x = dev.tensor([[[1.0, 0.0, 2.0], [2.0, 3.0, 4.0]]]);
+        let y = x.leaky_trace().upscale_2d::<2, 7, Bilinear>();
+        assert_close(
+            &y.array(),
+            &[[
+                [
+                    1.0000000, 0.6666666, 0.3333333, 0.0000000, 0.6666667, 1.3333335, 2.0000000,
+                ],
+                [
+                    2.0000000, 2.3333333, 2.6666665, 3.0000000, 3.3333335, 3.6666667, 4.0000000,
+                ],
+            ]],
+        );
+
+        let g = y.exp().mean().backward();
+        assert_close(
+            &g.get(&x).array(),
+            &[[
+                [0.3201411, 0.3673356, 0.7548153],
+                [1.3615142, 4.6318388, 6.4302063],
+            ]],
+        );
+    }
+}
diff --git a/src/tensor_ops/upscale2d/upscale2d.cu b/src/tensor_ops/upscale2d/upscale2d.cu
new file mode 100644
index 000000000..5e571bcc3
--- /dev/null
+++ b/src/tensor_ops/upscale2d/upscale2d.cu
@@ -0,0 +1,255 @@
+#include "cuda_utils.cuh"
+
+struct Upscale2dOp {
+    size_t batch;
+    size_t chan;
+    size_t h_in;
+    size_t h_out;
+    size_t w_in;
+    size_t w_out;
+};
+
+template<typename T>
+__device__ void nearest_upscale2d_fwd(
+    const Upscale2dOp op,
+    const size_t *inp_strides,
+    const size_t *inp_sizes,
+    const size_t *out_strides,
+    const size_t *out_sizes,
+    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;
+    }
+
+    float h_scale = ((float)inp_sizes[2])/out_sizes[2];
+    float w_scale = ((float)inp_sizes[3])/out_sizes[3];
+
+    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;
+
+    size_t ih = h_scale * oh;
+    size_t iw = w_scale * ow;
+
+    size_t inp_i = b * inp_strides[0] + c * inp_strides[1] + ih * inp_strides[2] + iw * inp_strides[3];
+    
+    out[i] = inp[inp_i];
+}
+
+template<typename T>
+__device__ void nearest_upscale2d_bwd(
+    const Upscale2dOp op,
+    const size_t *inp_strides,
+    const size_t *inp_sizes,
+    const size_t *out_strides,
+    const size_t *out_sizes,
+    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;
+    }
+
+    float h_scale = ((float)inp_sizes[2])/out_sizes[2];
+    float w_scale = ((float)inp_sizes[3])/out_sizes[3];
+
+    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;
+
+    // Probably isn't efficient, but it works
+    size_t oh_s = 0;
+    size_t ow_s = 0;
+    size_t oh_e = op.h_out-1;
+    size_t ow_e = op.w_out-1;
+    while (oh_s*h_scale < y) {
+        oh_s++;
+    }
+    while (ow_s*w_scale < x) {
+        ow_s++;
+    }
+    while (oh_e*h_scale >= y+1) {
+        oh_e--;
+    }
+    while (ow_e*w_scale >= x+1) {
+        ow_e--;
+    }
+
+    for (int oh = oh_s; oh <= oh_e; oh++) {
+        for (int ow = ow_s; ow <= ow_e; ow++) {
+            size_t out_i = b * out_strides[0] + c * out_strides[1] + oh * out_strides[2] + ow * out_strides[3];
+            grad_inp[i] += grad_out[out_i];
+        }
+    }
+}
+
+template<typename T>
+__device__ void bilinear_upscale2d_fwd(
+    const Upscale2dOp op,
+    const size_t *inp_strides,
+    const size_t *inp_sizes,
+    const size_t *out_strides,
+    const size_t *out_sizes,
+    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;
+    }
+
+    float h_scale = ((float)inp_sizes[2]-1)/(out_sizes[2]-1);
+    float w_scale = ((float)inp_sizes[3]-1)/(out_sizes[3]-1);
+
+    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;
+
+    size_t ih = h_scale * oh;
+    size_t iw = w_scale * ow;
+
+    T hs = h_scale * oh - ih;
+    T ws = w_scale * ow - iw;
+
+    T ll = inp[b * inp_strides[0] + c * inp_strides[1] + ih * inp_strides[2] + iw * inp_strides[3]] * (1-hs) * (1-ws);
+    T lh = ws != 0 ? inp[b * inp_strides[0] + c * inp_strides[1] + ih * inp_strides[2] + (iw+1) * inp_strides[3]] * (1-hs) * ws : 0;
+    T hl = hs != 0 ? inp[b * inp_strides[0] + c * inp_strides[1] + (ih+1) * inp_strides[2] + iw * inp_strides[3]] * hs * (1-ws) : 0;
+    T hh = hs != 0 && ws != 0 ? inp[b * inp_strides[0] + c * inp_strides[1] + (ih+1) * inp_strides[2] + (iw+1) * inp_strides[3]] * hs * ws : 0;
+
+    out[i] = ll + lh + hl + hh;
+}
+
+template<typename T>
+__device__ void bilinear_upscale2d_bwd(
+    const Upscale2dOp op,
+    const size_t *inp_strides,
+    const size_t *inp_sizes,
+    const size_t *out_strides,
+    const size_t *out_sizes,
+    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;
+    }
+
+    float h_scale = ((float)inp_sizes[2]-1)/(out_sizes[2]-1);
+    float w_scale = ((float)inp_sizes[3]-1)/(out_sizes[3]-1);
+
+    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;
+
+    // Probably isn't efficient, but it works
+    size_t oh_s = 0;
+    size_t ow_s = 0;
+    size_t oh_e = op.h_out-1;
+    size_t ow_e = op.w_out-1;
+    while (ceil(oh_s*h_scale) < y) {
+        oh_s++;
+    }
+    while (ceil(ow_s*w_scale) < x) {
+        ow_s++;
+    }
+    while (floor(oh_e*h_scale) > y) {
+        oh_e--;
+    }
+    while (floor(ow_e*w_scale) > x) {
+        ow_e--;
+    }
+
+    for (int oh = oh_s; oh <= oh_e; oh++) {
+        for (int ow = ow_s; ow <= ow_e; ow++) {
+            size_t out_i = b * out_strides[0] + c * out_strides[1] + oh * out_strides[2] + ow * out_strides[3];
+
+            T hs = abs(h_scale * oh - y);
+            T ws = abs(w_scale * ow - x);
+
+            grad_inp[i] += grad_out[out_i] * (1-hs)*(1-ws);
+        }
+    }
+}
+
+#define UPSCALE_OP(TYPENAME, fwd, bwd, fwd_FN, bwd_FN) \
+extern "C" __global__ void fwd( \
+    const Upscale2dOp op, \
+    const size_t *inp_strides, \
+    const size_t *inp_sizes, \
+    const size_t *out_strides, \
+    const size_t *out_sizes, \
+    const TYPENAME *inp, \
+    TYPENAME *out \
+) { \
+    fwd_FN(op, inp_strides, inp_sizes, out_strides, out_sizes, inp, out); \
+} \
+extern "C" __global__ void bwd( \
+    const Upscale2dOp op, \
+    const size_t *inp_strides, \
+    const size_t *inp_sizes, \
+    const size_t *out_strides, \
+    const size_t *out_sizes, \
+    const TYPENAME *inp, \
+    TYPENAME *grad_inp, \
+    const TYPENAME *out, \
+    const TYPENAME *grad_out \
+) { \
+    bwd_FN(op, inp_strides, inp_sizes, out_strides, out_sizes, inp, grad_inp, out, grad_out); \
+}
+
+UPSCALE_OP(
+    float,
+    nearest_upscale2d_fwd_f32, nearest_upscale2d_bwd_f32,
+    nearest_upscale2d_fwd, nearest_upscale2d_bwd
+);
+UPSCALE_OP(
+    double,
+    nearest_upscale2d_fwd_f64, nearest_upscale2d_bwd_f64,
+    nearest_upscale2d_fwd, nearest_upscale2d_bwd
+);
+UPSCALE_OP(
+    float,
+    bilinear_upscale2d_fwd_f32, bilinear_upscale2d_bwd_f32,
+    bilinear_upscale2d_fwd, bilinear_upscale2d_bwd
+);
+UPSCALE_OP(
+    double,
+    bilinear_upscale2d_fwd_f64, bilinear_upscale2d_bwd_f64,
+    bilinear_upscale2d_fwd, bilinear_upscale2d_bwd
+);
\ No newline at end of file


Test patch: