ndarray icon indicating copy to clipboard operation
ndarray copied to clipboard

Exact Chunks Parallel Iterator

Open benjaminrwilson opened this issue 2 years ago • 1 comments

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?

benjaminrwilson avatar Jul 31 '22 15:07 benjaminrwilson

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.

jturner314 avatar Aug 01 '22 00:08 jturner314