Repository: rust-ndarray/ndarray
Pull request #898 (2021-01-20 03:39:45)

Implement co-broadcasting in operator overloading
This PR did the following to implement co_broadcasting in operator overloading
1. Implement DimMax trait, which uses the same broadcast mechanism as Numpy, for all Dimension so as to get the returned array type in operator overloading. 
2. Use map_collect() and map_collect_owned() in operator overloading to avoid redundant array traversal.


bluss: Co-broadcasting is cool and interesting! I can see that you are attacking our fixme's in code one by one and that's gratifying, but it doesn't always tell the whole story of what might be the most important issues to solve. (*)

I think the type system is still limited in what can be expressed, and we seem to have a typenum-like problem here (not really the fault of typenum, just the limits of the type system). Expressing one generic arithmetic operation is ok.. but what about two? will it be usable for the user?  Imagine this:

```rust
type Mat<D> = Array<f64, D>;
fn foo<D>(x: Mat<D>, y: Mat<D>, z: Mat<D>) -> Mat<D> {
   x + &y * &z
}
```

How would we write the actual return type of the function? Maybe in some cases this makes it harder for the user of ndarray, or can we find a great solution?

It could be that co-broadcasting should be a feature that ndarray implements, without using it by default in the arithmetic ops.

(*) Specifically see the issue about #699 - there are lots of things we want to spend the - implementation, complexity, etc - budget on in arithmetic ops. 
SparrowLii: > How would we write the actual return type of the function? Maybe in some cases this change makes it harder for the user of ndarray, or can we find a great solution?
If the dimension types of the parameters are all D, the result can also be written directly as D. But if the dimensions are not the same, it will be really troublesome. That's right, if the cost of using operators is that additional restrictions must be added, it maybe better to risk panic. I'm sorry I can't think of a better way for the time being. Thanks for pointing out.
bluss: Changed the title from "fix" since it's a new feature, not a fix. There is more feedback and thoughts to be given, when I have time. 🙂 
bluss: Here's a trick that actually works to resolve the `mean` situation. Just to show you a neat trick to use. Ignore the `broadcast_shape` method change for now - that one can be fixed later (debatable if it should be a method at all).

What we do here is that we put the `BroadcastShape<Self>` and `BroadcastShape<Ix0>` requirements into requirements of the `Dimension` itself. That way we rearrange the situation so that instead every type that implements Dimension has to prove - for itself - that it implements these two broadcasts. It seems to sidestep the overlapping problem.

See this temporary branch for that diff: https://github.com/rust-ndarray/ndarray/pull/new/bluss/pr-898-broadcast-shape
jturner314: > Here's a trick that actually works to resolve the `mean` situation.

That's a really nice solution. Is it possible to add `BroadcastShape<Self::Smaller, BroadcastOutput=Self>` and `BroadcastShape<Self::Larger, BroadcastOutput=Self::Larger>` bounds too? If so, I would think all four bounds would avoid a need for users to add `BroadcastShape` bounds to their functions in most cases.
SparrowLii: > What we do here is that we put the BroadcastShape<Self> and BroadcastShape<Ix0> requirements into requirements of the Dimension itself. That way we rearrange the situation so that instead every type that implements Dimension has to prove - for itself - that it implements these two broadcasts. It seems to sidestep the overlapping problem.

That is amazing! It's so surprising it works! However, we still can’t add statements like BroadcastShape<Self::Smaller, BroadcastOutput=Self> because it will cause loops in analysis passes. But the exciting thing is that I tried this way in the definition of Dimension:

    /// Next smaller dimension (if applicable)
    type Smaller: Dimension + BroadcastShape<Self, BroadcastOutput=Self>;
    /// Next larger dimension
    type Larger: Dimension + RemoveAxis + BroadcastShape<Self, BroadcastOutput=Self::Larger>;

And it did run successfully in my nightly version (cargo 1.49.0-nightly (d5556aeb8 2020-11-04) ), I think we can achieve this amazing work as soon as the 1.49 version is stable.
bluss: Note that `+ BroadcastShape<Self::Smaller, BroadcastOutput=Self>` does work it just has to be written like `+ BroadcastShape<<Self as Dimension>::Smaller, BroadcastOutput=Self>` to help the compiler.
jturner314: We could also add this bound to `RemoveAxis`: `BroadcastShape<<<Self as Dimension>::Smaller as Dimension>::Larger, BroadcastOutput = Self>`.
SparrowLii: I used uninitialized() to create an uninitialized array, and added the zip_mut_from_pair method to impl_methods.rs to get the result. Now no matter whether we need co_broadcast or not, we only need to traverse self and rhs once each. And ArcArray will not be turned into Array. But there is still a problem: both uninitialized() and maybe_uninit() methods require elements to implement the copy trait (I don’t know the reason behind this). So it is still unavoidable to change the signature of the method in impl_numeric.rs. 
bluss: uninitialized is hard to use correctly. I haven't removed it yet because I think I have managed to put it on sound ground - it *is possible* to use it soundly.

Unfortunately this method

1. Has some usage by actual users (github code search confirms)
2. This usage is from casual inspection almost always technically incorrect

Please use `maybe_uninit`. 🙂 If you do, then I don't have to go through and point out where your usage of `uninitialized` breaks the soundness rules 🙂 
SparrowLii: > Please use maybe_uninit that I mentioned. 🙂 If you do, then I don't have to go through and point out where your usage of uninitialized breaks the soundness rules 🙂

I tried to use maybe_uninit(), but I can’t think of a simple way to generate an ArrayBase with the data type MaybeUninit<A>. So I added a feature called MaybeUninitSubst to get the mapping from ArrayBase with element A to MaybeUnint<A>. In addition used Zip to get the results. 
bluss: Good fighting. The type related problems with the maybe_uninit constructor are not so fun, we're really feeling that we miss the "GAT" feature here (still experimental).
bluss: I'm working on fixing the maybe uninit functionality so that it can do what we need.
SparrowLii: > I'm working on fixing the maybe uninit functionality so that it can do what we need.

Cool! It will be very exciting if it and co_broadcast both land.
SparrowLii: @bluss Thank you so much for the amazing work! I have rebased the pr and used map_collect_owned according to your guide to avoid unsafe code. 
SparrowLii: > The PR is basically ready, a few minor things we want to change and then we merge. A few of the other issues mentioned can be fixed after merging.

Thank you very much for reviewing. It does have points that need improvement. I added a function called broadcast_with, but if this is not the focus at the moment, maybe we can consider it in the future.
bluss: Ok that's good, please make broadcast_with `pub(crate)` for now - otherwise we have to design name, signature and doc comment more. Let's get to the merge.
bluss: I'll just make a list of "known issues" that we can fix **after** the PR review and merge

- Rename BroadcastShape to `DimMax` - we also move the broadcast method out of the trait (it can just be a free function)
- It is a wish that the `D: Dimension + BroadcastShape<E>` bound in arith ops becomes just `D: Dimension`, i.e. the DimMax with associated type becomes implied by `D: Dimension`. At this point it is not certain how this can be done.
- Preserve the allocation of the input array in arithmetic ops in more cases (when only the left or right side needs broadcast)
- Improve performance of "no-broadcasting-needed" scenario in &array + &array operation
- Potentially provide public function or method similar to `broadcast_with` - that does the broadcasting of two arrays together.
SparrowLii: @bluss It has been modified accordingly.  I very much look forward to our ability to solve the above problems in the future. Especially removing the restriction like BroadCast<E>, it may represent a new language feature IMO, which will be so cool and exciting.

Patch:
diff --git a/src/data_traits.rs b/src/data_traits.rs
index 1e191c46..7ac63d54 100644
--- a/src/data_traits.rs
+++ b/src/data_traits.rs
@@ -16,9 +16,7 @@ use std::ptr::NonNull;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
 
-use crate::{
-    ArrayBase, CowRepr, Dimension, OwnedArcRepr, OwnedRepr, RawViewRepr, ViewRepr,
-};
+use crate::{ArrayBase, CowRepr, Dimension, OwnedArcRepr, OwnedRepr, RawViewRepr, ViewRepr};
 
 /// Array representation trait.
 ///
