Repository: coreylowman/dfdx (language: rust)

Issue:
Enable upscale2d on runtime dimensions
This will involve a bit of rewrite of the Upscale2D traits to use `Dim` and `ConstDim`, similar to how Broadcast is implemented


Discussion / hints:
@opfromthestart @nkoppel either of you have any interest in this one? Otherwise I can
It won't be quite as easy as conv2d/convtrans2d because the output dimensions can't be inferred, but I'll take a shot at this. Upscale also needs documentation.
Awesome thanks! Yeah we can stick with opfromthestart's approach of passing the dimensions in. I'm imagining:

```rust
t.upscale2d(); // const case
t.upscale2d_like(height, width); // runtime case
```

Fix patch:
diff --git a/src/nn/upscale.rs b/src/nn/upscale.rs
index 9efa80738..16e9968ef 100644
--- a/src/nn/upscale.rs
+++ b/src/nn/upscale.rs
@@ -1,6 +1,5 @@
-#[cfg(feature = "nightly")]
-use crate::prelude::{Const, Dim, Dtype, HasErr, Tape, Tensor, Upscale2DKernel, ZerosTensor};
-use crate::prelude::{ConstUpscale2D, NearestNeighbor, UpscaleMethod};
+use crate::prelude::{Const, NearestNeighbor, Upscale2DWithMethod, UpscaleMethod};
+use crate::prelude::{Dim, Dtype, HasErr, Tape, Tensor, Upscale2DKernel, ZerosTensor};
 
 #[allow(unused)]
 use super::{BuildModule, Module, NonMutableModule, ZeroSizedModule};
@@ -11,10 +10,10 @@ pub struct Upscale2D<const OH: usize, const OW: usize = OH, M: UpscaleMethod = N
 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>
+impl<const OH: usize, const OW: usize, M: UpscaleMethod, Img: Upscale2DWithMethod<M>> Module<Img>
     for Upscale2D<OH, OW, M>
 {
-    type Output = Img::Output<OH, OW>;
+    type Output = Img::Output<Const<OH>, Const<OW>>;
     type Error = Img::Err;
 
     fn try_forward(&self, x: Img) -> Result<Self::Output, Img::Err> {
@@ -22,13 +21,10 @@ impl<const OH: usize, const OW: usize, M: UpscaleMethod, Img: ConstUpscale2D<M>>
     }
 }
 
-#[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")]
@@ -84,6 +80,52 @@ where
     }
 }
 
+impl<
+        const H: usize,
+        const W: usize,
+        C: Dim,
+        E: Dtype,
+        M: UpscaleMethod,
+        D: Upscale2DKernel<E, M> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > Module<Tensor<(C, usize, usize), E, D, T>> for Upscale2DBy<H, W, M>
+{
+    type Output = Tensor<(C, usize, usize), E, D, T>;
+    type Error = <Self::Output as HasErr>::Err;
+
+    fn try_forward(
+        &self,
+        x: Tensor<(C, usize, usize), E, D, T>,
+    ) -> Result<Self::Output, Self::Error> {
+        let shape = x.shape;
+        x.try_upscale2d_like(shape.1 * H, shape.2 * W)
+    }
+}
+
+impl<
+        const H: usize,
+        const W: usize,
+        B: Dim,
+        C: Dim,
+        E: Dtype,
+        M: UpscaleMethod,
+        D: Upscale2DKernel<E, M> + ZerosTensor<E>,
+        T: 'static + Tape<E, D>,
+    > Module<Tensor<(B, C, usize, usize), E, D, T>> for Upscale2DBy<H, W, M>
+where
+    Tensor<(B, C, usize, usize), E, D, T>: Sized,
+{
+    type Output = Tensor<(B, C, usize, usize), E, D, T>;
+    type Error = <Self::Output as HasErr>::Err;
+
+    fn try_forward(
+        &self,
+        x: Tensor<(B, C, usize, usize), E, D, T>,
+    ) -> Result<Self::Output, Self::Error> {
+        let shape = x.shape;
+        x.try_upscale2d_like(shape.2 * H, shape.3 * W)
+    }
+}
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
index 5d1a8bd73..14d7fa910 100644
--- a/src/tensor_ops/mod.rs
+++ b/src/tensor_ops/mod.rs
@@ -268,9 +268,8 @@ mod convtrans2d;
 pub use convtrans2d::{ConvTransAlgebra, TryConvTrans2D, TryConvTrans2DTo};
 
 mod upscale2d;
-#[cfg(feature = "nightly")]
 pub(crate) use upscale2d::Upscale2DKernel;
-pub use upscale2d::{Bilinear, ConstUpscale2D, NearestNeighbor, TryUpscale2D, UpscaleMethod};
+pub use upscale2d::{Bilinear, NearestNeighbor, TryUpscale2D, Upscale2DWithMethod, UpscaleMethod};
 
 #[cfg(feature = "nightly")]
 mod pool2d;
diff --git a/src/tensor_ops/upscale2d/mod.rs b/src/tensor_ops/upscale2d/mod.rs
index 52d9f4a04..433a83944 100644
--- a/src/tensor_ops/upscale2d/mod.rs
+++ b/src/tensor_ops/upscale2d/mod.rs
@@ -70,54 +70,88 @@ pub trait Upscale2DKernel<E: Unit, M: UpscaleMethod>: DeviceStorage {
     ) -> Result<(), Self::Err>;
 }
 
