I'm having difficulty performing an operation using ndarray - Rust (probably because I am very familiar with numpy and expecting things to just work) and could use some help.
Briefly, I have 2D array for which I would like to multiply all the elements in the bottom row by the same scalar value, and then add assign the results to the top row.
In numpy, this looks like:
arr[0, :] += x * arr[1, :]
# or, equivalently
arr[0] += x * arr[1]
Edited: Fixed indexing in Python example
I've tried a few things in Rust, but I keep running into either the borrow checker being unhappy about me mutably borrowing the top row while trying to reference the bottom row, or there are issues with the add/assignment. Here is one attempt which fails to compile due to the mutable and immutable borrows.
use ndarray::{arr2, s};
fn main(){
let mut arr = arr2(&[[1, 2],
[3, 4]]);
let x = 2;
// Mutable and immutable slices occur here
let mut top = arr.slice_mut(s![0, ..]);
let bottom = arr.slice(s![1, ..]);
top += &(&bottom * x);
let expected = arr2(&[[7, 10],
[3, 4]]);
assert_eq!(arr, expected);
}
Any ideas on how to do this without unnecessary copying? Copilot seems stumped. Thanks!
The easiest way to do this would be if ndarray had an equivalent of slice::get_many_mut. Unfortunately it doesn't.
What you can do instead is use a mutable iterator to enable having multiple mutable slices of disjoint parts of the array at the same time:
(Edit: Modified to update top without allocating a temporary containing bottom * x)
use ndarray::{arr2, Axis,};
fn main(){
let mut arr = arr2(&[[1, 2],
[3, 4]]);
let x = 2;
let mut rows = arr.axis_iter_mut(Axis(0));
let mut top = rows.next().unwrap();
let bottom = rows.next().unwrap();
top.zip_mut_with(&bottom, |t, b| *t += b * x);
println!("arr {:?}", arr);
let expected = arr2(&[[7, 10],
[3, 4]]);
assert_eq!(arr, expected);
}
In general, any time you work with a collection of some kind and need to have mutable references to logically disjoint parts of it, that will require using a collection method that internally "promises" the compiler that the views are disjoint.
If performance is not a proven issue, copying the slice that you are going to multiply could be a simpler/easier solution.
My solution doesn't copy/allocate any more than anything else in this thread. I just moved the operations around to make the immutable borrow no longer than it needs to be.
I don't see where you found additional copying in my code, compared to your original code. My code does the exact same thing as yours but pulls out the &bottom * x temporary into a named variable.
I've corrected my code sample from yesterday. I meant to show that you can update top without creating a temporary, but missed that it was still computing bottom * x which allocates a new owned tensor. The updated version uses zip_mut_with to avoid creating the temporary.
Thanks for the correction. I'm starting to see now that doing this without any additional memory allocations is actually not straight-forward when coming from a numpy background.
It does allocate a new array, but not because it's a variable instead of a temporary value, but because the expression &bottom * x itself allocates a new array. It has to – there's no way to return a dynamically-sized value otherwise.