Fix bug in #451

src/tensor_ops/max_to/cuda_kernel.rs

@@ -1,7 +1,7 @@
 use crate::{
     shapes::*,
     tensor::cuda::{Cuda, CudaArray},
-    tensor_ops::internal_reshapes::permute_for_reductions,
+    tensor_ops::reduction_utils::*,
 };
 
 use cudarc::driver::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig};
@@ -61,7 +61,9 @@ where
         let strides: CudaSlice<usize> = self.dev.take_async(strides)?;
 
         let physical_numel = inp.data.len();
-        let chunk_len = physical_numel / dst.num_elements();
+        let (dst_physical_numel, dst_strides) =
+            reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
+        let chunk_len = physical_numel / dst_physical_numel;
 
         let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
         let params = (
@@ -77,7 +79,7 @@ where
         Ok(CudaArray {
             data: Arc::new(storage),
             shape: dst,
-            strides: dst.strides(),
+            strides: dst_strides,
         })
     }
 
@@ -100,8 +102,11 @@ where
         let out_strides: CudaSlice<usize> = self.dev.take_async(out_strides.into())?;
 
         let physical_numel = grad_inp.data.len();
-        let virtual_numel = grad_inp.shape.num_elements();
-        let elems_per_thread = E::from_usize(virtual_numel / physical_numel).unwrap();
+        let elems_per_thread = E::from_usize(reduction_elems_per_thread::<Ax, Src>(
+            grad_inp.shape.concrete(),
+            grad_inp.strides,
+        ))
+        .unwrap();
 
         let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
         let params = (

src/tensor_ops/min_to/cuda_kernel.rs

@@ -1,7 +1,7 @@
 use crate::{
     shapes::*,
     tensor::cuda::{Cuda, CudaArray},
-    tensor_ops::internal_reshapes::permute_for_reductions,
+    tensor_ops::reduction_utils::*,
 };
 
 use cudarc::driver::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig};
@@ -61,7 +61,9 @@ where
         let strides: CudaSlice<usize> = self.dev.take_async(strides)?;
 
         let physical_numel = inp.data.len();
-        let chunk_len = physical_numel / dst.num_elements();
+        let (dst_physical_numel, dst_strides) =
+            reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
+        let chunk_len = physical_numel / dst_physical_numel;
 
         let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
         let params = (
@@ -77,7 +79,7 @@ where
         Ok(CudaArray {
             data: Arc::new(storage),
             shape: dst,
-            strides: dst.strides(),
+            strides: dst_strides,
         })
     }
 
@@ -100,8 +102,11 @@ where
         let out_strides: CudaSlice<usize> = self.dev.take_async(out_strides.into())?;
 
         let physical_numel = grad_inp.data.len();
-        let virtual_numel = grad_inp.shape.num_elements();
-        let elems_per_thread = E::from_usize(virtual_numel / physical_numel).unwrap();
+        let elems_per_thread = E::from_usize(reduction_elems_per_thread::<Ax, Src>(
+            grad_inp.shape.concrete(),
+            grad_inp.strides,
+        ))
+        .unwrap();
 
         let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
         let params = (

src/tensor_ops/sum_to/cuda_kernel.rs

@@ -1,7 +1,7 @@
 use crate::{
     shapes::*,
     tensor::cuda::{Cuda, CudaArray},
-    tensor_ops::internal_reshapes::permute_for_reductions,
+    tensor_ops::reduction_utils::*,
 };
 
 use cudarc::driver::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits};
@@ -50,11 +50,16 @@ where
 
         let mut storage = self.dev.alloc_zeros_async::<E>(dst.num_elements())?;
 
-        let physical_numel = inp.data.len();
-        let virtual_numel = inp.shape.num_elements();
-        let elems_per_thread = E::from_usize(virtual_numel / physical_numel).unwrap();
+        let elems_per_thread = E::from_usize(reduction_elems_per_thread::<Ax, Src>(
+            inp.shape.concrete(),
+            inp.strides,
+        ))
+        .unwrap();
 
-        let chunk_len = physical_numel / dst.num_elements();
+        let physical_numel = inp.data.len();
+        let (dst_physical_numel, dst_strides) =
+            reduction_output_strides::<Ax, Src, Dst>(inp.strides, dst);
+        let chunk_len = physical_numel / dst_physical_numel;
 
         let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
         let params = (
@@ -71,7 +76,7 @@ where
         Ok(CudaArray {
             data: Arc::new(storage),
             shape: dst,
-            strides: dst.strides(),
+            strides: dst_strides,
         })
     }
 
