nki.isa.tensor_partition_reduce — AWS Neuron Documentation (original) (raw)

This document is relevant for: Inf2, Trn1, Trn2

nki.isa.tensor_partition_reduce#

nki.isa.tensor_partition_reduce(op, data, *, mask=None, dtype=None, **kwargs)[source]#

Apply a reduction operation across partitions of an input data tile using GpSimd Engine.

Parameters:

• op – the reduction operator (add, max, bitwise_or, bitwise_and)
• data – the input tile to be reduced
• mask – (optional) a compile-time constant predicate that controls whether/how this instruction is executed (see NKI API Masking for details)
• 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:

output tile with reduced result

Example:

import neuronxcc.nki.isa as nisa import neuronxcc.nki.language as nl import numpy as np ...

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

Example 1: reduce add tile a of shape (128, 32, 4)

in the partition dimension and return

reduction result in tile b of shape (1, 32, 4)

################################################################## a = nl.load(a_tensor[0:128, 0:32, 0:4])
b = nisa.tensor_partition_reduce(np.add, a) nl.store(b_tensor[0:1, 0:32, 0:4], b)

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

Example 2: reduce add tile a of shape (b, p, f1, ...)

in the partition dimension p and return

reduction result in tile b of shape (b, 1, f1, ...)

################################################################## for i in nl.affine_range(a_tensor.shape[0]): a = nl.load(a_tensor[i]) b = nisa.tensor_partition_reduce(np.add, a) nl.store(b_tensor[i], b)

This document is relevant for: Inf2, Trn1, Trn2