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

Address review comment.

Browse Source
This commit is contained in:
Peter Hawkins 2019-08-15 12:33:36 -04:00
parent e57a5c42c5
commit e28e73b38f
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/lax/lax.py
View File

@ -2735,9 +2735,10 @@ def _dynamic_slice_jvp(primals, tangents, slice_sizes, operand_shape):
def _dynamic_slice_transpose_rule(t, operand, *start_indices, slice_sizes=None,
operand_shape=None):
assert operand is None
assert all(s is not None for s in start_indices)
zeros = full(operand_shape, tie_in(t, _zero(t)))
return ([dynamic_update_slice(zeros, t, start_indices)] +
[ad_util.zero] * len(start_indices))
[None] * len(start_indices))
def _batch_dynamic_slice_indices(indices, bdims):
size = next((x.shape[i] for x, i in zip(indices, bdims) if i is not None), -1)