dlpark 0.7.0

dlpack Rust binding for Python
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

#[derive(Debug, Eq, PartialEq)]
pub enum MemoryOrder {
    RowMajorContiguous,
    ColumnMajorContiguous,
    NonContiguous,
}

pub fn make_row_major_strides(shape: &[i64]) -> Vec<i64> {
    let rank = shape.len();
    let mut strides = vec![1; rank];
    for i in (0..rank - 1).rev() {
        strides[i] = strides[i + 1] * shape[i + 1];
    }
    strides
}

pub fn make_column_major_strides(shape: &[i64]) -> Vec<i64> {
    let rank = shape.len();
    let mut strides = vec![1; rank];
    for i in 0..rank - 1 {
        strides[i + 1] = strides[i] * shape[i];
    }
    strides
}

impl std::fmt::Display for MemoryOrder {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MemoryOrder::RowMajorContiguous => write!(f, "RowMajorContiguous"),
            MemoryOrder::ColumnMajorContiguous => write!(f, "ColumnMajorContiguous"),
            MemoryOrder::NonContiguous => write!(f, "NonContiguous"),
        }
    }
}

impl MemoryOrder {
    pub fn new(shape: &[i64], strides: &[i64]) -> Self {
        if is_row_major_contiguous(shape, strides) {
            MemoryOrder::RowMajorContiguous
        } else if is_column_major_contiguous(shape, strides) {
            MemoryOrder::ColumnMajorContiguous
        } else {
            MemoryOrder::NonContiguous
        }
    }

    pub fn is_contiguous(&self) -> bool {
        match self {
            MemoryOrder::RowMajorContiguous | MemoryOrder::ColumnMajorContiguous => true,
            MemoryOrder::NonContiguous => false,
        }
    }
}

pub fn is_row_major_contiguous(shape: &[i64], strides: &[i64]) -> bool {
    let mut expected = 1;
    for (dim_size, stride) in shape.iter().zip(strides.iter()).rev() {
        if *dim_size == 0 {
            continue;
        }
        if *stride != expected {
            return false;
        }
        expected *= dim_size;
    }
    true
}

pub fn is_column_major_contiguous(shape: &[i64], strides: &[i64]) -> bool {
    let mut expected = 1;
    for (dim_size, stride) in shape.iter().zip(strides.iter()) {
        if *dim_size == 0 {
            continue;
        }
        if *stride != expected {
            return false;
        }
        expected *= dim_size;
    }
    true
}

// Generated by copilot.
#[cfg(test)]
mod tests {
    use super::*;

    // test make_contiguous_strides
    #[test]
    fn test_row_major_strides() {
        let shape = vec![1, 2, 3];
        let strides = make_row_major_strides(&shape);
        assert_eq!(strides, vec![6, 3, 1]);
        assert!(is_row_major_contiguous(&shape, &strides));
    }

    // test make_column_major_strides
    #[test]
    fn test_column_major_strides() {
        let shape = vec![1, 2, 3];
        let strides = make_column_major_strides(&shape);
        assert_eq!(strides, vec![1, 1, 2]);
        assert!(is_column_major_contiguous(&shape, &strides));
    }
}