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

jnp.moveaxis: fix bug when axes are integer dtype

Browse Source
This commit is contained in:
Jake VanderPlas 2020-09-21 16:32:11 -07:00
parent ada6f30f59
commit 2cf8d49f5b
1 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
jax/numpy/lax_numpy.py
View File

@ -1256,10 +1256,15 @@ def swapaxes(a, axis1, axis2):
@_wraps(np.moveaxis)
def moveaxis(a, source, destination):
if isinstance(source, int):
source = (source,)
if isinstance(destination, int):
destination = (destination,)
_check_arraylike("moveaxis", a)
try:
source = (operator.index(source),)
except TypeError:
pass
try:
destination = (operator.index(destination),)
except TypeError:
pass
source = tuple(_canonicalize_axis(i, ndim(a)) for i in source)
destination = tuple(_canonicalize_axis(i, ndim(a)) for i in destination)
if len(source) != len(destination):