This document is relevant for: Inf2, Trn1, Trn1n

nki.isa.reciprocal#

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

Compute reciprocal of the input data tile. Uses Scalar Engine or Vector Engine based on the number of elements per partition in data.

Estimated instruction cost:

if N >= 64:

N Scalar Engine cycles, where N is the number of elements per partition in data.

else:

8*N Vector Engine cycles, where N is the number of elements per partition in data.

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, Trn1n