Merge pull request #900 from rust-ndarray/fix-zip-collect-bug-with-f

Fix bug in layout and broken debug assertion in collect to f-order array

Author: bluss

src/layout/mod.rs

@@ -78,6 +78,30 @@ mod tests {
     use crate::NdProducer;
 
     type M = Array2<f32>;
+    type M1 = Array1<f32>;
+    type M0 = Array0<f32>;
+
+    macro_rules! assert_layouts {
+        ($mat:expr, $($layout:expr),*) => {{
+            let layout = $mat.view().layout();
+            $(
+            assert!(layout.is($layout), "Assertion failed: array {:?} is not layout {}",
+                $mat,
+                stringify!($layout));
+            )*
+        }}
+    }
+
+    macro_rules! assert_not_layouts {
+        ($mat:expr, $($layout:expr),*) => {{
+            let layout = $mat.view().layout();
+            $(
+            assert!(!layout.is($layout), "Assertion failed: array {:?} show not have layout {}",
+                $mat,
+                stringify!($layout));
+            )*
+        }}
+    }
 
     #[test]
     fn contig_layouts() {
@@ -91,6 +115,34 @@ mod tests {
         assert!(af.is(FORDER) && af.is(FPREFER));
     }
 
+    #[test]
+    fn contig_cf_layouts() {
+        let a = M::zeros((5, 1));
+        let b = M::zeros((1, 5).f());
+        assert_layouts!(a, CORDER, CPREFER, FORDER, FPREFER);
+        assert_layouts!(b, CORDER, CPREFER, FORDER, FPREFER);
+
+        let a = M1::zeros(5);
+        let b = M1::zeros(5.f());
+        assert_layouts!(a, CORDER, CPREFER, FORDER, FPREFER);
+        assert_layouts!(b, CORDER, CPREFER, FORDER, FPREFER);
+
+        let a = M0::zeros(());
+        assert_layouts!(a, CORDER, CPREFER, FORDER, FPREFER);
+
+        let a = M::zeros((5, 5));
+        let b = M::zeros((5, 5).f());
+        let arow = a.slice(s![..1, ..]);
+        let bcol = b.slice(s![.., ..1]);
+        assert_layouts!(arow, CORDER, CPREFER, FORDER, FPREFER);
+        assert_layouts!(bcol, CORDER, CPREFER, FORDER, FPREFER);
+
+        let acol = a.slice(s![.., ..1]);
+        let brow = b.slice(s![..1, ..]);
+        assert_not_layouts!(acol, CORDER, CPREFER, FORDER, FPREFER);
+        assert_not_layouts!(brow, CORDER, CPREFER, FORDER, FPREFER);
+    }
+
     #[test]
     fn stride_layouts() {
         let a = M::zeros((5, 5));

src/zip/mod.rs

@@ -59,7 +59,8 @@ where
     pub(crate) fn layout_impl(&self) -> Layout {
         let n = self.ndim();
         if self.is_standard_layout() {
-            if n <= 1 {
+            // effectively one-dimensional => C and F layout compatible
+            if n <= 1 || self.shape().iter().filter(|&&len| len > 1).count() <= 1 {
                 Layout::one_dimensional()
             } else {
                 Layout::c()
@@ -1185,8 +1186,19 @@ macro_rules! map_impl {
                 {
                     // Get the last producer; and make a Partial that aliases its data pointer
                     let (.., ref output) = &self.parts;
-                    debug_assert!(output.layout().is(CORDER | FORDER));
-                    debug_assert_eq!(output.layout().tendency() >= 0, self.layout_tendency >= 0);
+
+                    // debug assert that the output is contiguous in the memory layout we need
+                    if cfg!(debug_assertions) {
+                        let out_layout = output.layout();
+                        assert!(out_layout.is(CORDER | FORDER));
+                        assert!(
+                            (self.layout_tendency <= 0 && out_layout.tendency() <= 0) ||
+                            (self.layout_tendency >= 0 && out_layout.tendency() >= 0),
+                            "layout tendency violation for self layout {:?}, output layout {:?},\
+                            output shape {:?}",
+                            self.layout, out_layout, output.raw_dim());
+                    }
+
                     let mut partial = Partial::new(output.as_ptr());
 
                     // Apply the mapping function on this zip

tests/par_zip.rs

@@ -107,20 +107,26 @@ fn test_zip_small_collect() {
 
     for m in 0..32 {
         for n in 0..16 {
-            let dim = (m, n);
-            let b = Array::from_shape_fn(dim, |(i, j)| 1. / (i + 2 * j + 1) as f32);
-            let c = Array::from_shape_fn(dim, |(i, j)| f32::ln((1 + i + j) as f32));
-
-            {
-                let a = Zip::from(&b).and(&c).par_map_collect(|x, y| x + y);
-
-                assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
-                assert_eq!(a.strides(), b.strides());
+            for &is_f in &[false, true] {
+                let dim = (m, n).set_f(is_f);
+                let b = Array::from_shape_fn(dim, |(i, j)| 1. / (i + 2 * j + 1) as f32);
+                let c = Array::from_shape_fn(dim, |(i, j)| f32::ln((1 + i + j) as f32));
+
+                {
+                    let a = Zip::from(&b).and(&c).par_map_collect(|x, y| x + y);
+
+                    assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+                    if m > 1 && n > 1 {
+                        assert_eq!(a.strides(), b.strides(),
+                            "Failure for {}x{} c/f: {:?}", m, n, is_f);
+                    }
+                }
             }
         }
     }
 }
 
+
 #[test]
 #[cfg(feature = "approx")]
 fn test_zip_assign_into() {