Dimension::strides_equivalent

Author: nilgoyette

src/dimension/dimension_trait.rs

@@ -229,6 +229,25 @@ pub trait Dimension:
         !end_iteration
     }
 
+    /// Returns `true` iff `strides1` and `strides2` are equivalent for the
+    /// shape `self`.
+    ///
+    /// The strides are equivalent if, for each axis with length > 1, the
+    /// strides are equal.
+    ///
+    /// Note: Returns `false` if any of the ndims don't match.
+    #[doc(hidden)]
+    fn strides_equivalent<D>(&self, strides1: &Self, strides2: &D) -> bool
+    where
+        D: Dimension,
+    {
+        let shape_ndim = self.ndim();
+        shape_ndim == strides1.ndim()
+            && shape_ndim == strides2.ndim()
+            && izip!(self.slice(), strides1.slice(), strides2.slice())
+                .all(|(&d, &s1, &s2)| d <= 1 || s1 as isize == s2 as isize)
+    }
+
     #[doc(hidden)]
     /// Return stride offset for index.
     fn stride_offset(index: &Self, strides: &Self) -> isize {

src/impl_methods.rs

@@ -1917,8 +1917,7 @@ where
     {
         debug_assert_eq!(self.shape(), rhs.shape());
 
-        // Same shape and order should have same strides
-        if self.strides() == rhs.strides() {
+        if self.dim.strides_equivalent(&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) {