This document is relevant for: Inf2, Trn1, Trn2

nki.isa.reciprocal#

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

Compute reciprocal of each element in the input data tile using Vector Engine.

Estimated instruction cost:

max(MIN_II, 8*N) Vector Engine cycles, where N is the number of elements per partition in data, and MIN_II is the minimum instruction initiation interval for small input tiles. MIN_II is roughly 64 engine cycles.

Parameters:

  • data – the input tile

  • 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.

  • mask – (optional) a compile-time constant predicate that controls whether/how this instruction is executed (see NKI API Masking for details)

Returns:

an output tile of reciprocal computation

Example:

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

x = nl.load(in_tensor[nl.mgrid[0:128, 0:512]])

y = nisa.reciprocal(x)

This document is relevant for: Inf2, Trn1, Trn2