This document is relevant for: Trn1, Trn2, Trn3

nki.language.rand#

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

Create a new tensor of given shape and dtype on the specified buffer, filled with random values.

Values are sampled from a uniform distribution between 0 and 1.

Warning

This API is experimental and may change in future releases.

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 (sbuf, psum, hbm), defaults to sbuf.
name – the name of the tensor.

Returns:

a new tensor allocated on the buffer with random values.

Examples:

import nki.language as nl

# nki.language.rand -- generate random values in [0, 1)
a = nl.rand((128, 512), dtype=nl.float32)

This document is relevant for: Trn1, Trn2, Trn3

nki.language.rand