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