This document is relevant for: Trn2, Trn3

nki.language.copy#

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

Create a copy of the input tile on SBUF.

Note

Using a dtype different from the input’s data type may result in precision loss.

Parameters:

  • x – the source of copy, must be a tile in SBUF or PSUM.

  • 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 new tile with the same layout as x, allocated on SBUF.

Examples:

import nki.language as nl

# nki.language.copy -- copy an nl.matmul result from PSUM back to SBUF.
# nl.matmul writes its result to PSUM; nl.copy brings it to SBUF.
x = nl.load(x_tensor)
y = nl.load(y_tensor)
result_psum = nl.matmul(x, y, transpose_x=True)  # x.T @ y, in PSUM
result = nl.copy(result_psum)  # PSUM -> SBUF

This document is relevant for: Trn2, Trn3