Repository: google/jax

Issue:
Inconsistent transpose behaviour between jax.numpy and numpy 
Example:
```
# Fine.
np.arange(6).reshape(1, 3, 2).transpose(2, 1, 0)

# TypeError: transpose() takes from 1 to 2 positional arguments but 4 were given
jnp.arange(6).reshape(1, 3, 2).transpose(2, 1, 0)
```


Discussion:
Thanks for the report - the multi-argument version of this is not yet implemented. Until we rectify that, you can pass a tuple of axes:
```python
jnp.arange(6).reshape(1, 3, 2).transpose((2, 1, 0))
```

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -30,7 +30,7 @@
 import operator
 import os
 import types
-from typing import Sequence, FrozenSet, Optional, Tuple, Union, cast
+from typing import Any, Sequence, FrozenSet, Optional, Tuple, Union, cast
 from textwrap import dedent as _dedent
 import warnings
 
@@ -1302,6 +1302,22 @@ def _reshape(a, *args, order="C"):
   else:
     raise ValueError("Unexpected value for 'order' argument: {}.".format(order))
 
+def _ensure_index_tuple(x: Any) -> Tuple[int, ...]:
+  """Convert x to a tuple of indices."""
+  try:
+    return (operator.index(x),)
+  except TypeError:
+    return tuple(map(operator.index, x))
+
+def _transpose(a, *args):
+  if not args:
+    axis = None
+  elif len(args) == 1:
+    axis = args[0] if args[0] is None else _ensure_index_tuple(args[0])
+  else:
+    axis = _ensure_index_tuple(args)
+  return transpose(a, axis)
+
 @_wraps(np.ravel)
 def ravel(a, order="C"):
   _check_arraylike("ravel", a)
@@ -5334,7 +5350,7 @@ def _operator_round(number, ndigits=None):
 _diff_methods = ["clip", "conj", "conjugate", "cumprod", "cumsum",
                  "diagonal", "dot", "max", "mean", "min", "prod", "ptp",
                  "ravel", "repeat", "sort", "squeeze", "std", "sum",
-                 "swapaxes", "take", "tile", "trace", "transpose", "var"]
+                 "swapaxes", "take", "tile", "trace", "var"]
 
 # These methods are mentioned explicitly by nondiff_methods, so we create
 # _not_implemented implementations of them here rather than in __init__.py.
@@ -5351,6 +5367,7 @@ def _operator_round(number, ndigits=None):
 for method_name in _nondiff_methods + _diff_methods:
   setattr(ShapedArray, method_name, core.aval_method(globals()[method_name]))
 setattr(ShapedArray, "reshape", core.aval_method(_reshape))
+setattr(ShapedArray, "transpose", core.aval_method(_transpose))
 setattr(ShapedArray, "flatten", core.aval_method(ravel))
 setattr(ShapedArray, "T", core.aval_property(transpose))
 setattr(ShapedArray, "real", core.aval_property(real))
@@ -5368,6 +5385,7 @@ def _operator_round(number, ndigits=None):
   for method_name in _nondiff_methods + _diff_methods:
     setattr(device_array, method_name, globals()[method_name])
   setattr(device_array, "reshape", _reshape)
+  setattr(device_array, "transpose", _transpose)
   setattr(device_array, "flatten", ravel)
   setattr(device_array, "T", property(transpose))
   setattr(device_array, "real", property(real))


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1661,6 +1661,28 @@ def testPiecewise(self, shape, dtype, ncond, nfunc):
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, check_cache_misses=False)
 
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+    {"testcase_name": "{}_perm={}_{}".format(
+      jtu.format_shape_dtype_string(shape, dtype), perm, arg_type),
+     "dtype": dtype, "shape": shape, "perm": perm, "arg_type": arg_type}
+    for dtype in default_dtypes
+    for shape in array_shapes
+    for arg_type in ["splat", "value"]
+    for perm in [None, tuple(np.random.RandomState(0).permutation(np.zeros(shape).ndim))]))
+  def testTransposeTuple(self, shape, dtype, perm, arg_type):
+    rng = jtu.rand_some_zero(self.rng())
+    args_maker = lambda: [rng(shape, dtype)]
+    if arg_type == "value":
+      np_fun = lambda x: x.transpose(perm)
+      jnp_fun = lambda x: jnp.array(x).transpose(perm)
+    else:
+      np_fun = lambda x: x.transpose(*(perm or ()))
+      jnp_fun = lambda x: jnp.array(x).transpose(*(perm or ()))
+
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "{}_trim={}".format(
       jtu.format_shape_dtype_string(a_shape, dtype), trim),