Repository: google/jax

Issue:
jax.random bug if "jax_disable_jit" == True
jax.random.choice does not work as expected if jit is disabled.

```python
import jax 
jax.config.update('jax_disable_jit', True)
jax.random.choice(jax.random.PRNGKey(1234), 8, shape=(100,))
```

Output:
```
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/random.py", line 619, in choice
    ind = randint(key, shape, 0, n_inputs)
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/random.py", line 438, in randint
    return _randint(key, shape, minval, maxval, dtype)
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/random.py", line 446, in _randint
    minval = _asarray(minval)
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/random.py", line 92, in _asarray
    raise TypeError(f"Function requires array input, got {x} of type {type(x)}.")
TypeError: Function requires array input, got 0 of type <class 'int'>.
```



Discussion:
Thanks for the report – I think this check is too restrictive. We should probably use the `_check_arraylike` utility from `lax_numpy.py`.

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -25,7 +25,7 @@
 from jax._src import dtypes
 from jax.core import NamedShape
 from jax._src.api import jit, vmap
-from jax._src.numpy.lax_numpy import _constant_like, _convert_and_clip_integer, asarray
+from jax._src.numpy.lax_numpy import _constant_like, _convert_and_clip_integer, _check_arraylike
 from jax.lib import xla_bridge
 from jax.lib import xla_client
 from jax.lib import cuda_prng
@@ -86,14 +86,6 @@ def _is_prng_key(key: jnp.ndarray) -> bool:
 ### utilities
 
 
-# TODO(mattjj,jakevdp): add more info to error message, use this utility more
-def _asarray(x):
-  """A more restrictive jnp.asarray, only accepts JAX arrays and np.ndarrays."""
-  if not isinstance(x, (np.ndarray, jnp.ndarray)):
-    raise TypeError(f"Function requires array input, got {x} of type {type(x)}.")
-  return jnp.asarray(x)
-
-
 def _make_rotate_left(dtype):
   if not jnp.issubdtype(dtype, np.integer):
     raise TypeError("_rotate_left only accepts integer dtypes.")
@@ -143,7 +135,7 @@ def rolled_loop_step(i, state):
   x, ks, rotations = state
   for r in rotations[0]:
     x = apply_round(x, r)
-  new_x = [x[0] + ks[0], x[1] + ks[1] + asarray(i + 1, dtype=np.uint32)]
+  new_x = [x[0] + ks[0], x[1] + ks[1] + jnp.asarray(i + 1, dtype=np.uint32)]
   return new_x, rotate_list(ks), rotate_list(rotations)
 
 def _threefry2x32_lowering(key1, key2, x1, x2, use_rolled_loops=True):
@@ -445,8 +437,9 @@ def _randint(key, shape, minval, maxval, dtype):
   if not jnp.issubdtype(dtype, np.integer):
     raise TypeError(f"randint only accepts integer dtypes, got {dtype}")
 
-  minval = _asarray(minval)
-  maxval = _asarray(maxval)
+  _check_arraylike("randint", minval, maxval)
+  minval = jnp.asarray(minval)
+  maxval = jnp.asarray(maxval)
   if not jnp.issubdtype(minval.dtype, np.integer):
     minval = minval.astype(int)
   if not jnp.issubdtype(maxval.dtype, np.integer):
@@ -603,10 +596,11 @@ def choice(key: jnp.ndarray,
                     f"got {shape}")
   if np.ndim(a) not in [0, 1]:
     raise ValueError("a must be an integer or 1-dimensional")
+  _check_arraylike("choice", a)
   if np.ndim(a) == 0:
-    a = int(a)
+    a = core.concrete_or_error(int, a, "The error occurred in jax.random.choice()")
   else:
-    a = _asarray(a)
+    a = jnp.asarray(a)
   n_inputs = int(a) if np.ndim(a) == 0 else len(a)  # type: ignore[arg-type]
   n_draws = prod(shape)
   if n_draws == 0:


Test patch: