tweak handling of random.permutation scalar case

jax/random.py

@@ -444,11 +444,14 @@ def permutation(key, x):
   Returns:
     A shuffled version of x or array range
   """
-  if isinstance(x, (int, onp.integer)):
-    arr = _shuffle(key, onp.arange(x), 0)
+  if not onp.shape(x):
+    # scalar case, must be a concrete integer
+    if not onp.issubdtype(lax.dtype(x), onp.integer):
+      raise TypeError("x must be an integer or at least 1-dimensional")
+    x = int(x)  # TODO(mattjj): concrete tracer error from core.py
+    return _shuffle(key, np.arange(x), 0)
   else:
-    arr = _shuffle(key, x, 0)
-  return arr
+    return _shuffle(key, x, 0)
 
 
 @partial(jit, static_argnums=(2,))

tests/random_test.py

@@ -25,6 +25,7 @@
 import scipy.stats
 
 from jax import api
+from jax import core
 from jax import grad
 from jax import lax
 from jax import numpy as np
@@ -202,6 +203,13 @@ def testPermutationInteger(self):
     self.assertFalse(onp.all(perm1 == onp.arange(100)))  # seems unlikely!
     self.assertTrue(onp.all(onp.sort(perm1) == onp.arange(100)))
 
+  def testPermutationErrors(self):
+    key = random.PRNGKey(0)
+    with self.assertRaises(TypeError):
+      random.permutation(key, 10.)
+    with self.assertRaises(core.ConcretizationTypeError):
+      api.jit(random.permutation)(key, 10)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_p={}_{}".format(p, dtype),
        "p": p, "dtype": onp.dtype(dtype).name}