pub struct PyTensor { /* private fields */ }Expand description
A tensor imported from Python via the DLPack protocol.
This type wraps a DLManagedTensor received from a Python object
(typically a PyTorch, JAX, or NumPy tensor) and provides safe access
to the tensor’s metadata and data pointer.
§Lifetime
The tensor data is valid as long as this PyTensor is alive.
When dropped, the tensor’s deleter is called to notify the producer.
§Example
use pyo3::prelude::*;
use pyo3_dlpack::PyTensor;
#[pyfunction]
fn process(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<()> {
let tensor = PyTensor::from_pyany(py, obj)?;
println!("Shape: {:?}", tensor.shape());
println!("Device: {:?}", tensor.device());
println!("Dtype: {:?}", tensor.dtype());
if tensor.device().is_cpu() {
// Safe to access data on CPU
let ptr = tensor.data_ptr() as *const f32;
// ...
}
Ok(())
}Implementations§
Source§impl PyTensor
impl PyTensor
Sourcepub fn from_pyany(_py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<Self>
pub fn from_pyany(_py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<Self>
Create a PyTensor from a Python object that supports the DLPack protocol.
This calls __dlpack__() on the object to get a DLPack capsule,
then extracts the tensor information.
§Arguments
py- Python GIL tokenobj- A Python object that implements__dlpack__()(e.g., PyTorch tensor)
§Errors
Returns an error if:
- The object doesn’t have a
__dlpack__method - The returned capsule is invalid
- The capsule doesn’t contain a valid DLManagedTensor
Sourcepub fn from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyResult<Self>
pub fn from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyResult<Self>
Sourcepub fn dtype(&self) -> DLDataType
pub fn dtype(&self) -> DLDataType
Get the data type of the tensor elements.
Sourcepub fn strides(&self) -> Option<&[i64]>
pub fn strides(&self) -> Option<&[i64]>
Get the strides as a slice, or None for contiguous tensors.
Strides are in number of elements (not bytes).
If None, the tensor is assumed to be in compact row-major order.
Sourcepub fn is_contiguous(&self) -> bool
pub fn is_contiguous(&self) -> bool
Check if the tensor is contiguous in row-major (C) order.
Sourcepub fn data_ptr(&self) -> *mut c_void
pub fn data_ptr(&self) -> *mut c_void
Get the raw data pointer.
For GPU tensors, this is a device pointer that cannot be directly dereferenced on the CPU.
The pointer is adjusted by byte_offset().
Sourcepub fn data_ptr_raw(&self) -> *mut c_void
pub fn data_ptr_raw(&self) -> *mut c_void
Get the raw data pointer without byte offset adjustment.
Sourcepub fn byte_offset(&self) -> u64
pub fn byte_offset(&self) -> u64
Get the byte offset from the raw data pointer.