Primitive Sort

Sorts arrays along the given axis.

Sorting is determined by the first array only: it is the sort key, and any additional arrays are reordered with the same permutation that sorts the first. This enables idioms like argsort (sort x paired with an iota and read off the second output) and key-value sorts (sort keys paired with values).

1-D slices along axis are sorted independently; other axes are preserved.

Usage

prim_sort(xs, axis, descending = FALSE, is_stable = FALSE)

Arguments

xs

(list of arrayish)
One or more arrays to sort. The first is the sort key; the rest are carried along under the same permutation. All must share the same shape.

axis

(integer(1))
Axis along which to sort. Negative values count from the end, i.e. -1 refers to the last axis.

descending

(logical(1))
If TRUE, sort the key in descending order (largest first). Default FALSE. Additional arrays are reordered by the same permutation regardless.

is_stable

(logical(1))
If TRUE, the sort is stable: the relative order of equal keys is preserved. Default FALSE.

Value

(list of arrayish)
One sorted output per element of xs, in the same order, each with the shape and data type of its input.

Implemented Rules

• stablehlo

• reverse

StableHLO

Lowers to hlo_sort(), specified under sort. The comparator uses hlo_compare() (LT for ascending, GT for descending) on the first array. For float keys it uses compare_type = "TOTALORDER" and canonicalizes -0/+0 and -NaN/+NaN to their positive form before comparing, so all NaN values land at one end of the result regardless of sign. Integer keys use SIGNED / UNSIGNED as appropriate.

See also

nv_sort(), nv_argsort(), nv_top_k(), nv_median()

Examples

x <- nv_array(c(3, 1, 4, 1, 5))
prim_sort(list(x), axis = 1L)[[1L]]
#> AnvlArray
#>  1
#>  1
#>  3
#>  4
#>  5
#> [ CPUf32{5} ] 

# sort indices by the values (argsort): pair x with iota and read off
# the second result
idx <- nv_iota(axis = 1L, dtype = "i64", shape = 5L)
out <- prim_sort(list(x, idx), axis = 1L)
out[[1L]] # sorted x
#> AnvlArray
#>  1
#>  1
#>  3
#>  4
#>  5
#> [ CPUf32{5} ] 
out[[2L]] # permutation indices
#> AnvlArray
#>  2
#>  4
#>  1
#>  3
#>  5
#> [ CPUi64{5} ]