Module tvm::ndarray

This module implements the NDArray type for working with TVM tensors or coverting from a Rust’s ndarray to TVM NDArray.

One can create an empty NDArray given the shape, device and dtype using [empty]. To create an NDArray from a mutable buffer in cpu use copy_from_buffer. To copy an NDArray to different device use copy_to_device.

Given a Rust's dynamic ndarray, one can convert it to TVM NDArray as follows:

Example

use std::convert::TryFrom;

let a = Array::from_shape_vec((2, 2), vec![1f32, 2., 3., 4.])
    .unwrap()
    .into_dyn(); // Rust's ndarray
let nd = NDArray::from_rust_ndarray(&a, Device::cpu(0), DataType::from_str("float32").unwrap()).unwrap();
assert_eq!(nd.shape(), &[2, 2]);
let rnd: ArrayD<f32> = ArrayD::try_from(&nd).unwrap();
assert!(rnd.all_close(&a, 1e-8f32));

Structs

NDArray

NDArrayContainer

See the module-level documentation for more details.

Traits

Num32

A trait for the supported 32-bits numerical types in frontend.