@@ -414,7 +412,6 @@ pub unsafe trait DataOwned: Data {
     /// Corresponding owned data with MaybeUninit elements
     type MaybeUninit: DataOwned<Elem = MaybeUninit<Self::Elem>>
         + RawDataSubst<Self::Elem, Output=Self>;
-
     #[doc(hidden)]
     fn new(elements: Vec<Self::Elem>) -> Self;
 
@@ -440,6 +437,7 @@ unsafe impl<A> DataOwned for OwnedRepr<A> {
     fn new(elements: Vec<A>) -> Self {
         OwnedRepr::from(elements)
     }
+
     fn into_shared(self) -> OwnedArcRepr<A> {
         OwnedArcRepr(Arc::new(self))
     }
@@ -622,3 +620,4 @@ impl<'a, A: 'a, B: 'a> RawDataSubst<B> for ViewRepr<&'a mut A> {
         ViewRepr::new()
     }
 }
+
diff --git a/src/dimension/broadcast.rs b/src/dimension/broadcast.rs
new file mode 100644
index 00000000..dc1513f0
--- /dev/null
+++ b/src/dimension/broadcast.rs
@@ -0,0 +1,143 @@
+use crate::error::*;
+use crate::{Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
+
+/// Calculate the common shape for a pair of array shapes, that they can be broadcasted
+/// to. Return an error if the shapes are not compatible.
+///
+/// Uses the [NumPy broadcasting rules]
+//  (https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html#general-broadcasting-rules).
+pub(crate) fn co_broadcast<D1, D2, Output>(shape1: &D1, shape2: &D2) -> Result<Output, ShapeError>
+where
+    D1: Dimension,
+    D2: Dimension,
+    Output: Dimension,
+{
+    let (k, overflow) = shape1.ndim().overflowing_sub(shape2.ndim());
+    // Swap the order if d2 is longer.
+    if overflow {
+        return co_broadcast::<D2, D1, Output>(shape2, shape1);
+    }
+    // The output should be the same length as shape1.
+    let mut out = Output::zeros(shape1.ndim());
+    for (out, s) in izip!(out.slice_mut(), shape1.slice()) {
+        *out = *s;
+    }
+    for (out, s2) in izip!(&mut out.slice_mut()[k..], shape2.slice()) {
+        if *out != *s2 {
+            if *out == 1 {
+                *out = *s2
+            } else if *s2 != 1 {
+                return Err(from_kind(ErrorKind::IncompatibleShape));
+            }
+        }
+    }
+    Ok(out)
+}
+
+pub trait DimMax<Other: Dimension> {
+    /// The resulting dimension type after broadcasting.
+    type Output: Dimension;
+}
+
+/// Dimensions of the same type remain unchanged when co_broadcast.
+/// So you can directly use D as the resulting type.
+/// (Instead of <D as DimMax<D>>::BroadcastOutput)
+impl<D: Dimension> DimMax<D> for D {
+    type Output = D;
+}
+
+macro_rules! impl_broadcast_distinct_fixed {
+    ($smaller:ty, $larger:ty) => {
+        impl DimMax<$larger> for $smaller {
+            type Output = $larger;
+        }
+
+        impl DimMax<$smaller> for $larger {
+            type Output = $larger;
+        }
+    };
+}
+
+impl_broadcast_distinct_fixed!(Ix0, Ix1);
+impl_broadcast_distinct_fixed!(Ix0, Ix2);
+impl_broadcast_distinct_fixed!(Ix0, Ix3);
+impl_broadcast_distinct_fixed!(Ix0, Ix4);
+impl_broadcast_distinct_fixed!(Ix0, Ix5);
+impl_broadcast_distinct_fixed!(Ix0, Ix6);
+impl_broadcast_distinct_fixed!(Ix1, Ix2);
+impl_broadcast_distinct_fixed!(Ix1, Ix3);
+impl_broadcast_distinct_fixed!(Ix1, Ix4);
+impl_broadcast_distinct_fixed!(Ix1, Ix5);
+impl_broadcast_distinct_fixed!(Ix1, Ix6);
+impl_broadcast_distinct_fixed!(Ix2, Ix3);
+impl_broadcast_distinct_fixed!(Ix2, Ix4);
+impl_broadcast_distinct_fixed!(Ix2, Ix5);
+impl_broadcast_distinct_fixed!(Ix2, Ix6);
+impl_broadcast_distinct_fixed!(Ix3, Ix4);
+impl_broadcast_distinct_fixed!(Ix3, Ix5);
+impl_broadcast_distinct_fixed!(Ix3, Ix6);
+impl_broadcast_distinct_fixed!(Ix4, Ix5);
+impl_broadcast_distinct_fixed!(Ix4, Ix6);
+impl_broadcast_distinct_fixed!(Ix5, Ix6);
+impl_broadcast_distinct_fixed!(Ix0, IxDyn);
+impl_broadcast_distinct_fixed!(Ix1, IxDyn);
+impl_broadcast_distinct_fixed!(Ix2, IxDyn);
+impl_broadcast_distinct_fixed!(Ix3, IxDyn);
+impl_broadcast_distinct_fixed!(Ix4, IxDyn);
+impl_broadcast_distinct_fixed!(Ix5, IxDyn);
+impl_broadcast_distinct_fixed!(Ix6, IxDyn);
+
+
+#[cfg(test)]
+#[cfg(feature = "std")]
+mod tests {
+    use super::co_broadcast;
+    use crate::{Dimension, Dim, DimMax, ShapeError, Ix0, IxDynImpl, ErrorKind};
+
+    #[test]
+    fn test_broadcast_shape() {
+        fn test_co<D1, D2>(
+            d1: &D1,
+            d2: &D2,
+            r: Result<<D1 as DimMax<D2>>::Output, ShapeError>,
+        ) where
+            D1: Dimension + DimMax<D2>,
+            D2: Dimension,
+        {
+            let d = co_broadcast::<D1, D2, <D1 as DimMax<D2>>::Output>(&d1, d2);
+            assert_eq!(d, r);
+        }
+        test_co(&Dim([2, 3]), &Dim([4, 1, 3]), Ok(Dim([4, 2, 3])));
+        test_co(
+            &Dim([1, 2, 2]),
+            &Dim([1, 3, 4]),
+            Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)),
+        );
+        test_co(&Dim([3, 4, 5]), &Ix0(), Ok(Dim([3, 4, 5])));
+        let v = vec![1, 2, 3, 4, 5, 6, 7];
+        test_co(
+            &Dim(vec![1, 1, 3, 1, 5, 1, 7]),
+            &Dim([2, 1, 4, 1, 6, 1]),
+            Ok(Dim(IxDynImpl::from(v.as_slice()))),
+        );
+        let d = Dim([1, 2, 1, 3]);
+        test_co(&d, &d, Ok(d));
+        test_co(
+            &Dim([2, 1, 2]).into_dyn(),
+            &Dim(0),
+            Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)),
+        );
+        test_co(
+            &Dim([2, 1, 1]),
+            &Dim([0, 0, 1, 3, 4]),
+            Ok(Dim([0, 0, 2, 3, 4])),
+        );
+        test_co(&Dim([0]), &Dim([0, 0, 0]), Ok(Dim([0, 0, 0])));
+        test_co(&Dim(1), &Dim([1, 0, 0]), Ok(Dim([1, 0, 0])));
+        test_co(
+            &Dim([1, 3, 0, 1, 1]),
+            &Dim([1, 2, 3, 1]),
+            Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)),
+        );
+    }
+}
diff --git a/src/dimension/dimension_trait.rs b/src/dimension/dimension_trait.rs
index 92f24118..6007f93a 100644
--- a/src/dimension/dimension_trait.rs
+++ b/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, DimMax};
 use crate::IntoDimension;
 use crate::RemoveAxis;
 use crate::{ArrayView1, ArrayViewMut1};
