ndarray
ndarray copied to clipboard
Exact Chunks Parallel Iterator
Hello,
Thank you for this fantastic package.
I would like to split an MxN
array into k^2
chunks of size (M/k)x(N/k)
and then process these in parallel. I'm aware of exact_chunks
, but it doesn't natively support par_iter
, right? What is the idiomatic way to do this?
We should add a ParallelIterator
implementation for Parallel<ExactChunks/Mut>
.
In the meantime, a workaround is to Zip
's parallel methods:
use ndarray::prelude::*;
use ndarray::Zip;
fn main() {
let mut a = Array2::<u8>::zeros([6, 6]);
Zip::from(a.exact_chunks_mut([3, 3])).par_for_each(|mut chunk| chunk.diag_mut().fill(1));
assert_eq!(
a,
array![
[1, 0, 0, 1, 0, 0],
[0, 1, 0, 0, 1, 0],
[0, 0, 1, 0, 0, 1],
[1, 0, 0, 1, 0, 0],
[0, 1, 0, 0, 1, 0],
[0, 0, 1, 0, 0, 1],
],
);
}
The par_azip!
macro may be a bit more convenient.
Regardless, make sure to enable to rayon
feature.