This document is relevant for: Trn2, Trn3

nki.language.ndarray#

nki.language.ndarray(shape, dtype, buffer=MemoryRegion.sbuf, name='', address=None)[source]#

Create a new tensor of given shape and dtype on the specified buffer.

Parameters:

  • shape – the shape of the tensor.

  • dtype – the data type of the tensor.

  • buffer – the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf.

  • name – the name of the tensor, used in scheduling.

  • address – optional memory address (partition_offset, free_offset).

Returns:

a new NkiTensor allocated on the buffer.

This document is relevant for: Trn2, Trn3