fix bug in reshape on cuda; fix 02-ops on cuda (#622)

coreylowman · Mar 27, 2023 · cf6e11b · cf6e11b
1 parent 67f4884
commit cf6e11b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/examples/02-ops.rs b/examples/02-ops.rs
@@ -57,7 +57,7 @@ fn main() {
     // these operations are equal across devices
     #[cfg(feature = "cuda")]
     {
-        use dfdx::nn::ToDevice;
+        use dfdx::{nn::ToDevice, tensor::Cpu};
 
         let cpu = Cpu::default();
 

diff --git a/src/tensor_ops/reshape_to/cuda_kernel.rs b/src/tensor_ops/reshape_to/cuda_kernel.rs
@@ -47,7 +47,7 @@ where
             self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?;
         }
 
-        let numel = inp.data.len();
+        let numel = inp.shape.num_elements();
         let mut storage = unsafe { self.dev.alloc::<E>(numel) }?;
 
         let inp_dims = self.dev.htod_copy(inp.shape.concrete().into())?;

diff --git a/src/tensor_ops/reshape_to/mod.rs b/src/tensor_ops/reshape_to/mod.rs
@@ -187,4 +187,18 @@ mod tests {
             ],
         )
     }
+
+    #[test]
+    fn test_reshape_broadcasted() {
+        let dev: TestDevice = Default::default();
+        let a: Tensor<Rank2<2, 3>, TestDtype, _> = dev.tensor([1., 2., 3.]).broadcast();
+        let b: Tensor<Rank2<3, 2>, TestDtype, _> = a.clone().reshape();
+
+        #[cfg(feature = "cuda")]
+        use cudarc::driver::DeviceSlice;
+
+        assert_eq!(b.data.len(), 6);
+        assert_eq!(a.as_vec(), b.as_vec());
+        assert_eq!(b.array(), [[1., 2.], [3., 1.], [2., 3.]]);
+    }
 }