mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-25 03:06:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

jit lax_numpy.roll (#2392)

Browse Source 
This was making tracing slow for code with lots of rolls.
This commit is contained in:
Stephan Hoyer 2020-03-09 13:21:30 -07:00 committed by GitHub
parent f3f0abb53e
commit 863576c5c1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
jax/numpy/lax_numpy.py
View File

@ -2672,8 +2672,8 @@ def msort(a):
return sort(a, axis=0)
@_wraps(onp.roll)
def roll(a, shift, axis=None):
@partial(jit, static_argnums=(2,))
def _roll(a, shift, axis):
a = asarray(a)
a_shape = shape(a)
if axis is None:
@ -2696,6 +2696,11 @@ def roll(a, shift, axis=None):
return a
@_wraps(onp.roll)
def roll(a, shift, axis=None):
return _roll(a, shift, axis)
@_wraps(onp.take)
def take(a, indices, axis=None, out=None, mode=None):
if out: