-
Notifications
You must be signed in to change notification settings - Fork 310
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
Inplace operations (sum_axis, ...) #425
Comments
This is one way to get the raw
I'm curious how your implementation works, though. It seems to me that in general it's a difficult problem to be able to shrink the underlying
then in row-major memory layout, this would be
in memory, and for column-major memory layout it would be
In general, you can have arrays with weird custom memory layouts with arbitrary strides. Let's say you're trying to perform a
If you want to be able to shrink the underlying
You can then shrink the
in memory. However, you need the results in the first three elements in memory, not scattered throughout, in order to be able to shrink the Another approach is to immediately place the results in the first three elements while being careful about the sequence you perform the in-place sums. You do have to be careful, though, because for example you don't want to overwrite
which you can shrink to the first three elements. For general strides, this becomes a very complex problem to avoid overwriting elements with results before you need to read them to compute sums. There's probably some algorithm that would work by swapping elements around to keep them from being overwritten before they're needed, but I'm not really sure. All this becomes much simpler if you don't care about shrinking the underlying #[macro_use]
extern crate ndarray;
extern crate num_traits;
use ndarray::prelude::*;
use ndarray::RemoveAxis;
use num_traits::Zero;
use std::ops::{Add, AddAssign};
/// Sums along the axis without allocating a new array.
///
/// **Panics** if `axis` is out-of-bounds or if `arr.len_of(axis) == 0`.
fn sum_axis_inplace<A, D>(mut arr: Array<A, D>, axis: Axis) -> Array<A, D::Smaller>
where
A: Clone + Add<Output = A> + AddAssign + Zero,
D: Dimension + RemoveAxis,
{
{
// Get mutable views of the first subview along `axis` and the rest of the array.
let (first, rest) = arr.view_mut().split_at(axis, 1);
// Remove the extra axis to simplify things a bit.
let mut first = first.remove_axis(axis);
// Zip the pieces together and perform the in-place sums, saving the results to `first`.
azip!(mut first, ref rest (rest.lanes(axis)) in {
*first += rest.scalar_sum()
});
}
// Select the first subview (where the results were stored).
arr.into_subview(axis, 0)
}
fn main() {
let arr = array![[0, 1, 2], [3, 4, 5]];
assert_eq!(sum_axis_inplace(arr, Axis(0)), array![3, 5, 7]);
let arr = array![[0, 1, 2], [3, 4, 5]];
assert_eq!(sum_axis_inplace(arr, Axis(1)), array![3, 12]);
let arr = array![[[0, 1], [2, 3]], [[4, 5], [6, 7]]];
assert_eq!(
sum_axis_inplace(arr.clone(), Axis(0)),
array![[4, 6], [8, 10]],
);
assert_eq!(
sum_axis_inplace(arr.clone(), Axis(1)),
array![[2, 4], [10, 12]],
);
assert_eq!(
sum_axis_inplace(arr.clone(), Axis(2)),
array![[1, 5], [9, 13]],
);
} This uses Does my comment answer your question? What would you suggest for "easier access to the elements of the array (Vec, shape, strides)"? With the various methods I talked about in the list at the beginning of the comment, I think it's possible to access everything. However, I do think that the Edit: "[Are] there any plans to [provide in-place axis operations]?" At the moment, I don't have any plans to do this myself, but maybe @bluss does. In most cases, I'd rather just create a new array with the correct number of elements (the behavior of the current |
Thanks @jturner314 for you quick and detailed response. I think your solution is pretty good. A few clarifications:
As you were curious about my implementation, I leave it here. Is only defined for dynamic arrays. It uses a bunch of unsafe code and is very, very ugly, trying to workaround the ownership rules that rustc enforces (bad idea, I know). This implementation does not work because it does not modify the shape and strides of the resulting array. However, it locates the results in the first positions of the Vec. Probably, the worst Rust code ever written: #[macro_use]
extern crate ndarray;
extern crate num_traits;
extern crate core;
use num_traits::Zero;
use std::ops::{Add, AddAssign};
use std::fmt::Debug;
use ndarray::{Ix, Ixs, ArrayView, ShapeBuilder, Axis, ArrayViewD, ArrayD};
/// Copied from ndarray's dimension/mod.rs
/// Calculate offset from `Ix` stride converting sign properly
#[inline(always)]
fn stride_offset(n: Ix, stride: Ixs) -> isize {
(n as isize) * (stride as isize)
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TabularValues<V>(pub ArrayD<V>);
impl<V> TabularValues<V>
where V: Debug + Zero + Clone + AddAssign<V>
{
/// Returns the shape and stride arrays with the axis removed.
fn shape_remove_axis(&self, axis: &Axis) -> (Vec<Ix>, Vec<Ix>) {
let shape = self.0.shape().clone();
let new_shape = shape.iter()
.enumerate()
// Removing the axis
.filter(|&(index, _)| index != axis.index())
.map(|(_, &dim)| dim)
.collect::<Vec<Ix>>();
let strides = self.0.strides().clone();
let new_strides = strides.iter()
.enumerate()
// Removing the axis
.filter(|&(index, _)| index != axis.index())
.map(|(_, &stride)| stride as usize)
.collect::<Vec<Ix>>();
(new_shape, new_strides)
}
/// Push ArrayViewDs of each index of the given axis to vec_subviews. With the new_shape and
/// new_strides, we can obtain views of the given axis. This is the same as filling a Vec with
/// the result of self.0.subview(axis, i); with i from 1 to ax_len.
/// This avoids rustc complaining again..
fn get_subviews(&self, axis: Axis, new_shape: Vec<Ix>, new_strides: Vec<Ix>, vec_subviews: &mut Vec<ArrayViewD<V>>) {
let ax_len = self.0.len_of(axis);
// We include all the subviews from 1 to ax_len because all these subviews will be summed to the
// subview with index 0.
for i in 1..ax_len {
let off = stride_offset(i, self.0.strides()[axis.index()]);
let mut ptr = self.0.as_ptr();
unsafe {
ptr = ptr.offset(off);
let subview = ArrayView::from_shape_ptr(new_shape.clone().strides(new_strides.clone()), ptr);
vec_subviews.push(subview);
}
}
}
pub fn sum_axis_inplace(&mut self, axis: Axis)
{
let ax_len = self.0.len_of(axis);
let (new_shape, new_strides) = self.shape_remove_axis(&axis);
// This vector should be initialized here because if it is created by the function get_subviews()
// there is a rustc error of inmutable/mutable reference to self.
let mut vec_subviews = Vec::with_capacity(ax_len - 1);
self.get_subviews(axis, new_shape.clone(), new_strides.clone(), &mut vec_subviews);
{
// Do the sum over a view in the axis with index 0.
let mut subview_mut = self.0.subview_mut(axis, 0);
vec_subviews.iter()
.for_each(|sub_view| {
subview_mut += sub_view
});
}
let mut ptr = self.0.as_ptr();
// Relocate the summed elements in the first positions of the vec.
unsafe {
// This is the same as self.0.subview(axis, 0);
// It avoids an inmutable/mutable reference error with self in the next line.
let subview = ArrayView::from_shape_ptr(new_shape.clone().strides(new_strides.clone()), ptr);
self.0.iter_mut().zip(subview.iter())
.for_each(|(mut_ref, subview_value)| {
*mut_ref = subview_value.clone()
});
}
// Modify shape and strides
// ....
}
}
fn main() {
let mut arr = TabularValues(array![[0, 1, 2], [3, 4, 5]].into_dyn());
arr.sum_axis_inplace(Axis(0));
assert_eq!(arr, TabularValues(array![3, 5, 7].into_dyn()));
let mut arr = TabularValues(array![[0, 1, 2], [3, 4, 5]].into_dyn());
arr.sum_axis_inplace(Axis(1));
assert_eq!(arr, TabularValues(array![3, 12].into_dyn()));
let mut arr = TabularValues(array![[[0, 1], [2, 3]], [[4, 5], [6, 7]]].into_dyn());
let mut arr_test = arr.clone();
arr_test.sum_axis_inplace(Axis(0));
assert_eq!(
arr_test,
TabularValues(array![[4, 6], [8, 10]].into_dyn()),
);
let mut arr_test = arr.clone();
arr_test.sum_axis_inplace((Axis(1)));
assert_eq!(
arr_test,
TabularValues(array![[2, 4], [10, 12]].into_dyn()),
);
let mut arr_test = arr.clone();
arr_test.sum_axis_inplace((Axis(2)));
assert_eq!(
arr_test,
TabularValues(array![[1, 5], [9, 13]].into_dyn()),
);
} As offtopic, I found myself also solving the "be able to sum over multiple axes". That was easier, but it would be cool to have it implemented. |
Welcome to Rust and
I understand now. I'm hesitant to provide direct mutable access to those things just because it's so easy to make a mistake, but I suppose we can make this For shape and strides, a safe approach would be to add a method For the
Oh, okay. Here's a (somewhat hacky but safe) solution using the currently available methods: #[macro_use]
extern crate ndarray;
extern crate num_traits;
use ndarray::prelude::*;
use num_traits::Zero;
use std::ops::{Add, AddAssign};
/// Sums along the axis without allocating a new array.
///
/// **Panics** if `axis` is out-of-bounds or if `arr.len_of(axis) == 0`.
fn sum_axis_inplace<A>(arr: &mut ArrayD<A>, axis: Axis)
where
A: Clone + Add<Output = A> + AddAssign + Zero,
{
{
// Get mutable views of the first subview along `axis` and the rest of the array.
let (first, rest) = arr.view_mut().split_at(axis, 1);
// Remove the extra axis to simplify things a bit.
let mut first = first.remove_axis(axis);
// Zip the pieces together and perform the in-place sums, saving the results to `first`.
azip!(mut first, ref rest (rest.lanes(axis)) in {
*first += rest.scalar_sum()
});
}
// Take ownership of the array by swapping a temporary into the `arr` reference.
// (This doesn't require any heap allocations.)
let owned = ::std::mem::replace(arr, Array::from_vec(Vec::new()).into_dyn());
// Select the first subview (where the results were stored) and put it into the `arr` reference.
::std::mem::replace(arr, owned.into_subview(axis, 0));
}
fn main() {
let mut arr = array![[0, 1, 2], [3, 4, 5]].into_dyn();
sum_axis_inplace(&mut arr, Axis(0));
assert_eq!(arr, array![3, 5, 7].into_dyn());
let mut arr = array![[0, 1, 2], [3, 4, 5]].into_dyn();
sum_axis_inplace(&mut arr, Axis(1));
assert_eq!(arr, array![3, 12].into_dyn());
let mut arr = array![[[0, 1], [2, 3]], [[4, 5], [6, 7]]].into_dyn();
sum_axis_inplace(&mut arr, Axis(0));
assert_eq!(arr, array![[4, 6], [8, 10]].into_dyn());
let mut arr = array![[[0, 1], [2, 3]], [[4, 5], [6, 7]]].into_dyn();
sum_axis_inplace(&mut arr, Axis(1));
assert_eq!(arr, array![[2, 4], [10, 12]].into_dyn());
let mut arr = array![[[0, 1], [2, 3]], [[4, 5], [6, 7]]].into_dyn();
sum_axis_inplace(&mut arr, Axis(2));
assert_eq!(arr, array![[1, 5], [9, 13]].into_dyn());
} Note that one option is to not remove the extra axis and instead just set the length of that axis to 1 by using That said, I do think it would be good to add in-place variants of
I realized that this problem is actually a lot easier than I initially thought. In fact, I think it would be useful to add a If all of the strides are positive and you keep the order of the elements in memory the same, you can just shift the used elements to the front of the
with column-major memory layout, so it looks like this in memory:
Note that the strides are 1, 2, 4. If you want to sum over axis 1, you can store the results in the first subview along axis 1, so after performing the sums you'd have
which is
in memory. Now, you just need to shift the sums to the front of the
Then, you can drop the last four elements and you're done. The resulting array is column-major layout with strides 1, 2. This appears to be the approach that the "Relocate the summed elements in the first positions of the vec." section in your implementation is generally trying to follow. You do have to be careful with a couple of things, though:
How do you determine strides for the shrunken array that maintain the memory order of the original array, particularly when the original array can have weird custom strides? If all of the strides are positive, I think this will work:
For example, let's say the shape was 2, 3, 4, the strides were 1, 2, 6, and you're removing axis 1. The results of the steps would be:
I don't think this will work unchanged for negative strides; they require some more thought. However, I think it will work for weird custom strides as long as they're all positive. Since you're new to Rust and
Anyway, in summary to all this:
Since we've already had a lot of discussion in this thread, I'd like to create separate issues for the API changes. Does the code I provided resolve your use-case for this issue? What do you think of the API changes I listed? |
Thank you very much for you kindly support and comments. I implemented your version without further problems (nice trick using mem::replace() to take ownership). I don't shrink the vector to avoid dealing with unsafe code and pointer arithmetics. It would be nice to have it implemented in ndarray directly, so you don't have to deal with the movement of the results at the start of the vector and the computation of shape/strides. If it is implemented in the future, I certainly will use it. With the rest of the comments, your suggestions seems nice to me.
|
Okay, it sounds like this issue is resolved, so I'm closing it. I've created new issues for the potential API changes. |
The issue is closed, but I thought it would be interesting to share some benchmark results comparing the The results:
We can see an important performance improvement, especially in large arrays. I think that in those instances, the heap allocations dominate the computation, so the larger the array, the larger the improvements. |
(Hi!) I just want to chime in and say that we don't want to expose something that locks us in to using a Vec for holding the array elements. That we can convert to/from a Vec is ok, but not that we show a Vec behind a reference. |
It would be really useful to have inplace versions for some operations.
In one of my projects I am using sum_axis(). However, I would prefer to make the operation inplace. I don't think is necessary to allocate more memory. Is there any plans to do such a thing?
I've tried to develop my own version of the sum_axis_inplace by wrapping the ArrayD inside a struct:
pub struct TabularValues<V>(pub ArrayD<V>);
However, making inplace changes in the Array involves a bunch of unsafe code (in my implementation I used multiple ArrayViews of the data in the deleted axis with ArrayView::from_shape_ptr()). Still, my implementation is not complete because I don't know how to modify the shape and strides. In fact, I would also like to call shrink_to_fit in the inner Vec, but I don't know how to do that.
Having easier access to the elements of the array (Vec, shape, strides) would be useful.
The text was updated successfully, but these errors were encountered: