Repository: rust-ndarray/ndarray (rust)

Issue:
and_broadcast_mut for Zip
Currently Zip has no mutable equivalent for the `and_broadcast` method, which is presumably to allow it to safely be `split()` and parallelized. However this makes some calculations difficult. The examples below are for backprop through a broadcast-able element-wise multiplication, but the problem appears as a pattern often.

Currently I'm looking at using this:
```
unsafe{
	let input2_grad = data.get_mut(&self.input2_id.gradient_id())?;

	// do not split/parallelise this Zip!
	Zip::from(&input1)
		.and(&output_grad)
		.and_broadcast(&input2_grad)
		.apply(|input1, out_grad, input2_grad| {
			let input2_grad = input2_grad as *const f32 as *mut f32;
			*input2_grad += input1 * out_grad;
		});
}
```
which works correctly, but also has UB written all over it depending on how Zip is implemented internally (I realize Zip owes me nothing here). 

The safe alternative I've found is:
```
let iter = input1.exact_chunks(input2.shape()).into_iter()
	.zip(output_grad.exact_chunks(input2.shape()));

for (input1_chunk, out_grad_chunk) in iter {
	Zip::from(&mut input2_grad)
		.and(&input1_chunk)
		.and(&out_grad_chunk)
		.apply(|input2_grad, input1, out_grad| {
			*input2_grad += input1 * out_grad;
		});
}
```
which is catastrophically (over 30x) slower for some cases.

So my questions are:
 * is there a safe high performance workaround that I've missed?
 * would it be possible to add a `and_broadcast_mut()` method which would set a flag or something which makes `split()` panic?

Pull request description:
Add methods .cell_view() and .into_cell_view()
Add `ArrayBase::cell_view` (similar to methods `.view()` and `.view_mut()`) and the conversion method `ArrayViewMut::into_cell_view()`

We've wanted to add these methods since we added `.cast::<T>()` to the raw views.

The translation of `ArrayViewMut<'a, T, D>` -> `ArrayView<'a, MathCell<T>, D>` works under the same principle
as the method `Cell:from_mut` in the standard library.

These methods should give a lot of flexibility - suddenly the same array view can be traversed multiple times
in the same `Zip` and mutated while being traversed in a different way than before. The price to pay for this
is using the Cell methods to manipulate the elements.

We use a new cell type - `MathCell` - which is just a transparent wrapper around `Cell`. This is so that
we can implement arithmetic ops for cells. However, finding the right implementations requires more design
work so that will be in a later PR.

This fixes #478 - the use case is fixed by using `Zip::from(a).and_broadcast(b.cell_view())` - i.e composing `cell_view`
with the existing broadcasting method.

Patch:
diff --git a/src/argument_traits.rs b/src/argument_traits.rs
index a93e33a1..82d4869a 100644
--- a/src/argument_traits.rs
+++ b/src/argument_traits.rs
@@ -1,6 +1,7 @@
 use std::cell::Cell;
 use std::mem::MaybeUninit;
 
+use crate::math_cell::MathCell;
 
 /// A producer element that can be assigned to once
 pub trait AssignElem<T> {
@@ -22,6 +23,13 @@ impl<'a, T> AssignElem<T> for &'a Cell<T> {
     }
 }
 
+/// Assignable element, simply `self.set(input)`.
+impl<'a, T> AssignElem<T> for &'a MathCell<T> {
+    fn assign_elem(self, input: T) {
+        self.set(input);
+    }
+}
+
 /// Assignable element, the item in the MaybeUninit is overwritten (prior value, if any, is not
 /// read or dropped).
 impl<'a, T> AssignElem<T> for &'a mut MaybeUninit<T> {
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index cccc195a..df48ea1c 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -20,6 +20,7 @@ use crate::dimension::{
     abs_index, axes_of, do_slice, merge_axes, size_of_shape_checked, stride_offset, Axes,
 };
 use crate::error::{self, ErrorKind, ShapeError};
+use crate::math_cell::MathCell;
 use crate::itertools::zip;
 use crate::zip::Zip;
 
@@ -151,6 +152,20 @@ where
         unsafe { ArrayViewMut::new(self.ptr, self.dim.clone(), self.strides.clone()) }
     }
 
+    /// Return a shared view of the array with elements as if they were embedded in cells.
+    ///
+    /// The cell view requires a mutable borrow of the array. Once borrowed the
+    /// cell view itself can be copied and accessed without exclusivity.
+    ///
+    /// The view acts "as if" the elements are temporarily in cells, and elements
+    /// can be changed through shared references using the regular cell methods.
+    pub fn cell_view(&mut self) -> ArrayView<'_, MathCell<A>, D>
+    where
+        S: DataMut,
+    {
+        self.view_mut().into_cell_view()
+    }
+
     /// Return an uniquely owned copy of the array.
     ///
     /// If the input array is contiguous and its strides are positive, then the
diff --git a/src/impl_views/conversions.rs b/src/impl_views/conversions.rs
index 863d26dd..8c423010 100644
--- a/src/impl_views/conversions.rs
+++ b/src/impl_views/conversions.rs
@@ -13,6 +13,7 @@ use crate::imp_prelude::*;
 use crate::{Baseiter, ElementsBase, ElementsBaseMut, Iter, IterMut};
 
 use crate::iter::{self, AxisIter, AxisIterMut};
+use crate::math_cell::MathCell;
 use crate::IndexLonger;
 
 /// Methods for read-only array views.
@@ -117,6 +118,21 @@ where
     pub fn into_slice(self) -> Option<&'a mut [A]> {
         self.try_into_slice().ok()
     }
+
+    /// Return a shared view of the array with elements as if they were embedded in cells.
+    ///
+    /// The cell view itself can be copied and accessed without exclusivity.
+    ///
+    /// The view acts "as if" the elements are temporarily in cells, and elements
+    /// can be changed through shared references using the regular cell methods.
+    pub fn into_cell_view(self) -> ArrayView<'a, MathCell<A>, D> {
+        // safety: valid because
+        // A and MathCell<A> have the same representation
+        // &'a mut T is interchangeable with &'a Cell<T> -- see method Cell::from_mut in std
+        unsafe {
+            self.into_raw_view_mut().cast::<MathCell<A>>().deref_into_view()
+        }
+    }
 }
 
 /// Private array view methods
diff --git a/src/lib.rs b/src/lib.rs
index c0a3b2e8..89862c78 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -139,6 +139,7 @@ pub use crate::linalg_traits::LinalgScalar;
 
 pub use crate::stacking::{concatenate, stack, stack_new_axis};
 
+pub use crate::math_cell::MathCell;
 pub use crate::impl_views::IndexLonger;
 pub use crate::shape_builder::{Shape, ShapeBuilder, StrideShape};
 
@@ -180,6 +181,7 @@ mod layout;
 mod linalg_traits;
 mod linspace;
 mod logspace;
+mod math_cell;
 mod numeric_util;
 mod partial;
 mod shape_builder;
diff --git a/src/math_cell.rs b/src/math_cell.rs
new file mode 100644
index 00000000..f0f8da40
--- /dev/null
+++ b/src/math_cell.rs
@@ -0,0 +1,102 @@
+
+use std::cell::Cell;
+use std::cmp::Ordering;
+use std::fmt;
+
+use std::ops::{Deref, DerefMut};
+
+/// A transparent wrapper of [`Cell<T>`](std::cell::Cell) which is identical in every way, except
+/// it will implement arithmetic operators as well.
+///
+/// The purpose of `MathCell` is to be used from [.cell_view()](crate::ArrayBase::cell_view).
+/// The `MathCell` derefs to `Cell`, so all the cell's methods are available.
+#[repr(transparent)]
+#[derive(Default)]
+pub struct MathCell<T>(Cell<T>);
+
+impl<T> MathCell<T> {
+    /// Create a new cell with the given value
+    #[inline(always)]
+    pub const fn new(value: T) -> Self { MathCell(Cell::new(value)) }
+
+    /// Return the inner value
+    pub fn into_inner(self) -> T { Cell::into_inner(self.0) }
+
+    /// Swap value with another cell
+    pub fn swap(&self, other: &Self) {
+        Cell::swap(&self.0, &other.0)
+    }
+}
+
+impl<T> Deref for MathCell<T> {
+    type Target = Cell<T>;
+    #[inline(always)]
+    fn deref(&self) -> &Self::Target { &self.0 }
+}
+
+impl<T> DerefMut for MathCell<T> {
+    #[inline(always)]
+    fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 }
+}
+
+impl<T> Clone for MathCell<T>
+    where T: Copy
+{
+    fn clone(&self) -> Self {
+        MathCell::new(self.get())
+    }
+}
+
+impl<T> PartialEq for MathCell<T>
+    where T: Copy + PartialEq
+{
+    fn eq(&self, rhs: &Self) -> bool {
+        self.get() == rhs.get()
+    }
+}
+
+impl<T> Eq for MathCell<T>
+    where T: Copy + Eq
+{ }
+
+impl<T> PartialOrd for MathCell<T>
+    where T: Copy + PartialOrd
+{
+    fn partial_cmp(&self, rhs: &Self) -> Option<Ordering> {
+        self.get().partial_cmp(&rhs.get())
+    }
+
+    fn lt(&self, rhs: &Self) -> bool { self.get().lt(&rhs.get()) }
+    fn le(&self, rhs: &Self) -> bool { self.get().le(&rhs.get()) }
+    fn gt(&self, rhs: &Self) -> bool { self.get().gt(&rhs.get()) }
+    fn ge(&self, rhs: &Self) -> bool { self.get().ge(&rhs.get()) }
+}
+
+impl<T> Ord for MathCell<T>
+    where T: Copy + Ord
+{
+    fn cmp(&self, rhs: &Self) -> Ordering {
+        self.get().cmp(&rhs.get())
+    }
+}
+
+impl<T> fmt::Debug for MathCell<T>
+    where T: Copy + fmt::Debug
+{
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        self.get().fmt(f)
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::MathCell;
+
+    #[test]
+    fn test_basic() {
+        let c = &MathCell::new(0);
+        c.set(1);
+        assert_eq!(c.get(), 1);
+    }
+}


Test patch:
diff --git a/tests/views.rs b/tests/views.rs
new file mode 100644
index 00000000..216e5540
--- /dev/null
+++ b/tests/views.rs
@@ -0,0 +1,16 @@
+use ndarray::prelude::*;
+use ndarray::Zip;
+
+#[test]
+fn cell_view() {
+    let mut a = Array::from_shape_fn((10, 5), |(i, j)| (i * j) as f32);
+    let answer = &a + 1.;
+
+    {
+        let cv1 = a.cell_view();
+        let cv2 = cv1;
+
+        Zip::from(cv1).and(cv2).apply(|a, b| a.set(b.get() + 1.));
+    }
+    assert_eq!(a, answer);
+}