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

fix symbolic zero handling in _pad_transpose

Browse Source 
tested manually against example from @matthewdhoffman
This commit is contained in:
Matthew Johnson 2019-07-31 13:27:19 -07:00
parent 6434d9d029
commit 0600b738f4
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/lax/lax.py
View File

@ -2386,8 +2386,11 @@ def _pad_shape_rule(operand, padding_value, padding_config):
return tuple(out_shape)
def _pad_transpose(t, operand, padding_value, padding_config):
lo, hi, interior = zip(*padding_config)
if t is ad_util.zero:
return [ad_util.zero if operand is None else None,
ad_util.zero if padding_value is None else None]
lo, hi, interior = zip(*padding_config)
total = lambda x: _reduce_sum(x, list(range(t.ndim)))
def t_op():