@@ -46,6 +46,11 @@ pub trait Dimension:
     + MulAssign
     + for<'x> MulAssign<&'x Self>
     + MulAssign<usize>
+    + DimMax<Ix0, Output=Self>
+    + DimMax<Self, Output=Self>
+    + DimMax<IxDyn, Output=IxDyn>
+    + DimMax<<Self as Dimension>::Smaller, Output=Self>
+    + DimMax<<Self as Dimension>::Larger, Output=<Self as Dimension>::Larger>
 {
     /// For fixed-size dimension representations (e.g. `Ix2`), this should be
     /// `Some(ndim)`, and for variable-size dimension representations (e.g.
diff --git a/src/dimension/mod.rs b/src/dimension/mod.rs
index 1359b8f3..2505681b 100644
--- a/src/dimension/mod.rs
+++ b/src/dimension/mod.rs
@@ -12,6 +12,7 @@ use num_integer::div_floor;
 
 pub use self::axes::{axes_of, Axes, AxisDescription};
 pub use self::axis::Axis;
+pub use self::broadcast::DimMax;
 pub use self::conversion::IntoDimension;
 pub use self::dim::*;
 pub use self::dimension_trait::Dimension;
@@ -28,6 +29,7 @@ use std::mem;
 mod macros;
 mod axes;
 mod axis;
+pub(crate) mod broadcast;
 mod conversion;
 pub mod dim;
 mod dimension_trait;
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index 9943e7c8..958fc3f1 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -14,14 +14,15 @@ use rawpointer::PointerExt;
 
 use crate::imp_prelude::*;
 
-use crate::arraytraits;
+use crate::{arraytraits, DimMax};
 use crate::dimension;
 use crate::dimension::IntoDimension;
 use crate::dimension::{
     abs_index, axes_of, do_slice, merge_axes, move_min_stride_axis_to_last,
     offset_from_ptr_to_memory, size_of_shape_checked, stride_offset, Axes,
 };
-use crate::error::{self, ErrorKind, ShapeError};
+use crate::dimension::broadcast::co_broadcast;
+use crate::error::{self, ErrorKind, ShapeError, from_kind};
 use crate::math_cell::MathCell;
 use crate::itertools::zip;
 use crate::zip::Zip;
@@ -1766,6 +1767,28 @@ where
         unsafe { Some(ArrayView::new(self.ptr, dim, broadcast_strides)) }
     }
 
