This document is relevant for: Inf2, Trn1, Trn1n

nki.kernels.select_and_scatter_kernel#

nki.kernels.select_and_scatter_kernel(operand_tensor, source_tensor, out_tensor)[source]#

Implementation of a select-and-scatter kernel.

It selects an element from each window of operand_tensor, and then scatters source_tensor to the indices of the selected positions to construct out_tensor with the same shape as the operand_tensor.

This kernel assumes that

  • windows dimensions: (3, 3)

  • windows strides: (2, 2)

  • padding: (1, 1)

  • init value: 0

  • select computation: greater-than

  • scatter computation: add

IO Tensor layouts:

  • operand_tensor: shape (n, c, h, w)

  • source_tensor : shape (n, c, src_h, src_w)

  • out_tensor : shape (n, c, h, w)

IO tensor dtypes:

  • This kernel assumes all IO tensors have the same dtype

This document is relevant for: Inf2, Trn1, Trn1n