@@ -92,8 +97,11 @@ where
         let out_strides: CudaSlice<usize> = self.dev.take_async(out_strides.into())?;
 
         let physical_numel = grad_inp.data.len();
-        let virtual_numel = grad_inp.shape.num_elements();
-        let elems_per_thread = E::from_usize(virtual_numel / physical_numel).unwrap();
+        let elems_per_thread = E::from_usize(reduction_elems_per_thread::<Ax, Src>(
+            grad_inp.shape.concrete(),
+            grad_inp.strides,
+        ))
+        .unwrap();
 
         let cfg = LaunchConfig::for_num_elems(physical_numel as u32);
         let params = (

src/tensor_ops/sum_to/mod.rs

@@ -150,4 +150,13 @@ mod tests {
         let g = r.sum().backward();
         assert_close(&g.get(&t).array(), &t.array());
     }
+
+    #[test]
+    fn test_sum_reduce_to_more_than_physical_elements() {
+        let dev: TestDevice = Default::default();
+        let a: Tensor<_, TestDtype, _> = dev.tensor([1.0, 2.0, 3.0]);
+        let b = a.broadcast::<Rank3<4, 3, 2>, _>();
+        let c = b.sum::<Rank2<4, 3>, _>();
+        assert_eq!(c.array(), [[2.0, 4.0, 6.0]; 4]);
+    }
 }

src/tensor_ops/utilities/mod.rs

@@ -3,8 +3,8 @@ pub(crate) mod cpu_kernels;
 #[cfg(feature = "cuda")]
 pub(crate) mod cuda_kernels;
 mod device;
-pub(crate) mod internal_reshapes;
 pub(crate) mod ops;
+pub(crate) mod reduction_utils;
 
 pub use backward::Backward;
 pub use device::Device;

src/tensor_ops/utilities/reduction_utils.rs

@@ -0,0 +1,76 @@
+#[cfg(feature = "cuda")]
+use crate::prelude::{Axes, Shape};
+#[cfg(feature = "cuda")]
+use std::vec::Vec;
+
+/// Moves all axes in Ax to the end of dims and strides and removes broadcasted dimensions
+/// so that a cuda kernel called for each physical element of the input tensor will place elements
+/// to be reduced with each other next to each other in memory.
+#[cfg(feature = "cuda")]
+pub(crate) fn permute_for_reductions<I, Ax: Axes>(dims: I, strides: I) -> (Vec<usize>, Vec<usize>)
+where
+    I: IntoIterator<Item = usize>,
+{
+    let mut tmp = dims
+        .into_iter()
+        .zip(strides.into_iter())
+        .map(|x| (false, x))
+        .collect::<Vec<_>>();
+
+    for i in Ax::as_array().into_iter() {
+        tmp[i as usize].0 = true;
+    }
+
+    // requires stable sorting to keep non-summed axes in the correct order
+    tmp.sort_by_key(|x| x.0);
+
+    tmp.into_iter()
+        .map(|(_, x)| x)
+        .filter(|(_, stride)| *stride != 0)
+        .unzip()
+}
+
+/// Returns the physical number of elements and strides of dst so that broadcasted dimensions in
+/// src are also broadcasted in dst
+#[cfg(feature = "cuda")]
+pub(crate) fn reduction_output_strides<Ax: Axes, Src: Shape, Dst: Shape>(
+    src_strides: Src::Concrete,
+    dst: Dst,
+) -> (usize, Dst::Concrete) {
+    let dst_dims = dst.concrete();
+    let mut dst_strides = dst.strides();
+    let mut numel = 1;
+    let mut j = Dst::NUM_DIMS;
+
+    for i in (0..Src::NUM_DIMS).rev() {
+        if !Ax::as_array().into_iter().any(|x| x as usize == i) {
+            j -= 1;
+            if src_strides[i] == 0 {
+                dst_strides[j] = 0;
+            } else {
+                dst_strides[j] = numel;
+                numel *= dst_dims[j];
+            }
+        }
+    }
+
+    (numel, dst_strides)
+}
+
+/// Gives the product of all dimensions that are being reduced and are broadcasted.
+#[cfg(feature = "cuda")]
+pub(crate) fn reduction_elems_per_thread<Ax: Axes, S: Shape>(
+    dims: S::Concrete,
+    strides: S::Concrete,
+) -> usize {
+    Ax::as_array()
+        .into_iter()
+        .map(|ax| {
+            if strides[ax as usize] == 0 {
+                dims[ax as usize]
+            } else {
+                1
+            }
+        })
+        .product()
+}