This document is relevant for: Trn1, Trn2, Trn3

nki.language.broadcast_to#

nki.language.broadcast_to(x, shape, dtype=None)[source]#

Broadcast a tile to a new shape following numpy broadcasting rules.

((Similar to numpy.broadcast_to))

Warning

This API is experimental and may change in future releases.

If x.shape is already the same as shape, returns x unchanged (or a dtype-cast copy if dtype differs).

Parameters:

  • x – the source tile in SBUF or PSUM.

  • shape – the target shape. Must have the same rank as x. Each dimension must either match or be broadcast from size 1.

  • dtype – (optional) data type to cast the output type to (see Supported Data Types for more information); if not specified, it will default to be the same as the data type of the input tile.

Returns:

a tile with the target shape containing broadcast values from x.

This document is relevant for: Trn1, Trn2, Trn3