-pub trait ConstUpscale2D<M: UpscaleMethod>: HasErr {
-    type Output<const OH: usize, const OW: usize>;
+pub trait Upscale2DWithMethod<M: UpscaleMethod>: HasErr {
+    type Output<OH: Dim, OW: Dim>;
     fn try_upscale2d<const OH: usize, const OW: usize>(
         self,
+    ) -> Result<Self::Output<Const<OH>, Const<OW>>, Self::Err> {
+        self.try_upscale2d_like(Const::<OH>, Const::<OW>)
+    }
+
+    fn try_upscale2d_like<OH: Dim, OW: Dim>(
+        self,
+        height: OH,
+        width: OW,
     ) -> 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>
+    ) -> <Self as Upscale2DWithMethod<M>>::Output<Const<OH>, Const<OW>>
     where
-        Self: ConstUpscale2D<M>,
+        Self: Upscale2DWithMethod<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>
+    ) -> Result<<Self as Upscale2DWithMethod<M>>::Output<Const<OH>, Const<OW>>, Self::Err>
+    where
+        Self: Upscale2DWithMethod<M>,
+    {
+        Upscale2DWithMethod::try_upscale2d(self)
+    }
+    fn upscale_2d_like<OH: Dim, OW: Dim, M: UpscaleMethod>(
+        self,
+        height: OH,
+        width: OW,
+    ) -> <Self as Upscale2DWithMethod<M>>::Output<OH, OW>
+    where
+        Self: Upscale2DWithMethod<M>,
+    {
+        self.try_upscale2d_like(height, width).unwrap()
+    }
+    fn try_upscale_2d_like<OH: Dim, OW: Dim, M: UpscaleMethod>(
+        self,
+        height: OH,
+        width: OW,
+    ) -> Result<<Self as Upscale2DWithMethod<M>>::Output<OH, OW>, Self::Err>
     where
-        Self: ConstUpscale2D<M>,
+        Self: Upscale2DWithMethod<M>,
     {
-        ConstUpscale2D::try_upscale2d(self)
+        Upscale2DWithMethod::try_upscale2d_like(self, height, width)
     }
 }
 impl<T> TryUpscale2D for T {}
 
 impl<
         C: Dim,
-        const H: usize,
-        const W: usize,
+        H: Dim,
+        W: Dim,
         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>
+    > Upscale2DWithMethod<M> for Tensor<(C, H, W), E, D, T>
 {
-    type Output<const OH: usize, const OW: usize> = Tensor<(C, Const<OH>, Const<OW>), E, D, T>;
+    type Output<OH: Dim, OW: Dim> = Tensor<(C, OH, OW), E, D, T>;
 
-    fn try_upscale2d<const OH: usize, const OW: usize>(
+    fn try_upscale2d_like<OH: Dim, OW: Dim>(
         self,
+        out_height: OH,
+        out_width: OW,
     ) -> Result<Self::Output<OH, OW>, Self::Err> {
+        let in_height = self.shape.1;
+        let in_width = self.shape.2;
+
         let &(chan, _, _) = self.shape();
-        let op = Upscale2DOp::new([1, chan.size(), H, W], [OH, OW]);
+        let op = Upscale2DOp::new(
+            [1, chan.size(), in_height.size(), in_width.size()],
+            [out_height.size(), out_width.size()],
+        );
         let (inp, mut tape) = self.split_tape();
-        let mut out = inp
-            .device
-            .try_zeros_like(&(chan, Default::default(), Default::default()))?;
+        let mut out = inp.device.try_zeros_like(&(chan, out_height, out_width))?;
         inp.device.forward(op, &inp, &mut out)?;
         let phantom_out = out.clone();
         tape.try_alloc_grad(&inp)?;
@@ -134,25 +168,33 @@ impl<
 impl<
         B: Dim,
         C: Dim,
-        const H: usize,
-        const W: usize,
+        H: Dim,
+        W: Dim,
         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>
+    > Upscale2DWithMethod<M> for Tensor<(B, C, H, W), E, D, T>
 {
-    type Output<const OH: usize, const OW: usize> = Tensor<(B, C, Const<OH>, Const<OW>), E, D, T>;
+    type Output<OH: Dim, OW: Dim> = Tensor<(B, C, OH, OW), E, D, T>;
 
-    fn try_upscale2d<const OH: usize, const OW: usize>(
+    fn try_upscale2d_like<OH: Dim, OW: Dim>(
         self,
+        out_height: OH,
+        out_width: OW,
     ) -> Result<Self::Output<OH, OW>, Self::Err> {
+        let in_height = self.shape.2;
+        let in_width = self.shape.3;
+
         let &(batch, chan, _, _) = self.shape();
-        let op = Upscale2DOp::new([batch.size(), chan.size(), H, W], [OH, OW]);
+        let op = Upscale2DOp::new(
+            [batch.size(), chan.size(), in_height.size(), in_width.size()],
+            [out_height.size(), out_width.size()],
+        );
         let (inp, mut tape) = self.split_tape();
-        let mut out =
-            inp.device
-                .try_zeros_like(&(batch, chan, Default::default(), Default::default()))?;
+        let mut out = inp
+            .device
+            .try_zeros_like(&(batch, chan, out_height, out_width))?;
         inp.device.forward(op, &inp, &mut out)?;
         let phantom_out = out.clone();
         tape.try_alloc_grad(&inp)?;


Test patch: