-
Notifications
You must be signed in to change notification settings - Fork 311
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add accumulate_axis_inplace method and implement NdProducer for RawArrayView/Mut #611
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -2231,4 +2231,60 @@ where | |
}) | ||
} | ||
} | ||
|
||
/// Iterates over pairs of consecutive elements along the axis. | ||
/// | ||
/// The first argument to the closure is an element, and the second | ||
/// argument is the next element along the axis. Iteration is guaranteed to | ||
/// proceed in order along the specified axis, but in all other respects | ||
/// the iteration order is unspecified. | ||
/// | ||
/// # Example | ||
/// | ||
/// For example, this can be used to compute the cumulative sum along an | ||
/// axis: | ||
/// | ||
/// ``` | ||
/// use ndarray::{array, Axis}; | ||
/// | ||
/// let mut arr = array![ | ||
/// [[1, 2], [3, 4], [5, 6]], | ||
/// [[7, 8], [9, 10], [11, 12]], | ||
/// ]; | ||
/// arr.accumulate_axis_inplace(Axis(1), |&prev, curr| *curr += prev); | ||
/// assert_eq!( | ||
/// arr, | ||
/// array![ | ||
/// [[1, 2], [4, 6], [9, 12]], | ||
/// [[7, 8], [16, 18], [27, 30]], | ||
/// ], | ||
/// ); | ||
/// ``` | ||
pub fn accumulate_axis_inplace<F>(&mut self, axis: Axis, mut f: F) | ||
where | ||
F: FnMut(&A, &mut A), | ||
S: DataMut, | ||
{ | ||
if self.len_of(axis) <= 1 { | ||
return; | ||
} | ||
let mut curr = self.raw_view_mut(); // mut borrow of the array here | ||
let mut prev = curr.raw_view(); // derive further raw views from the same borrow | ||
prev.slice_axis_inplace(axis, Slice::from(..-1)); | ||
curr.slice_axis_inplace(axis, Slice::from(1..)); | ||
// This implementation relies on `Zip` iterating along `axis` in order. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good point, I guess we'll think about if that's something we can guarantee. Probably not publically. |
||
Zip::from(prev).and(curr).apply(|prev, curr| unsafe { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is a particularly simple Zip - the prev and curr pointer are always just the stride of |
||
// These pointer dereferences and borrows are safe because: | ||
// | ||
// 1. They're pointers to elements in the array. | ||
// | ||
// 2. `S: DataMut` guarantees that elements are safe to borrow | ||
// mutably and that they don't alias. | ||
// | ||
// 3. The lifetimes of the borrows last only for the duration | ||
// of the call to `f`, so aliasing across calls to `f` | ||
// cannot occur. | ||
f(&*prev, &mut *curr) | ||
}); | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This method is a neat demonstration of an array Zip-ing with itself. It seems a bit niche but we can of course include it. But the question of course goes out, how do we generalize this. Apart from the raw pointer Zip, of course? How to generalize the Self-zip in a safe Rust package?
Apart from returning to the idea of an array view of Cells, which is not hard to do.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we really need to wait until Rust has generic associated types before we can write good self-Zip abstractions. We can do things like this today:
but without generic associated types, we can't abstract this into a
WindowsMut
producer that implements a trait, because the item type (which is an associated type of the trait) has a temporary lifetime.