+    /// For two arrays or views, find their common shape if possible and
+    /// broadcast them as array views into that shape.
+    ///
+    /// Return `ShapeError` if their shapes can not be broadcast together.
+    #[allow(clippy::type_complexity)]
+    pub(crate) fn broadcast_with<'a, 'b, B, S2, E>(&'a self, other: &'b ArrayBase<S2, E>) ->
+        Result<(ArrayView<'a, A, DimMaxOf<D, E>>, ArrayView<'b, B, DimMaxOf<D, E>>), ShapeError>
+    where
+        S: Data<Elem=A>,
+        S2: Data<Elem=B>,
+        D: Dimension + DimMax<E>,
+        E: Dimension,
+    {
+        let shape = co_broadcast::<D, E, <D as DimMax<E>>::Output>(&self.dim, &other.dim)?;
+        if let Some(view1) = self.broadcast(shape.clone()) {
+            if let Some(view2) = other.broadcast(shape) {
+                return Ok((view1, view2));
+            }
+        }
+        Err(from_kind(ErrorKind::IncompatibleShape))
+    }
+
     /// Swap axes `ax` and `bx`.
     ///
     /// This does not move any data, it just adjusts the array’s dimensions
@@ -2013,7 +2036,7 @@ where
         self.map_inplace(move |elt| *elt = x.clone());
     }
 
-    fn zip_mut_with_same_shape<B, S2, E, F>(&mut self, rhs: &ArrayBase<S2, E>, mut f: F)
+    pub(crate) fn zip_mut_with_same_shape<B, S2, E, F>(&mut self, rhs: &ArrayBase<S2, E>, mut f: F)
     where
         S: DataMut,
         S2: Data<Elem = B>,
@@ -2443,3 +2466,5 @@ unsafe fn unlimited_transmute<A, B>(data: A) -> B {
     let old_data = ManuallyDrop::new(data);
     (&*old_data as *const A as *const B).read()
 }
+
+type DimMaxOf<A, B> = <A as DimMax<B>>::Output;
diff --git a/src/impl_ops.rs b/src/impl_ops.rs
index 256bee3e..d38cb566 100644
--- a/src/impl_ops.rs
+++ b/src/impl_ops.rs
@@ -6,6 +6,8 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use crate::dimension::DimMax;
+use crate::Zip;
 use num_complex::Complex;
 
 /// Elements that can be used as direct operands in arithmetic with arrays.
