FEAT: Add method squeeze_into

This method can squeeze into a particular dimensionality.

Squeezing means removing axes of length 1. When squeezing to a
particular dimensionality, we may have to still pad out the shape with
extra 1-shape axes to fill the dimensionality.

Author: bluss

src/dimension/mod.rs

@@ -762,6 +762,53 @@ where
     *strides = new_strides;
 }
 
+/// Remove axes with length one, except never removing the last axis.
+pub(crate) fn squeeze_into<D, E>(dim: &D, strides: &D) -> Result<(E, E), ShapeError>
+where
+    D: Dimension,
+    E: Dimension,
+{
+    debug_assert_eq!(dim.ndim(), strides.ndim());
+
+    // Count axes with dim == 1; we keep axes with d == 0 or d > 1
+    let mut ndim_new = 0;
+    for &d in dim.slice() {
+        if d != 1 { ndim_new += 1; }
+    }
+    let mut fill_ones = 0;
+    if let Some(e_ndim) = E::NDIM {
+        if e_ndim < ndim_new {
+            return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape));
+        }
+        fill_ones = e_ndim - ndim_new;
+        ndim_new = e_ndim;
+    } else {
+        // dynamic-dimensional
+        // use minimum one dimension unless input has less than one dim
+        if dim.ndim() > 0 && ndim_new == 0 {
+            ndim_new = 1;
+            fill_ones = 1;
+        }
+    }
+
+    let mut new_dim = E::zeros(ndim_new);
+    let mut new_strides = E::zeros(ndim_new);
+    let mut i = 0;
+    while i < fill_ones {
+        new_dim[i] = 1;
+        new_strides[i] = 1;
+        i += 1;
+    }
+    for (&d, &s) in izip!(dim.slice(), strides.slice()) {
+        if d != 1 {
+            new_dim[i] = d;
+            new_strides[i] = s;
+            i += 1;
+        }
+    }
+    Ok((new_dim, new_strides))
+}
+
 
 /// Sort axes to standard/row major order, i.e Axis(0) has biggest stride and Axis(n - 1) least
 /// stride
@@ -1148,6 +1195,91 @@ mod test {
         assert_eq!(s, sans);
     }
 
+    #[test]
+    #[cfg(feature = "std")]
+    fn test_squeeze_into() {
+        use super::squeeze_into;
+
+        let dyndim = Dim::<&[usize]>;
+
+        // squeeze to ixdyn
+        let d = dyndim(&[1, 2, 1, 1, 3, 1]);
+        let s = dyndim(&[!0, !0, !0, 9, 10, !0]);
+        let dans = dyndim(&[2, 3]);
+        let sans = dyndim(&[!0, 10]);
+        let (d2, s2) = squeeze_into::<_, IxDyn>(&d, &s).unwrap();
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+
+        // squeeze to ixdyn does not go below 1D
+        let d = dyndim(&[1, 1]);
+        let s = dyndim(&[3, 4]);
+        let dans = dyndim(&[1]);
+        let sans = dyndim(&[1]);
+        let (d2, s2) = squeeze_into::<_, IxDyn>(&d, &s).unwrap();
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+
+        let d = Dim([1, 1]);
+        let s = Dim([3, 4]);
+        let dans = Dim([1]);
+        let sans = Dim([1]);
+        let (d2, s2) = squeeze_into::<_, Ix1>(&d, &s).unwrap();
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+
+        // squeeze to zero-dim
+        let (d2, s2) = squeeze_into::<_, Ix0>(&d, &s).unwrap();
+        assert_eq!(d2, Ix0());
+        assert_eq!(s2, Ix0());
+
+        let d = Dim([0, 1, 3, 4]);
+        let s = Dim([2, 3, 4, 5]);
+        let dans = Dim([0, 3, 4]);
+        let sans = Dim([2, 4, 5]);
+        let (d2, s2) = squeeze_into::<_, Ix3>(&d, &s).unwrap();
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+
+        // Pad with ones
+        let d = Dim([0, 1, 3, 1]);
+        let s = Dim([2, 3, 4, 5]);
+        let dans = Dim([1, 0, 3]);
+        let sans = Dim([1, 2, 4]);
+        let (d2, s2) = squeeze_into::<_, Ix3>(&d, &s).unwrap();
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+
+        // Try something that doesn't fit
+        let d = Dim([0, 1, 3, 1]);
+        let s = Dim([2, 3, 4, 5]);
+        let res = squeeze_into::<_, Ix1>(&d, &s);
+        assert!(res.is_err());
+        let res = squeeze_into::<_, Ix0>(&d, &s);
+        assert!(res.is_err());
+
+        // Squeeze 0d to 0d
+        let d = Dim([]);
+        let s = Dim([]);
+        let res = squeeze_into::<_, Ix0>(&d, &s);
+        assert!(res.is_ok());
+        // grow 0d to 2d
+        let dans = Dim([1, 1]);
+        let sans = Dim([1, 1]);
+        let (d2, s2) = squeeze_into::<_, Ix2>(&d, &s).unwrap();
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+
+        // Squeeze 0d to 0d dynamic
+        let d = dyndim(&[]);
+        let s = dyndim(&[]);
+        let (d2, s2) = squeeze_into::<_, IxDyn>(&d, &s).unwrap();
+        let dans = d;
+        let sans = s;
+        assert_eq!(d2, dans);
+        assert_eq!(s2, sans);
+    }
+
     #[test]
     fn test_merge_axes_from_the_back() {
         let dyndim = Dim::<&[usize]>;