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