This document is relevant for: Inf2, Trn1, Trn2

nki.tensor#

class nki.tensor[source]#

A tensor object represents a multidimensional, homogeneous array of fixed-size items

Methods

assert_shape

Assert that the tensor has the given shape.

astype

Copy of the tensor, cast to a specified type.

broadcast_to

Broadcast tensor to a new shape based on numpy broadcast rules.

expand_dims

Gives a new shape to a tensor by adding a dimension of size 1 at the specified position.

reshape

Gives a new shape to an array without changing its data.

view

Return a new view of the tensor, reinterpret to a specified type.

Attributes

dtype

Data type of the tensor.

itemsize

Length of one tensor element in bytes.

ndim

Number of dimensions of the tensor.

shape

Shape of the tensor.
assert_shape(shape)[source]#

Assert that the tensor has the given shape.

Parameters:

shape – The expected shape.

Returns:

The tensor.

astype(dtype)[source]#

Copy of the tensor, cast to a specified type.

Parameters:

dtype – The target dtype

Returns:

the tensor with new type. Copy ALWAYS occur

broadcast_to(shape)[source]#

Broadcast tensor to a new shape based on numpy broadcast rules. The tensor object must be a tile or can be implicitly converted to a tile. A tensor can be implicitly converted to a tile iff the partition dimension is the highest dimension.

Parameters:

shape – The new shape

Returns:

Return a new view of the tensor, no copy will occur

property dtype#

Data type of the tensor.

expand_dims(axis)[source]#

Gives a new shape to a tensor by adding a dimension of size 1 at the specified position.

Parameters:

axis – the position of the new dimension.

Returns:

Return a new tensor with expanded shape

property itemsize#

Length of one tensor element in bytes.

property ndim#

Number of dimensions of the tensor.

reshape(shape)[source]#

Gives a new shape to an array without changing its data.

Parameters:

shape – The new shape

Returns:

Return a new view of the tensor, no copy will occur

property shape#

Shape of the tensor.

view(dtype)[source]#

Return a new view of the tensor, reinterpret to a specified type.

Returns:

A new tensor object refer to the original tensor data, NO copy will occur

This document is relevant for: Inf2, Trn1, Trn2