pyo3_dlpack/lib.rs
1//! # pyo3-dlpack
2//!
3//! Zero-copy DLPack tensor interop for PyO3.
4//!
5//! This crate provides a safe and ergonomic way to exchange tensor data between
6//! Rust and Python ML frameworks (PyTorch, JAX, TensorFlow, CuPy, etc.) using
7//! the [DLPack](https://github.com/dmlc/dlpack) protocol.
8//!
9//! ## Features
10//!
11//! - **Zero-copy**: Tensors are shared directly without copying data
12//! - **PyO3 0.23+**: Uses the modern `IntoPyObject` trait (no deprecation warnings)
13//! - **Bidirectional**: Import tensors from Python and export tensors to Python
14//! - **Device-agnostic**: Works with CPU, CUDA, ROCm, and other devices
15//!
16//! ## Example: Importing a PyTorch tensor
17//!
18//! ```ignore
19//! use pyo3::prelude::*;
20//! use pyo3_dlpack::PyTensor;
21//!
22//! #[pyfunction]
23//! fn process_tensor(tensor: PyTensor) -> PyResult<()> {
24//! // Access tensor metadata
25//! println!("Shape: {:?}", tensor.shape());
26//! println!("Device: {:?}", tensor.device());
27//!
28//! // Get the raw data pointer (for GPU tensors, this is a device pointer)
29//! let ptr = tensor.data_ptr();
30//!
31//! Ok(())
32//! }
33//! ```
34//!
35//! ## Example: Exporting a tensor to Python
36//!
37//! ```ignore
38//! use pyo3::prelude::*;
39//! use pyo3_dlpack::{ExportConfig, IntoDLPack};
40//!
41//! struct MyGpuTensor {
42//! ptr: *mut f32,
43//! shape: Vec<i64>,
44//! device_id: i32,
45//! }
46//!
47//! impl IntoDLPack for MyGpuTensor {
48//! // ... implement the trait
49//! }
50//!
51//! #[pyfunction]
52//! fn create_tensor(py: Python<'_>) -> PyResult<PyObject> {
53//! let tensor = MyGpuTensor { /* ... */ };
54//! tensor.into_dlpack(py)
55//! }
56//! ```
57
58mod export;
59mod ffi;
60mod managed;
61
62// Re-export public API
63pub use export::{IntoDLPack, TensorInfo};
64pub use ffi::{DLDataType, DLDataTypeCode, DLDevice, DLDeviceType, DLManagedTensor, DLTensor};
65pub use managed::PyTensor;
66
67// Convenience constructors
68pub use ffi::{
69 cpu_device, cuda_device, dtype_bf16, dtype_bool, dtype_f16, dtype_f32, dtype_f64, dtype_i16,
70 dtype_i32, dtype_i64, dtype_i8, dtype_u16, dtype_u32, dtype_u64, dtype_u8, metal_device,
71};
72
73/// The DLPack capsule name for tensor exchange
74pub const DLPACK_CAPSULE_NAME: &std::ffi::CStr = c"dltensor";
75
76/// The DLPack capsule name after consumption (to prevent double-free)
77pub const DLPACK_CAPSULE_NAME_USED: &std::ffi::CStr = c"used_dltensor";