This document is relevant for: Trn2, Trn3

Gather Kernel API Reference#

Gather rows from input based on indices using indirect DMA load.

Equivalent to PyTorch’s input[index] for dim=0 with 2D input and 1D index.

Background#

The gather kernel gathers rows from a 2D input tensor based on a 1D index tensor using an indirect DMA load, producing output[i, :] = input[index[i], :].

API Reference#

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

gather#

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

Gather rows from input based on indices using indirect DMA load.

Parameters:

  • input (nl.ndarray) – [N, D], Source tensor to gather from

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

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

Returns:

[K, D], Gathered result where output[i, :] = input[index[i], :]

Return type:

nl.ndarray

Notes:

  • Input tensor must be 2D

  • Index tensor must be 1D

  • dim must be 0

  • Under LNC sharding (num_shards > 1), the index size K must be divisible by num_shards

Dimensions:

  • N: Number of rows in input tensor

  • D: Feature dimension size

  • K: Number of indices (output rows)

This document is relevant for: Trn2, Trn3