This document is relevant for: Trn2, Trn3

nki.isa.exponential#

nki.isa.exponential(dst: NkiTensor, src: NkiTensor, max_value: float | NkiTensor = 0.0, reduce_res: NkiTensor | None = None, reduce_cmd: reduce_cmd = reduce_cmd.idle, reduce_init: float | NkiTensor = 0.0, name=None)[source]#

Apply exponential function to each element after subtracting a max_value using Vector Engine.

Note

Available only on NeuronCore-v4 and newer.

This instruction computes exp(src - max_value) for each element. The instruction can optionally maintain a running sum of the exponential values using shared internal reduction registers in the Vector Engine.

The exponential operation is performed as:

dst[i] = exp(src[i] - max_value)

When accumulation is enabled through reduce_cmd, the instruction also computes:

reduce_res[i] = sum(dst[i])

The Vector Engine performs the computation in float32 precision internally and can output results in various data types as specified by the dst dtype field. See nki-type-casting for more information about type casting behavior.

Constraints

Supported engines: Vector.
src, dst must have the same number of elements in the partition dimension.
src, dst must have the same number of elements in the free dimensions.
src and dst can have at most three free dimensions (maximum rank 4, including the partition dimension).
reduce_init should be unset or set to 0.0 when reduce_cmd is not load_reduce.

Tensor indirection.

dst and src support tensor indirection (gather/scatter) by passing a view created with .indirect(index). max_value, reduce_res, and reduce_init do not support tensor indirection. Runs on the Vector engine. (exponential itself is only available on NeuronCore-v4 and later.)

When operands are manually allocated, their base partitions must satisfy:

the index of every .indirect() view starts on a quadrant boundary (a multiple of 32);
if src uses .indirect(), src starts on the same partition as its index;
if dst uses .indirect() and src is in SBUF, dst’s index starts on the same partition as src;
if dst uses .indirect() and src is in PSUM and uses .indirect(), dst’s index starts on the same partition as src’s index.

Parameters:

dst – The output tile with exponential function applied. Supported buffers: SBUF, PSUM. Supported dtypes: float8_e4m3, float8_e5m2, float16, bfloat16, float32, tfloat32, int8, int16, int32, uint8, uint16.
src – The input tile to apply exponential function on. Supported buffers: SBUF, PSUM. Supported dtypes: float8_e4m3, float8_e5m2, float16, bfloat16, float32, int8, int16, int32, uint8, uint16, uint32.
max_value – The maximum value to subtract from each element before applying exponential (for numerical stability). Can be a scalar or vector of shape (src.shape[0], 1). Supported dtypes: float32.
reduce_res –
Optional tile to store reduction results (sum of exponentials). Must have shape (src.shape[0], 1). Supported buffers: SBUF, PSUM. Supported dtypes: float8_e4m3, float8_e5m2, float16, bfloat16, float32, tfloat32.

Pass None to keep the reduction result in the Vector Engine’s internal accumulator without writing it out. This is useful when chaining multiple calls that reduce into the same accumulator — only the final call needs to pass a tile to retrieve the accumulated result.
reduce_cmd – Control the state of reduction registers for accumulating exponential results. Supported: idle, reset_reduce, reduce, load_reduce.
reduce_init – Initial value for reduction when using reduce_cmd.load_reduce. Supported dtypes: float32.

Accumulator behavior:

The Vector Engine maintains internal accumulator registers that can be controlled via the reduce_cmd parameter:

reduce_cmd.reset_reduce: Reset accumulators to 0, then accumulate the current results.
reduce_cmd.reduce: Continue accumulating without resetting (useful for multi-step reductions).
reduce_cmd.load_reduce: Load the values from reduce_init into the accumulator, then accumulate the current result on top of it.
reduce_cmd.idle: (default) No accumulation performed, accumulator state unknown.

Note

Even when reduce_cmd is set to idle, the accumulator state may still be modified. Always use reset_reduce after any Vector Engine operation that ran with idle mode to ensure consistent behavior.

Note

The accumulator registers are shared across Vector Engine accumulation instructions including exponential(), range_select(), scalar_tensor_tensor(), select_reduce(), tensor_scalar_cumulative(), and tensor_scalar_reduce().

Behavior

# Initialize reduction if requested
if reduce_cmd == reduce_cmd.reset_reduce:
    accumulator = 0
elif reduce_cmd == reduce_cmd.load_reduce:
    accumulator = reduce_init
elif reduce_cmd == reduce_cmd.idle:
    accumulator = undefined  # Not used

# Process each element
for i in range(num_elements):
    dst[i] = exp(src[i] - max_value)

    # Update reduction if active
    if reduce_cmd != reduce_cmd.idle:
        accumulator += dst[i]

This document is relevant for: Trn2, Trn3

nki.isa.exponential

Contents

nki.isa.exponential#