Crate pyo3_dlpack 

pyo3-dlpack

Zero-copy DLPack tensor interop for PyO3.

This crate provides a safe and ergonomic way to exchange tensor data between Rust and Python ML frameworks (PyTorch, JAX, TensorFlow, CuPy, etc.) using the DLPack protocol.

Features

• Zero-copy: Tensors are shared directly without copying data
• PyO3 0.23+: Uses the modern IntoPyObject trait (no deprecation warnings)
• Bidirectional: Import tensors from Python and export tensors to Python
• Device-agnostic: Works with CPU, CUDA, ROCm, and other devices

Example: Importing a PyTorch tensor

use pyo3::prelude::*;
use pyo3_dlpack::PyTensor;

#[pyfunction]
fn process_tensor(tensor: PyTensor) -> PyResult<()> {
    // Access tensor metadata
    println!("Shape: {:?}", tensor.shape());
    println!("Device: {:?}", tensor.device());

    // Get the raw data pointer (for GPU tensors, this is a device pointer)
    let ptr = tensor.data_ptr();

    Ok(())
}

Example: Exporting a tensor to Python

use pyo3::prelude::*;
use pyo3_dlpack::{ExportConfig, IntoDLPack};

struct MyGpuTensor {
    ptr: *mut f32,
    shape: Vec<i64>,
    device_id: i32,
}

impl IntoDLPack for MyGpuTensor {
    // ... implement the trait
}

#[pyfunction]
fn create_tensor(py: Python<'_>) -> PyResult<PyObject> {
    let tensor = MyGpuTensor { /* ... */ };
    tensor.into_dlpack(py)
}

Structs

DLDataType

Data type descriptor specifying the element type of a tensor.

DLDevice

A device descriptor specifying where tensor data resides.

DLManagedTensor

A managed tensor with ownership semantics.

DLTensor

The core DLTensor structure describing a tensor’s data and layout.

PyTensor

A tensor imported from Python via the DLPack protocol.

TensorInfo

Information about a tensor for DLPack export.

Enums

DLDataTypeCode

Data type codes as defined by DLPack.

DLDeviceType

Device type codes as defined by DLPack.

Constants

DLPACK_CAPSULE_NAME

The DLPack capsule name for tensor exchange

DLPACK_CAPSULE_NAME_USED

The DLPack capsule name after consumption (to prevent double-free)

Traits

IntoDLPack

Trait for types that can be exported as DLPack tensors.

Functions

cpu_device

Create a DLDevice for CPU.

cuda_device

Create a DLDevice for CUDA with the specified device ID.

dtype_bf16

Create a DLDataType for bf16 (bfloat16).

dtype_bool

Create a DLDataType for bool.

dtype_f16

Create a DLDataType for f16 (half precision float).

dtype_f32

Create a DLDataType for f32 (single precision float).

dtype_f64

Create a DLDataType for f64 (double precision float).

dtype_i8

Create a DLDataType for i8 (signed 8-bit integer).

dtype_i16

Create a DLDataType for i16 (signed 16-bit integer).

dtype_i32

Create a DLDataType for i32 (signed 32-bit integer).

dtype_i64

Create a DLDataType for i64 (signed 64-bit integer).

dtype_u8

Create a DLDataType for u8 (unsigned 8-bit integer).

dtype_u16

Create a DLDataType for u16 (unsigned 16-bit integer).

dtype_u32

Create a DLDataType for u32 (unsigned 32-bit integer).

dtype_u64

Create a DLDataType for u64 (unsigned 64-bit integer).

metal_device

Create a DLDevice for Metal (Apple GPU) with the specified device ID.