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

This document is relevant for: Inf2, Trn1, Trn2

nki.language.broadcast_to#

nki.language.broadcast_to(src, *, shape, **kwargs)[source]#

Broadcast the src tile to a new shape based on numpy broadcast rules. The src may also be a tensor object which may be implicitly converted to a tile. A tensor can be implicitly converted to a tile if the partition dimension is the outermost dimension.

Parameters:

• src – the source of broadcast, a tile in SBUF or PSUM. May also be a tensor object.
• shape – the target shape for broadcasting.

Returns:

a new tile broadcast along the partition dimension of src, this new tile will be in SBUF, but can be also assigned to a PSUM tensor.

import neuronxcc.nki.language as nl

##################################################################

Example 1: Load from in_tensor[P, F] that is on HBM and

copy into out_tile[P, F] that is on SBUF by broadcasting

################################################################## ...

...

broadcast into out_tile[P, F] that is on SBUF

from data_tile[P, F] that is on SBUF

in_tile = nl.load(in_tensor, dtype=in_tensor.dtype) out_tile = nl.broadcast_to(in_tile, shape=(128, in_tensor.shape[1]))

store output

nl.store(out_tensor, out_tile)

This document is relevant for: Inf2, Trn1, Trn2