Skip to contents

TopK Operator (CHLO)

Source: R/op-top_k.R
hlo_top_k.Rd

This op is from the CHLO dialect, a higher-level companion to stableHLO that is lowered to stableHLO during compilation. See https://openxla.org/stablehlo/generated/chlo#chlotop_k_chlotop_kop for details.

Usage

infer_types_top_k(operand, k)

hlo_top_k(operand, k)

Arguments

operand

(FuncValue)
Tensor of integer, unsigned integer, or floating-point type with rank >= 1.

k

(integer(1))
Number of top elements to return along the last dimension. Must satisfy 1 <= k <= dim(operand, -1).

Value

FuncValue

A list() of two FuncValues: the top-k values (same dtype as operand) and their indices into the last dimension (dtype i32). Ties are broken by lower index first.