dlpark

A pure Rust implementation of dmlc/dlpack.

Check example/from_numpy for usage.

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

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 HasData for PyRgbImage { fn data(&self) -> *mut std::ffi::cvoid { self.0.asptr() as *const cvoid as *mut cvoid } }

impl HasDevice for PyRgbImage { fn device(&self) -> Device { Device::CPU } }

impl HasDtype for PyRgbImage { fn dtype(&self) -> DataType { DataType::U8 } }

impl HasShape for PyRgbImage { fn shape(&self) -> Shape { Shape::Owned( [self.0.height(), self.0.width(), 3] .map(|x| x as i64) .to_vec(), ) } }

// Strides can be infered from Shape since it's compact and row-majored. impl HasStrides for PyRgbImage {} impl HasByteOffset for PyRgbImage {} ```

Then we can return a ManagerCtx<PyRgbImage>

```rust

[pyfunction]

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

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)) ```