dlpark

A pure Rust implementation of dmlc/dlpack.

Check example/with_pyo3 for usage.

This implementation focuses on transferring tensor from Rust to Python and vice versa.

It can also be used without pyo3 as a Rust library with default-features = false, check example/from_numpy.

Quick Start

We provide a simple example of how to transfer image::RgbImage to Python and torch.Tensor to Rust.

Full code is here.

Rust $\rightarrow$ Python

We have to implement some traits for a struct to be able to converted to PyObject

```rust use std::ffi::c_void; use dlpark::prelude::*;

struct PyRgbImage(image::RgbImage);

impl ToTensor for PyRgbImage { fn dataptr(&self) -> *mut std::ffi::cvoid { self.0.asptr() as *const cvoid as *mut c_void }

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

fn shape(&self) -> CowIntArray {
    CowIntArray::from_owned(
        [self.0.height(), self.0.width(), 3]
            .map(|x| x as i64)
            .to_vec(),
    )
}

fn device(&self) -> Device {
    Device::CPU
}

fn dtype(&self) -> DataType {
    DataType::U8
}

fn strides(&self) -> Option<CowIntArray> {
    None
}

} ```

Then we can return a ManagerCtx<PyRgbImage>

```rust

[pyfunction]

fn readimage(filename: &str) -> ManagerCtx { let img = image::open(filename).unwrap(); let rgbimg = img.torgb8(); ManagerCtx::new(PyRgbImage(rgbimg)) } ```

You can acess it in Python

```python import dlparkimg from torch.utils.dlpack import todlpack, fromdlpack import matplotlib.pyplot as plt

tensor = fromdlpack(dlparkimg.readimage("candy.jpg"))

print(tensor.shape) plt.imshow(tensor.numpy()) plt.show() ```

If you want to convert it to numpy.ndarray , you can make a simple wrapper like this

```python import numpy as np import dlparkimg

class FakeTensor:

def __init__(self, x):
    self.x = x

def __dlpack__(self):
    return self.x

arr = np.fromdlpack(FakeTensor(dlparkimg.readimage("candy.jpg"))) ```

Python $\rightarrow$ Rust

ManagedTensor holds the memory of tensor and provide methods to access the tensor's attributes.

```rust

[pyfunction]

fn writeimage(filename: &str, tensor: ManagedTensor) { let buf = tensor.asslice::();

let rgb_img = image::ImageBuffer::<Rgb<u8>, _>::from_raw(
    tensor.shape()[1] as u32,
    tensor.shape()[0] as u32,
    buf,
)
.unwrap();

rgb_img.save(filename).unwrap();

} ```

And you can call it in Python

```python import dlparkimg from torch.utils.dlpack import todlpack, fromdlpack

bgrimg = tensor[..., [2, 1, 0]] # [H, W, C=3] dlparkimg.writeimage('bgr.jpg', todlpack(bgrimg)) ```