Announcement 📢 Releasing dlpackrs

DLPack is the standard in-memory data format that facilitates zero-cost tensor transfer across major Deep Learning frameworks (PyTorch, TensorFlow and TVM) and the supported Python Array processing frameworks such as Numpy, CuPy. The dlpackrs provides a safe idiomatic Rust binding where Rust ndarray and tensor frameworks can use it to gain the same kind of benefits (if not done already) across the supported hardware types.

In Python, such conversion goes through PyCapsule (minimal example for TVM – numpy conversion is here) but Rust needs no such restrictions or workaround. In fact, the PyO3/numpy conversion to Rust ndarray is already zero-cost through pointers (taken from here) which is

impl<S, D, A> ToPyArray for ArrayBase<S, D>
where
    S: Data<Elem = A>,
    D: Dimension,
    A: Element,
{
    type Item = A;
    type Dim = D;

    fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
        let len = self.len();
        match self.order() {
            Some(flag) if A::IS_COPY => {
                // if the array is contiguous, copy it by `copy_nonoverlapping`.
                let strides = self.npy_strides();
                unsafe {
                    let array = PyArray::new_uninit(py, self.raw_dim(), strides.as_ptr(), flag);
                    ptr::copy_nonoverlapping(self.as_ptr(), array.data(), len); // <-- pointer copy
                    array
                }
            }
            _ => {
                // if the array is not contiguous, copy all elements by `ArrayBase::iter`.
                let dim = self.raw_dim();
                unsafe {
                    let array = PyArray::<A, _>::new(py, dim, false);
                    let mut data_ptr = array.data();
                    for item in self.iter() { // <-- pointer copy starts
                        data_ptr.write(item.clone());
                        data_ptr = data_ptr.add(1);
                    }
                    array
                }
            }
        }
    }
}



Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.

About Me

I’m a Machine Learning Engineer and have been inhaling in the Data space since 2013. Before that, I was devoted to the pure Mathematics. I like learning and working in public. I blog to manage and structure information. If you want to know more about me and my work, please check out my GitHub and Linkedin.

Newsletter

%d bloggers like this: