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

transpose shouldn't transpose with identity perm

Browse Source
This commit is contained in:
Matthew Johnson 2018-12-12 17:53:37 -08:00
parent 538850e271
commit 77d6fb4c01
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jax/lax.py
View File

@ -244,7 +244,11 @@ def _index_untake(axes, src, dst, *idxs):
return dst
def transpose(operand, permutation):
return transpose_p.bind(operand, permutation=tuple(permutation))
permutation = tuple(permutation)
if permutation == tuple(range(len(permutation))):
return operand
else:
return transpose_p.bind(operand, permutation=permutation)
def reduce(operand, init_value, computation, dimensions):
monoid_reducer = _get_monoid_reducer(computation, init_value)