dlpark 0.9.0-alpha.1

dlpack Rust binding for Python
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

use crate::{
    DlpackElement, ManagedTensorBase,
    allocation::dynamic,
    ffi::DLDevice,
    metadata::{Copied, Dynamic},
};
use ndarray::{ArrayBase, Dimension, OwnedRepr};
use std::os::raw::c_void;

/// Converts a boxed owned ndarray into an initialized DLPack allocation.
///
/// The array is not copied. Its shape and strides are converted to DLPack's
/// `i64` representation before the boxed array becomes the manager context.
/// Moving the array transfers ownership without copying its data allocation.
impl<T, D, M> TryFrom<Box<ArrayBase<OwnedRepr<T>, D>>> for dynamic::Initialized<M>
where
    T: DlpackElement + Send,
    D: Dimension,
    M: ManagedTensorBase,
{
    type Error = crate::metadata::Error;

    fn try_from(array: Box<ArrayBase<OwnedRepr<T>, D>>) -> Result<Self, Self::Error> {
        let data_ptr = if array.is_empty() {
            std::ptr::null_mut()
        } else {
            array.as_ptr() as *mut c_void
        };
        let prepared =
            Dynamic::new(Copied(array.shape()), Copied(array.strides())).prepare::<M>()?;
        let mut initialized = prepared.initialize(array)?;
        initialized.set_data(data_ptr);
        initialized.set_dtype(T::DTYPE);
        initialized.set_device(DLDevice::CPU);
        Ok(initialized)
    }
}