nki.language.ones — AWS Neuron Documentation (original) (raw)

This document is relevant for: Inf2, Trn1, Trn2

nki.language.ones#

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

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

((Similar to numpy.ones))

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.

Returns:

a new tensor allocated on the buffer.

This document is relevant for: Inf2, Trn1, Trn2