nki.language.ndarray#

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

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

((Similar to numpy.ndarray))

Parameters:

  • shape – the shape of the tensor.

  • dtype – the data type of the tensor (see Supported Data Types for more information).

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

  • name – the name of the tensor. The name parameter has to be unique for tensors on each Physical NeuronCore(PNC) within each Logical NeuronCore(LNC). It is optional for SRAM tensors, IO tensors, and any HBM tensors that are only visible to one Physical NeuronCore. For shared_hbm tensors that are not used as kernel inputs or outputs, name must be specified. In addition, the compiler uses the name to link non-IO shared_hbm tensors among PNCs. In other word, shared_hbm tensors will point to the same underlying memory as long as they have the same name, even if the tensors appear in different control flow.

Returns:

a new tensor allocated on the buffer.