coreylowman · coreylowman · Feb 15, 2023 · Feb 15, 2023
diff --git a/src/tensor_ops/reshape_to/cpu_kernel.rs b/src/tensor_ops/reshape_to/cpu_kernel.rs
@@ -1,15 +1,12 @@
-use crate::shapes::{Dtype, HasSameNumelAs, Shape};
+use crate::shapes::{Dtype, Shape};
 use crate::tensor::cpu::{Cpu, LendingIterator, StridedArray};
 
 impl<E: Dtype> super::ReshapeKernel<E> for Cpu {
     fn forward<Src: Shape, Dst: Shape>(
         &self,
         dst: Dst,
         inp: &Self::Storage<Src, E>,
-    ) -> Result<Self::Storage<Dst, E>, Self::Err>
-    where
-        Src: HasSameNumelAs<Dst>,
-    {
+    ) -> Result<Self::Storage<Dst, E>, Self::Err> {
         let mut out = StridedArray::new(dst)?;
         let mut inp_iter = inp.iter();
         let mut out_iter = out.iter_mut();
@@ -23,10 +20,7 @@ impl<E: Dtype> super::ReshapeKernel<E> for Cpu {
         &self,
         grad_inp: &mut Self::Storage<Src, E>,
         grad_out: &Self::Storage<Dst, E>,
-    ) -> Result<(), Self::Err>
-    where
-        Src: HasSameNumelAs<Dst>,
-    {
+    ) -> Result<(), Self::Err> {
         let mut inp_iter = grad_inp.iter_mut();
         let mut out_iter = grad_out.iter();
         while let Some((i, o)) = inp_iter.next().zip(out_iter.next()) {

diff --git a/src/tensor_ops/reshape_to/cuda_kernel.rs b/src/tensor_ops/reshape_to/cuda_kernel.rs
@@ -30,10 +30,7 @@ where
         &self,
         dst: Dst,
         inp: &Self::Storage<Src, E>,
-    ) -> Result<Self::Storage<Dst, E>, Self::Err>
-    where
-        Src: HasSameNumelAs<Dst>,
-    {
+    ) -> Result<Self::Storage<Dst, E>, Self::Err> {
         if !self.dev.has_func(Self::MOD, Self::FNS[0]) {
             self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?;
         }
@@ -72,10 +69,7 @@ where
         &self,
         grad_inp: &mut Self::Storage<Src, E>,
         grad_out: &Self::Storage<Dst, E>,
-    ) -> Result<(), Self::Err>
-    where
-        Src: HasSameNumelAs<Dst>,
-    {
+    ) -> Result<(), Self::Err> {
         let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap();
         let numel = grad_inp.data.len();
 

diff --git a/src/tensor_ops/reshape_to/mod.rs b/src/tensor_ops/reshape_to/mod.rs
@@ -10,16 +10,12 @@ pub trait ReshapeKernel<E: Dtype>: DeviceStorage {
         &self,
         dst: Dst,
         inp: &Self::Storage<Src, E>,
-    ) -> Result<Self::Storage<Dst, E>, Self::Err>
-    where
-        Src: HasSameNumelAs<Dst>;
+    ) -> Result<Self::Storage<Dst, E>, Self::Err>;
     fn backward<Src: Shape, Dst: Shape>(
         &self,
         grad_inp: &mut Self::Storage<Src, E>,
         grad_out: &Self::Storage<Dst, E>,
-    ) -> Result<(), Self::Err>
-    where
-        Src: HasSameNumelAs<Dst>;
+    ) -> Result<(), Self::Err>;
 }
 
 /// **Requires Nightly** Change the shape of a tensor moving data around.
@@ -32,17 +28,21 @@ pub trait ReshapeTo: HasErr + HasShape {
     }
     fn try_reshape<Dst: ConstShape>(self) -> Result<Self::WithShape<Dst>, Self::Err>
     where
-        Self::Shape: HasSameNumelAs<Dst>;
+        Self::Shape: HasSameNumelAs<Dst>,
+    {
+        self.try_reshape_like::<Dst>(&Default::default())
+    }
+    fn reshape_like<Dst: Shape>(self, dst: &Dst) -> Self::WithShape<Dst> {
+        self.try_reshape_like(dst).unwrap()
+    }
+    fn try_reshape_like<Dst: Shape>(self, dst: &Dst) -> Result<Self::WithShape<Dst>, Self::Err>;
 }
 
 impl<S: Shape, E: Dtype, D: ReshapeKernel<E>, T: Tape<D>> ReshapeTo for Tensor<S, E, D, T> {
-    fn try_reshape<Dst: ConstShape>(self) -> Result<Self::WithShape<Dst>, Self::Err>
-    where
-        Self::Shape: HasSameNumelAs<Dst>,
-    {
-        let dst: Dst = Default::default();
+    fn try_reshape_like<Dst: Shape>(self, dst: &Dst) -> Result<Self::WithShape<Dst>, Self::Err> {
+        assert_eq!(self.shape().num_elements(), dst.shape().num_elements());
         let (inp, mut tape) = self.split_tape();
-        let out = inp.device.upgrade(inp.device.forward(dst, &inp.storage)?);
+        let out = inp.device.upgrade(inp.device.forward(*dst, &inp.storage)?);
         let phantom_out = out.clone();
         tape.try_alloc_grad(&inp)?;
         tape.try_alloc_grad(&out)?;
@@ -63,6 +63,14 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    #[should_panic]
+    fn test_invalid_reshape() {
+        let dev: TestDevice = Default::default();
+        let t: Tensor<(usize,), TestDtype, _> = dev.zeros_like(&(5,));
+        let _ = t.reshape_like(&(7,));
+    }
+
     #[test]
     fn test_valid_reshapes() {
         let dev: TestDevice = Default::default();