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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

// use crate::error::Error;

// use crate::safe_managed_tensor::SafeManagedTensorVersioned;
// use crate::{
//     data_type::{DataType, InferDataType},
//     device::Device,
//     manager_context::TensorLike,
//     memory_layout::StridedLayout,
//     utils::make_contiguous_strides,
// };

use crate::traits::{InferDataType, StridedLayout, TensorLike, TensorView};
use crate::utils::make_row_major_strides;
use crate::{Result, ffi};
use crate::{SafeManagedTensor, SafeManagedTensorVersioned};

use ndarray::{ArrayBase, ArrayViewD, Dimension, RawData, ShapeBuilder};

impl<S, D> TensorLike<StridedLayout> for ArrayBase<S, D>
where
    S: RawData,
    S::Elem: InferDataType,
    D: Dimension,
{
    type Error = crate::Error;

    fn data_ptr(&self) -> *mut std::ffi::c_void {
        self.as_ptr() as *mut S::Elem as *mut std::ffi::c_void
    }

    fn memory_layout(&self) -> StridedLayout {
        let mut layout = StridedLayout::with_ndim(self.ndim());
        for i in 0..self.ndim() {
            layout.shape_mut()[i] = self.shape()[i] as i64;
            layout.strides_mut()[i] = self.strides()[i] as i64;
        }
        layout
    }

    fn device(&self) -> Result<ffi::Device> {
        Ok(ffi::Device::CPU)
    }

    fn data_type(&self) -> Result<ffi::DataType> {
        Ok(S::Elem::data_type())
    }

    fn byte_offset(&self) -> u64 {
        0
    }
}

impl<'a, A> TryFrom<&'a SafeManagedTensorVersioned> for ArrayViewD<'a, A> {
    type Error = crate::Error;

    fn try_from(value: &'a SafeManagedTensorVersioned) -> Result<Self> {
        let shape: Vec<usize> = value.shape().iter().map(|x| *x as usize).collect();
        let shape = match value.strides() {
            Some(s) => {
                let strides: Vec<usize> = s.iter().map(|x| *x as usize).collect();
                shape.strides(strides)
            }
            None => {
                let strides = make_row_major_strides(value.shape())
                    .into_iter()
                    .map(|x| x as usize)
                    .collect();
                shape.strides(strides)
            }
        };
        unsafe {
            Ok(ArrayViewD::from_shape_ptr(
                shape,
                value.as_slice::<A>()?.as_ptr(),
            ))
        }
    }
}

impl<'a, A> TryFrom<&'a SafeManagedTensor> for ArrayViewD<'a, A> {
    type Error = crate::Error;

    fn try_from(value: &'a SafeManagedTensor) -> Result<Self> {
        let shape: Vec<usize> = value.shape().iter().map(|x| *x as usize).collect();
        let shape = match value.strides() {
            Some(s) => {
                let strides: Vec<usize> = s.iter().map(|x| *x as usize).collect();
                shape.strides(strides)
            }
            None => {
                let strides = make_row_major_strides(value.shape())
                    .into_iter()
                    .map(|x| x as usize)
                    .collect();
                shape.strides(strides)
            }
        };
        unsafe {
            Ok(ArrayViewD::from_shape_ptr(
                shape,
                value.as_slice::<A>()?.as_ptr(),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use ndarray::IxDyn;

    use super::*;

    #[test]
    fn test_dlpack() {
        let arr =
            ndarray::ArrayD::from_shape_vec(IxDyn(&[1, 2, 3]), vec![1i32, 2, 3, 4, 5, 6]).unwrap();
        let arr2 = arr.clone();
        let mt = SafeManagedTensor::new(arr).unwrap();
        let view = ArrayViewD::<i32>::try_from(&mt).unwrap();
        assert_eq!(view.shape(), arr2.shape());
        assert_eq!(view.strides(), arr2.strides());
        assert_eq!(view.as_slice().unwrap(), arr2.as_slice().unwrap());
    }

    #[test]
    fn test_dlpack_versioned() {
        let arr =
            ndarray::ArrayD::from_shape_vec(IxDyn(&[1, 2, 3]), vec![1i32, 2, 3, 4, 5, 6]).unwrap();
        let arr2 = arr.clone();
        let mt = SafeManagedTensorVersioned::new(arr).unwrap();
        let view = ArrayViewD::<i32>::try_from(&mt).unwrap();
        assert_eq!(view.shape(), arr2.shape());
        assert_eq!(view.strides(), arr2.strides());
        assert_eq!(view.as_slice().unwrap(), arr2.as_slice().unwrap());
    }
}