Add BroadCastShape<Ix0> to the Dimension definition

Author: SparrowLii

src/dimension/broadcast.rs

@@ -41,7 +41,7 @@ where
     Ok(out)
 }
 
-pub trait BroadcastShape<Other: Dimension>: Dimension {
+pub trait BroadcastShape<Other: Dimension> {
     /// The resulting dimension type after broadcasting.
     type BroadcastOutput: Dimension;
 
@@ -51,26 +51,36 @@ pub trait BroadcastShape<Other: Dimension>: Dimension {
     ///
     /// Uses the [NumPy broadcasting rules]
     /// (https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html#general-broadcasting-rules).
-    fn broadcast_shape(&self, other: &Other) -> Result<Self::BroadcastOutput, ShapeError> {
-        broadcast_shape::<Self, Other, Self::BroadcastOutput>(self, other)
-    }
+    fn broadcast_shape(&self, other: &Other) -> Result<Self::BroadcastOutput, ShapeError>;
 }
 
 /// Dimensions of the same type remain unchanged when co_broadcast.
 /// So you can directly use D as the resulting type.
 /// (Instead of <D as BroadcastShape<D>>::BroadcastOutput)
 impl<D: Dimension> BroadcastShape<D> for D {
     type BroadcastOutput = D;
+
+    fn broadcast_shape(&self, other: &D) -> Result<Self::BroadcastOutput, ShapeError> {
+        broadcast_shape::<D, D, Self::BroadcastOutput>(self, other)
+    }
 }
 
 macro_rules! impl_broadcast_distinct_fixed {
     ($smaller:ty, $larger:ty) => {
         impl BroadcastShape<$larger> for $smaller {
             type BroadcastOutput = $larger;
+
+            fn broadcast_shape(&self, other: &$larger) -> Result<Self::BroadcastOutput, ShapeError> {
+                broadcast_shape::<Self, $larger, Self::BroadcastOutput>(self, other)
+            }
         }
 
         impl BroadcastShape<$smaller> for $larger {
             type BroadcastOutput = $larger;
+
+            fn broadcast_shape(&self, other: &$smaller) -> Result<Self::BroadcastOutput, ShapeError> {
+                broadcast_shape::<Self, $smaller, Self::BroadcastOutput>(self, other)
+            }
         }
     };
 }

src/dimension/dimension_trait.rs

@@ -15,7 +15,7 @@ use super::axes_of;
 use super::conversion::Convert;
 use super::{stride_offset, stride_offset_checked};
 use crate::itertools::{enumerate, zip};
-use crate::Axis;
+use crate::{Axis, BroadcastShape};
 use crate::IntoDimension;
 use crate::RemoveAxis;
 use crate::{ArrayView1, ArrayViewMut1};
@@ -46,6 +46,9 @@ pub trait Dimension:
     + MulAssign
     + for<'x> MulAssign<&'x Self>
     + MulAssign<usize>
+    + BroadcastShape<Ix0, BroadcastOutput=Self>
+    + BroadcastShape<Self, BroadcastOutput=Self>
+    + BroadcastShape<IxDyn, BroadcastOutput=IxDyn>
 {
     /// For fixed-size dimension representations (e.g. `Ix2`), this should be
     /// `Some(ndim)`, and for variable-size dimension representations (e.g.

src/numeric/impl_numeric.rs

@@ -13,7 +13,7 @@ use std::ops::{Add, Div, Mul};
 
 use crate::imp_prelude::*;
 use crate::itertools::enumerate;
-use crate::{numeric_util, BroadcastShape};
+use crate::numeric_util;
 
 /// # Numerical Methods for Arrays
 impl<A, S, D> ArrayBase<S, D>
@@ -283,11 +283,10 @@ where
     ///     a.mean_axis(Axis(0)).unwrap().mean_axis(Axis(0)).unwrap() == aview0(&3.5)
     /// );
     /// ```
-    pub fn mean_axis(&self, axis: Axis) -> Option<Array<A, <D::Smaller as BroadcastShape<Ix0>>::BroadcastOutput>>
+    pub fn mean_axis(&self, axis: Axis) -> Option<Array<A, D::Smaller>>
     where
         A: Clone + Zero + FromPrimitive + Add<Output = A> + Div<Output = A>,
         D: RemoveAxis,
-        D::Smaller: BroadcastShape<Ix0>,
     {
         let axis_length = self.len_of(axis);
         if axis_length == 0 {

tests/array.rs

@@ -1568,7 +1568,7 @@ fn test_broadcast_shape() {
         D1: Dimension + BroadcastShape<D2>,
         D2: Dimension,
     {
-        let d = d1.broadcast_shape(&d2);
+        let d = d1.broadcast_shape(d2);
         assert_eq!(d, r);
     }
     test_co(&Dim([2, 3]), &Dim([4, 1, 3]), Ok(Dim([4, 2, 3])));