jax.lax.sort#

jax.lax.sort(operand: Array, dimension: int = -1, is_stable: bool = True, num_keys: int = 1) → Array[源代码]#
jax.lax.sort(operand: Sequence[Array], dimension: int = -1, is_stable: bool = True, num_keys: int = 1) → tuple[Array, ...]

包装 XLA 的 Sort 操作符。

对于浮点数输入，-0.0 和 0.0 被视为等价，NaN 值被排序到数组的末尾。对于复数输入，排序顺序是实部和虚部的字典顺序，实部优先。

参数:

  • operand – 数组或数组序列

  • dimension – 整数维度，表示沿哪个维度进行排序。默认值：-1。

  • is_stable – 布尔值，指定是否使用稳定排序。默认值：True。

  • num_keys – 要视为排序键的操作数数量。默认值：1。对于 num_keys > 1，排序顺序将使用前 num_keys 个数组按字典顺序确定，第一个键为主键。其余的操作数将返回相同的排列。

返回:

输入的排序版本或排序后的输入。

返回类型:

operand