18 lines
425 B
Python
18 lines
425 B
Python
from functools import partial
|
|
|
|
import jax
|
|
from jax import numpy as jnp, jit
|
|
|
|
@partial(jit, static_argnames=['reverse'])
|
|
def rank_element(array, reverse=False):
|
|
"""
|
|
rank the element in the array.
|
|
if reverse is True, the rank is from large to small.
|
|
"""
|
|
if reverse:
|
|
array = -array
|
|
return jnp.argsort(jnp.argsort(array))
|
|
|
|
|
|
a = jnp.array([1 ,5, 3, 5, 2, 1, 0])
|
|
print(rank_element(a, reverse=True)) |