Class Tensor#

Inheritance Relationships#

Base Type#

Class Documentation#

class Tensor : public tvm::ffi::ObjectRef#

Managed Tensor (n-dimensional array). The tensor is backed by reference counted blocks.

Note

This class can be subclassed to implement downstream customized Tensor types that are backed by the same TensorObj storage type.

Public Functions

inline ShapeView shape() const#

Get the shape of the Tensor.

Returns:

The shape of the Tensor.

inline ShapeView strides() const#

Get the strides of the Tensor.

Returns:

The strides of the Tensor.

inline void *data_ptr() const#

Get the data pointer of the Tensor.

Returns:

The data pointer of the Tensor.

inline int32_t ndim() const#

Get the number of dimensions in the Tensor.

Returns:

The number of dimensions in the Tensor.

inline int64_t numel() const#

Get the number of elements in the Tensor.

Returns:

The number of elements in the Tensor.

inline DLDataType dtype() const#

Get the data type of the Tensor.

Returns:

The data type of the Tensor.

inline bool IsContiguous() const#

Check if the Tensor is contiguous.

Returns:

True if the Tensor is contiguous, false otherwise.

inline bool IsAligned(size_t alignment) const#

Check if the Tensor data is aligned to the given alignment.

Parameters:

alignment – The alignment to check.

Returns:

True if the Tensor data is aligned to the given alignment, false otherwise.

inline DLManagedTensor *ToDLPack() const#

Convert the Tensor to a DLPack managed tensor.

Returns:

The converted DLPack managed tensor.

inline DLManagedTensorVersioned *ToDLPackVersioned() const#

Convert the Tensor to a DLPack managed tensor.

Returns:

The converted DLPack managed tensor.

Public Static Functions

template<typename TNDAlloc, typename ...ExtraArgs>
static inline Tensor FromNDAlloc(TNDAlloc alloc, ffi::ShapeView shape, DLDataType dtype, DLDevice device, ExtraArgs&&... extra_args)#

Create a Tensor from a NDAllocator.

Parameters:
  • alloc – The NDAllocator.

  • shape – The shape of the Tensor.

  • dtype – The data type of the Tensor.

  • device – The device of the Tensor.

  • extra_args – Extra arguments to be forwarded to TNDAlloc.

Template Parameters:
  • TNDAlloc – The type of the NDAllocator, impelments Alloc and Free.

  • ExtraArgs – Extra arguments to be passed to Alloc.

Returns:

The created Tensor.

static inline Tensor FromDLPackAlloc(DLPackTensorAllocator allocator, ffi::Shape shape, DLDataType dtype, DLDevice device)#

Create a Tensor from a DLPackTensorAllocator.

This function can be used together with TVMFFIEnvSetTensorAllocator in the extra/c_env_api.h to create Tensor from the thread-local environment allocator.

ffi::Tensor tensor = ffi::Tensor::FromDLPackAlloc(
  TVMFFIEnvGetTensorAllocator(), shape, dtype, device
);
Parameters:
  • allocator – The DLPack allocator.

  • shape – The shape of the Tensor.

  • dtype – The data type of the Tensor.

  • device – The device of the Tensor.

Returns:

The created Tensor.

static inline Tensor FromDLPack(DLManagedTensor *tensor, size_t require_alignment = 0, bool require_contiguous = false)#

Create a Tensor from a DLPack managed tensor, pre v1.0 API.

Note

This function will not run any checks on flags.

Parameters:
  • tensor – The input DLPack managed tensor.

  • require_alignment – The minimum alignment requored of the data + byte_offset.

  • require_contiguous – Boolean flag indicating if we need to check for contiguity.

Returns:

The created Tensor.

static inline Tensor FromDLPackVersioned(DLManagedTensorVersioned *tensor, size_t require_alignment = 0, bool require_contiguous = false)#

Create a Tensor from a DLPack managed tensor, post v1.0 API.

Parameters:
  • tensor – The input DLPack managed tensor.

  • require_alignment – The minimum alignment requored of the data + byte_offset.

  • require_contiguous – Boolean flag indicating if we need to check for contiguity.

Returns:

The created Tensor.

Protected Functions

inline TensorObj *get_mutable() const#

Get mutable internal container pointer.

Returns:

a mutable container pointer.