zip_mut_with 'f' order

benches/iter.rs

@@ -370,3 +370,23 @@ fn iter_axis_chunks_5_iter_sum(bench: &mut Bencher) {
             .sum::<f32>()
     });
 }
+
+pub fn zip_mut_with(data: &Array3<f32>, out: &mut Array3<f32>) {
+    out.zip_mut_with(&data, |o, &i| {
+        *o = i;
+    });
+}
+
+#[bench]
+fn zip_mut_with_cc(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros((ISZ, ISZ, ISZ));
+    let mut out = Array3::zeros(data.dim());
+    b.iter(|| black_box(zip_mut_with(&data, &mut out)));
+}
+
+#[bench]
+fn zip_mut_with_ff(b: &mut Bencher) {
+    let data: Array3<f32> = Array3::zeros((ISZ, ISZ, ISZ).f());
+    let mut out = Array3::zeros(data.dim().f());
+    b.iter(|| black_box(zip_mut_with(&data, &mut out)));
+}

src/impl_methods.rs

@@ -6,7 +6,6 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use std::cmp;
 use std::ptr as std_ptr;
 use std::slice;
 
@@ -1917,18 +1916,20 @@ where
         F: FnMut(&mut A, &B),
     {
         debug_assert_eq!(self.shape(), rhs.shape());
-        if let Some(self_s) = self.as_slice_mut() {
-            if let Some(rhs_s) = rhs.as_slice() {
-                let len = cmp::min(self_s.len(), rhs_s.len());
-                let s = &mut self_s[..len];
-                let r = &rhs_s[..len];
-                for i in 0..len {
-                    f(&mut s[i], &r[i]);
+
+        // Same shape and order should have same strides
+        if self.strides() == rhs.strides() {
+            if let Some(self_s) = self.as_slice_memory_order_mut() {
+                if let Some(rhs_s) = rhs.as_slice_memory_order() {
+                    for (s, r) in self_s.iter_mut().zip(rhs_s) {
+                        f(s, &r);
+                    }
+                    return;
                 }
-                return;
             }
         }
-        // otherwise, fall back to the outer iter
+
+        // Otherwise, fall back to the outer iter
         self.zip_mut_with_by_rows(rhs, f);
     }