@@ -53,11 +55,11 @@ macro_rules! impl_binary_op(
 /// Perform elementwise
 #[doc=$doc]
 /// between `self` and `rhs`,
-/// and return the result (based on `self`).
+/// and return the result.
 ///
 /// `self` must be an `Array` or `ArcArray`.
 ///
-/// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
+/// If their shapes disagree, `self` is broadcast to their broadcast shape.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<A, B, S, S2, D, E> $trt<ArrayBase<S2, E>> for ArrayBase<S, D>
@@ -66,11 +68,11 @@ where
     B: Clone,
     S: DataOwned<Elem=A> + DataMut,
     S2: Data<Elem=B>,
-    D: Dimension,
+    D: Dimension + DimMax<E>,
     E: Dimension,
 {
-    type Output = ArrayBase<S, D>;
-    fn $mth(self, rhs: ArrayBase<S2, E>) -> ArrayBase<S, D>
+    type Output = ArrayBase<S, <D as DimMax<E>>::Output>;
+    fn $mth(self, rhs: ArrayBase<S2, E>) -> Self::Output
     {
         self.$mth(&rhs)
     }
@@ -79,9 +81,12 @@ where
 /// Perform elementwise
 #[doc=$doc]
 /// between `self` and reference `rhs`,
-/// and return the result (based on `self`).
+/// and return the result.
+///
+/// `rhs` must be an `Array` or `ArcArray`.
 ///
-/// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
+/// If their shapes disagree, `self` is broadcast to their broadcast shape,
+/// cloning the data if needed.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<'a, A, B, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for ArrayBase<S, D>
@@ -90,16 +95,55 @@ where
     B: Clone,
     S: DataOwned<Elem=A> + DataMut,
     S2: Data<Elem=B>,
-    D: Dimension,
+    D: Dimension + DimMax<E>,
     E: Dimension,
 {
-    type Output = ArrayBase<S, D>;
-    fn $mth(mut self, rhs: &ArrayBase<S2, E>) -> ArrayBase<S, D>
+    type Output = ArrayBase<S, <D as DimMax<E>>::Output>;
+    fn $mth(self, rhs: &ArrayBase<S2, E>) -> Self::Output
     {
-        self.zip_mut_with(rhs, |x, y| {
-            *x = x.clone() $operator y.clone();
-        });
-        self
+        if self.ndim() == rhs.ndim() && self.shape() == rhs.shape() {
+            let mut out = self.into_dimensionality::<<D as DimMax<E>>::Output>().unwrap();
+            out.zip_mut_with_same_shape(rhs, clone_iopf(A::$mth));
+            out
+        } else {
+            let (lhs, rhs) = self.broadcast_with(rhs).unwrap();
+            Zip::from(&lhs).and(&rhs).map_collect_owned(clone_opf(A::$mth))
+        }
+    }
+}
+
+/// Perform elementwise
+#[doc=$doc]
+/// between reference `self` and `rhs`,
+/// and return the result.
+///
+/// `rhs` must be an `Array` or `ArcArray`.
+///
+/// If their shapes disagree, `self` is broadcast to their broadcast shape,
+/// cloning the data if needed.
+///
+/// **Panics** if broadcasting isn’t possible.
+impl<'a, A, B, S, S2, D, E> $trt<ArrayBase<S2, E>> for &'a ArrayBase<S, D>
+where
+    A: Clone + $trt<B, Output=B>,
+    B: Clone,
+    S: Data<Elem=A>,
+    S2: DataOwned<Elem=B> + DataMut,
+    D: Dimension,
+    E: Dimension + DimMax<D>,
+{
+    type Output = ArrayBase<S2, <E as DimMax<D>>::Output>;
+    fn $mth(self, rhs: ArrayBase<S2, E>) -> Self::Output
+    where
+    {
+        if self.ndim() == rhs.ndim() && self.shape() == rhs.shape() {
+            let mut out = rhs.into_dimensionality::<<E as DimMax<D>>::Output>().unwrap();
+            out.zip_mut_with_same_shape(self, clone_iopf_rev(A::$mth));
+            out
+        } else {
+            let (rhs, lhs) = rhs.broadcast_with(self).unwrap();
+            Zip::from(&lhs).and(&rhs).map_collect_owned(clone_opf(A::$mth))
+        }
     }
 }
 
@@ -108,7 +152,8 @@ where
 /// between references `self` and `rhs`,
 /// and return the result as a new `Array`.
 ///
-/// If their shapes disagree, `rhs` is broadcast to the shape of `self`.
+/// If their shapes disagree, `self` and `rhs` is broadcast to their broadcast shape,
+/// cloning the data if needed.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<'a, A, B, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for &'a ArrayBase<S, D>
@@ -117,13 +162,13 @@ where
     B: Clone,
     S: Data<Elem=A>,
     S2: Data<Elem=B>,
-    D: Dimension,
+    D: Dimension + DimMax<E>,
     E: Dimension,
 {
-    type Output = Array<A, D>;
-    fn $mth(self, rhs: &'a ArrayBase<S2, E>) -> Array<A, D> {
-        // FIXME: Can we co-broadcast arrays here? And how?
-        self.to_owned().$mth(rhs)
+    type Output = Array<A, <D as DimMax<E>>::Output>;
+    fn $mth(self, rhs: &'a ArrayBase<S2, E>) -> Self::Output {
+        let (lhs, rhs) = self.broadcast_with(rhs).unwrap();
+        Zip::from(&lhs).and(&rhs).map_collect(clone_opf(A::$mth))
     }
 }
 
@@ -228,6 +273,18 @@ mod arithmetic_ops {
     use num_complex::Complex;
     use std::ops::*;
 
+    fn clone_opf<A: Clone, B: Clone, C>(f: impl Fn(A, B) -> C) -> impl FnMut(&A, &B) -> C {
+        move |x, y| f(x.clone(), y.clone())
+    }
+
+    fn clone_iopf<A: Clone, B: Clone>(f: impl Fn(A, B) -> A) -> impl FnMut(&mut A, &B) {
+        move |x, y| *x = f(x.clone(), y.clone())
+    }
+
+    fn clone_iopf_rev<A: Clone, B: Clone>(f: impl Fn(A, B) -> B) -> impl FnMut(&mut B, &A) {
+        move |x, y| *x = f(y.clone(), x.clone())
+    }
+
     impl_binary_op!(Add, +, add, +=, "addition");
     impl_binary_op!(Sub, -, sub, -=, "subtraction");
     impl_binary_op!(Mul, *, mul, *=, "multiplication");
diff --git a/src/lib.rs b/src/lib.rs
index 9cd7dc3f..c079b481 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -134,6 +134,7 @@ use std::marker::PhantomData;
 use alloc::sync::Arc;
 
 pub use crate::dimension::dim::*;
+pub use crate::dimension::DimMax;
 pub use crate::dimension::{Axis, AxisDescription, Dimension, IntoDimension, RemoveAxis};
 
 pub use crate::dimension::IxDynImpl;


Test patch:
diff --git a/tests/array.rs b/tests/array.rs
index 6b72bb5c..b0a28ca4 100644
--- a/tests/array.rs
+++ b/tests/array.rs
@@ -9,9 +9,9 @@
 
 use defmac::defmac;
 use itertools::{enumerate, zip, Itertools};
-use ndarray::indices;
 use ndarray::prelude::*;
 use ndarray::{arr3, rcarr2};
+use ndarray::indices;
 use ndarray::{Slice, SliceInfo, SliceOrIndex};
 
 macro_rules! assert_panics {
@@ -1565,6 +1565,49 @@ fn arithmetic_broadcast() {
     a.swap_axes(0, 1);
     let b = a.clone() / aview0(&1.);
     assert_eq!(a, b);
+
+    // reference
+    let a = arr2(&[[2], [3], [4]]);
+    let b = arr1(&[5, 6, 7]);
+    assert_eq!(&a + &b, arr2(&[[7, 8, 9], [8, 9, 10], [9, 10, 11]]));
+    assert_eq!(
+        a.clone() - &b,
+        arr2(&[[-3, -4, -5], [-2, -3, -4], [-1, -2, -3]])
+    );
+    assert_eq!(
+        a.clone() * b.clone(),
+        arr2(&[[10, 12, 14], [15, 18, 21], [20, 24, 28]])
+    );
+    assert_eq!(&b / a, arr2(&[[2, 3, 3], [1, 2, 2], [1, 1, 1]]));
+
+    // Negative strides and non-contiguous memory
+    let s = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
+    let s = Array3::from_shape_vec((2, 3, 2).strides((1, 4, 2)), s.to_vec()).unwrap();
+    let a = s.slice(s![..;-1,..;2,..]);
+    let b = s.slice(s![..2, -1, ..]);
+    let mut c = s.clone();
+    c.collapse_axis(Axis(2), 1);
+    let c = c.slice(s![1,..;2,..]);
+    assert_eq!(
+        &a.to_owned() + &b,
+        arr3(&[[[11, 15], [20, 24]], [[10, 14], [19, 23]]])
+    );
+    assert_eq!(
+        &a + b.into_owned() + c,
+        arr3(&[[[15, 19], [32, 36]], [[14, 18], [31, 35]]])
+    );
+
+    // shared array
+    let sa = a.to_shared();
+    let sa2 = sa.to_shared();
+    let sb = b.to_shared();
+    let sb2 = sb.to_shared();
+    let sc = c.to_shared();
+    let sc2 = sc.into_shared();
+    assert_eq!(
+        sa2 + &sb2 + sc2.into_owned(),
+        arr3(&[[[15, 19], [32, 36]], [[14, 18], [31, 35]]])
+    );
 }
 
 #[test]