This document is relevant for: Inf2, Trn1, Trn1n

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`	Get the data type of the tensor.
`ndim`	Get the number of dimensions of the tensor.
`shape`	Get the 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#

Get the data type of the tensor.

Returns:: The 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 ndim#: Get the 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#

Get the shape of the tensor.

Returns:: The 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, Trn1n

nki.tensor

Contents

nki.tensor#