tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

optimize import

Browse Source
This commit is contained in:
wls2002
2023-06-29 09:41:49 +08:00
parent d28cef1a87
commit 01b7731231
14 changed files with 29 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/jax_playground.py
View File

@@ -1,8 +1,8 @@
from functools import partial
import jax
from jax import numpy as jnp, jit
@partial(jit, static_argnames=['reverse'])
def rank_element(array, reverse=False):
"""
@@ -14,5 +14,5 @@ def rank_element(array, reverse=False):
return jnp.argsort(jnp.argsort(array))
a = jnp.array([1 ,5, 3, 5, 2, 1, 0])
print(rank_element(a, reverse=True))
a = jnp.array([1, 5, 3, 5, 2, 1, 0])
print(rank_element(a, reverse=True))