jax.numpy.transpose does not sufficiently validate inputs
For example:
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> x = np.zeros((2, 3, 4))

# jnp.transpose is silent on invalid permutations
>>> jnp.transpose(x, (0, 1, 2, 3))
array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])

# numpy transpose raises an error.
>>> np.transpose(x, (0, 1, 2, 3))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 5, in transpose
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 658, in transpose
    return _wrapfunc(a, 'transpose', axes)
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 58, in _wrapfunc
    return bound(*args, **kwds)
ValueError: axes don't match array
```
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index ddc3ef117..64e4300e5 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1216,7 +1216,7 @@ def transpose(operand: Array, permutation: Sequence[int]) -> Array:
   operator.
   """
   permutation = tuple(permutation)
-  if permutation == tuple(range(len(permutation))):
+  if permutation == tuple(range(np.ndim(operand))):
     return operand
   else:
     return transpose_p.bind(operand, permutation=permutation)
['jax/_src/lax/lax.py']