Repository: coreylowman/dfdx (language: rust)

Issue:
`Flatten2D` should work with dynamically sized tensors
Currently, `Flatten2D` only accepts statically size tensors. This breaks something like this:
```rust
#![feature(generic_const_exprs)]
use dfdx::prelude::*;

const N: usize = 8;
const M: usize = 16;

fn main() {
    let dev = Cpu::default();

    let mut model = dev.build_module::<Flatten2D, f32>();
    let value: Tensor<Rank3<1, M, N>, f32, Cpu> = dev.sample_normal();

    let pred = model.forward(value.clone()); // compiles

    let many_values = vec![value; 10].stack();
    let many_pred = model.forward(many_values); // should compile, but doesn't
}
```


Discussion / hints:

Fix patch:
diff --git a/src/nn/flatten.rs b/src/nn/flatten.rs
index 259a9fb23..1223ef041 100644
--- a/src/nn/flatten.rs
+++ b/src/nn/flatten.rs
@@ -29,21 +29,22 @@ where
 }
 
 #[cfg(feature = "nightly")]
-impl<const B: usize, const C: usize, const H: usize, const W: usize, D, E: Dtype, T>
-    Module<Tensor<Rank4<B, C, H, W>, E, D, T>> for Flatten2D
+impl<B: Dim, const C: usize, const H: usize, const W: usize, D, E: Dtype, T>
+    Module<Tensor<(B, Const<C>, Const<H>, Const<W>), E, D, T>> for Flatten2D
 where
     D: Device<E>,
     T: Tape<E, D>,
-    Rank2<B, { C * H * W }>: Sized,
+    (B, Const<{ C * H * W }>): Sized,
 {
-    type Output = Tensor<Rank2<B, { C * H * W }>, E, D, T>;
+    type Output = Tensor<(B, Const<{ C * H * W }>), E, D, T>;
     type Error = D::Err;
 
     fn try_forward(
         &self,
-        input: Tensor<Rank4<B, C, H, W>, E, D, T>,
+        input: Tensor<(B, Const<C>, Const<H>, Const<W>), E, D, T>,
     ) -> Result<Self::Output, D::Err> {
-        input.try_reshape()
+        let batch = input.shape.0;
+        input.try_reshape_like(&(batch, Const)).unwrap()
     }
 }
 
@@ -60,5 +61,8 @@ mod tests {
             Flatten2D.forward_mut(dev.zeros::<Rank3<10, 5, 2>>());
         let _: Tensor<Rank2<5, 24>, TestDtype, _> =
             Flatten2D.forward_mut(dev.zeros::<Rank4<5, 4, 3, 2>>());
+        let x = dev.zeros_like(&(5, Const::<4>, Const::<3>, Const::<2>));
+        let y: Tensor<(usize, Const<24>), TestDtype, _> = Flatten2D.forward_mut(x);
+        assert_eq!(y.shape.0, 5);
     }
 }


Test patch: