This document is relevant for: Trn2, Trn3

Scatter-Add Kernel API Reference#

Scatter-add from src into input based on indices using gather-accumulate-scatter pattern.

Equivalent to PyTorch’s input.scatter_add(dim=0, index=index, src=src). Performs: input[index[i], j] += src[i, j] for all i, j.

Background#

The scatter_add kernel scatter-adds rows from src into a 2D input tensor based on a 1D index tensor, using a gather-accumulate-scatter pattern. Indices within a tile of 128 rows should be unique for correctness.

API Reference#

Source code for this kernel API can be found at: scatter_add.py

scatter_add#

nkilib.experimental.misc.scatter_add(input: nl.ndarray, dim: int, index: nl.ndarray, src: nl.ndarray) → nl.ndarray#

Scatter-add from src into input based on indices using gather-accumulate-scatter pattern.

Parameters:

  • input (nl.ndarray) – [N, D], Destination tensor to accumulate into (modified in-place)

  • dim (int) – Dimension along which to scatter (must be 0)

  • index (nl.ndarray) – [K], 1D tensor of row indices into input

  • src (nl.ndarray) – [K, D], Source values to scatter-add

Returns:

[N, D], The input tensor with scattered values added

Return type:

nl.ndarray

Notes:

  • Input and src tensors must be 2D

  • Index tensor must be 1D

  • dim must be 0

  • Indices within a tile of 128 rows should be unique for correctness

Dimensions:

  • N: Number of rows in input tensor

  • D: Feature dimension size

  • K: Number of source rows / indices

This document is relevant for: Trn2, Trn3