mdarray 0.2.0

Multidimensional array for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

use crate::order::Order;
use std::marker::PhantomData;

pub trait Layout<const N: usize, O: Order> {
    fn shape(&self) -> &[usize; N];
}

pub struct StridedLayout<const N: usize, const M: usize, O: Order> {
    shape: [usize; N],
    strides: [isize; M],
    _marker: PhantomData<O>,
}

pub type DenseLayout<const N: usize, O> = StridedLayout<N, 0, O>;

impl<const N: usize, const M: usize, O: Order> StridedLayout<N, M, O> {
    pub const fn new(shape: [usize; N], strides: [isize; M]) -> Self {
        Self {
            shape,
            strides,
            _marker: PhantomData,
        }
    }

    pub fn stride(&self, dim: usize) -> isize {
        O::select(
            if dim < N - M {
                self.shape[..dim].iter().product::<usize>() as isize
            } else {
                self.strides[dim - (N - M)]
            },
            if dim < M {
                self.strides[dim]
            } else {
                self.shape[dim + 1..].iter().product::<usize>() as isize
            },
        )
    }

    pub fn strides(&self) -> &[isize; M] {
        &self.strides
    }
}

impl<const N: usize, const M: usize, O: Order> StridedLayout<N, M, O> {
    pub fn resize(&mut self, shape: [usize; N], strides: [isize; M]) {
        self.shape = shape;
        self.strides = strides;
    }
}

impl<const N: usize, const M: usize, O: Order> Layout<N, O> for StridedLayout<N, M, O> {
    fn shape(&self) -> &[usize; N] {
        &self.shape
    }
}