Skip to contents

Primitive Top-K

Source: R/primitives.R
prim_top_k.Rd

Returns the k largest values along the last axis, sorted in descending order, together with their indices into that axis.

For other axes, transpose so the target axis is last, call prim_top_k(), then transpose back. nv_top_k() does this.

Usage

prim_top_k(x, k)

Arguments

x

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

k

(integer(1))
Number of top elements. Must satisfy 1 <= k <= shape(x)[naxes(x)].

Value

list of two arrayish values:
The top-k values (same dtype as x) and their indices along the last axis (dtype i32, matching JAX). Both have the same shape as x with the last axis replaced by k. Ties are broken by lower index first.

Implemented Rules

  • stablehlo

  • reverse

StableHLO

Lowers to hlo_top_k().

See also

nv_top_k(), prim_sort()

Examples

x <- nv_array(c(3, 1, 4, 1, 5, 9, 2, 6))
prim_top_k(x, k = 3L)
#> [[1]]
#> AnvlArray
#>  9
#>  6
#>  5
#> [ CPUf32{3} ] 
#> 
#> [[2]]
#> AnvlArray
#>  6
#>  8
#>  5
#> [ CPUi32{3} ] 
#>