jax.lax.sort_key_val

Warning

This page was created from a pull request (#9655).

jax.lax.sort_key_val¶

jax.lax.sort_key_val(keys, values, dimension=- 1, is_stable=True)[source]¶

Sorts keys along dimension and applies the same permutation to values.

Parameters

keys (Any) –
values (Any) –
dimension (int) –
is_stable (bool) –

Return type

Tuple[Any, Any]

previous

jax.lax.sort

next

jax.lax.sqrt