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

Temporarily remove jit decorator on gather/scatter ops.

Browse Source
This commit is contained in:
Peter Hawkins 2019-09-16 13:57:07 -07:00
parent d691f81264
commit 45a02f39f0
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/numpy/lax_numpy.py
View File

@ -2407,7 +2407,9 @@ def _rewriting_take(arr, idx):
treedef, static_idx, dynamic_idx = _split_index_for_jit(idx)
return _gather(arr, treedef, static_idx, dynamic_idx)
@partial(jit, static_argnums=(1, 2))
# TODO(phawkins): re-enable jit after fixing excessive recompilation for
# slice indexes (e.g., slice(0, 5, None), slice(10, 15, None), etc.).
# @partial(jit, static_argnums=(1, 2))
def _gather(arr, treedef, static_idx, dynamic_idx):
idx = _merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx)
indexer = _index_to_gather(shape(arr), idx) # shared with _scatter_update

4
jax/ops/scatter.py
View File

@ -57,7 +57,9 @@ def _scatter_update(x, idx, y, scatter_op):
return _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx)
@partial(jit, static_argnums=(2, 3, 4))
# TODO(phawkins): re-enable jit after fixing excessive recompilation for
# slice indexes (e.g., slice(0, 5, None), slice(10, 15, None), etc.).
# @partial(jit, static_argnums=(2, 3, 4))
def _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx):
y = lax.convert_element_type(y, lax.dtype(x))