This document is relevant for: Trn2, Trn3

nki.isa.tensor_partition_reduce#

nki.isa.tensor_partition_reduce(dst: NkiTensor, op, data: NkiTensor, name=None)[source]#

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

The data and dst tiles can have at most four free dimensions (maximum rank 5, including the partition dimension).

Parameters:

  • dst – output tile with reduced result

  • op – the reduction operator (add, mean, max, bitwise_or, bitwise_and, bitwise_xor)

  • data – the input tile to be reduced

This document is relevant for: Trn2, Trn3