nki_samples.reference.vision.select_and_scatter_kernel

nki_samples.reference.vision.select_and_scatter_kernel = <neuronxcc.nki.compile.GenericKernel object>

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