jnp.moveaxis: fix bug when axes are integer dtype

[jakevdp]

Before:

In [1]: import jax.numpy as jnp                                                                                                          

In [2]: x = jnp.arange(6).reshape(2, 3) 

In [3]: jnp.moveaxis(x, jnp.int32(0), jnp.int32(1))                                                                                      
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-53fba72727a1> in <module>
----> 1 jnp.moveaxis(x, jnp.int32(0), jnp.int32(1))

~/github/google/jax/jax/numpy/lax_numpy.py in moveaxis(a, source, destination)
   1261   if isinstance(destination, int):
   1262     destination = (destination,)
-> 1263   source = tuple(_canonicalize_axis(i, ndim(a)) for i in source)
   1264   destination = tuple(_canonicalize_axis(i, ndim(a)) for i in destination)
   1265   if len(source) != len(destination):

~/github/google/jax/jax/interpreters/xla.py in __iter__(self)
   1060   def __iter__(self):
   1061     if self.ndim == 0:
-> 1062       raise TypeError("iteration over a 0-d array")  # same as numpy error
   1063     else:
   1064       return self._value.__iter__()

TypeError: iteration over a 0-d array

After:

In [1]: import jax.numpy as jnp                                                                                                          

In [2]: x = jnp.arange(6).reshape(2, 3)

In [3]: jnp.moveaxis(x, jnp.int32(0), jnp.int32(1))                                                                                      
Out[3]: 
DeviceArray([[0, 3],
             [1, 4],
             [2, 5]], dtype=int32)