Repository: google/jax
Issue: jax.random.permutation throws error when called on length zero input

When random.permutation is called with a length zero array, it should most naturally return a length zero array. Instead it throws an error message `cannot convert float infinity to integer`.

As an example, the code
```python
random.permutation(skey, np.ones((0,)))
```
returns an error message
```
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-23-13d34f8e8bed> in <module>()
----> 1 random.permutation(skey, np.ones((0,)))

12 frames
google3/third_party/py/jax/_src/random.py in permutation(key, x)
    539     return _shuffle(key, jnp.arange(x), 0)
    540   elif np.ndim(x) == 1:
--> 541     return _shuffle(key, x, 0)
    542   else:
    543     assert isinstance(x, jnp.ndarray)

google3/third_party/py/jax/api.py in cache_miss(*args, **kwargs)
    421         backend=backend,
    422         name=flat_fun.__name__,
--> 423         donated_invars=donated_invars)
    424     out_pytree_def = out_tree()
    425     out = tree_unflatten(out_pytree_def, out_flat)

google3/third_party/py/jax/core.py in bind(self, fun, *args, **params)
   1556 
   1557   def bind(self, fun, *args, **params):
-> 1558     return call_bind(self, fun, *args, **params)
   1559 
   1560   def process(self, trace, fun, tracers, params):

google3/third_party/py/jax/core.py in call_bind(primitive, fun, *args, **params)
   1547   tracers = map(top_trace.full_raise, args)
   1548   with maybe_new_sublevel(top_trace):
-> 1549     outs = primitive.process(top_trace, fun, tracers, params)
   1550   return map(full_lower, apply_todos(env_trace_todo(), outs))
   1551 

google3/third_party/py/jax/core.py in process(self, trace, fun, tracers, params)
   1559 
   1560   def process(self, trace, fun, tracers, params):
-> 1561     return trace.process_call(self, fun, tracers, params)
   1562 
   1563   def post_process(self, trace, out_tracers, params):

google3/third_party/py/jax/core.py in process_call(self, primitive, f, tracers, params)
    597 
    598   def process_call(self, primitive, f, tracers, params):
--> 599     return primitive.impl(f, *tracers, **params)
    600   process_map = process_call
    601 

google3/third_party/py/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    576 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
    577   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
--> 578                                *unsafe_map(arg_spec, args))
    579   try:
    580     return compiled_fun(*args)

google3/third_party/py/jax/linear_util.py in memoized_fun(fun, *args)
    258       fun.populate_stores(stores)
    259     else:
--> 260       ans = call(fun, *args)
    261       cache[key] = (ans, fun.stores)
    262 

google3/third_party/py/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    649 
    650   abstract_args, arg_devices = unzip2(arg_specs)
--> 651   jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, abstract_args, transform_name="jit")
    652   if any(isinstance(c, core.Tracer) for c in consts):
    653     raise core.UnexpectedTracerError("Encountered an unexpected tracer.")

google3/third_party/py/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals, transform_name)
   1207     main.source_info = fun_sourceinfo(fun.f, transform_name)  # type: ignore
   1208     main.jaxpr_stack = ()  # type: ignore
-> 1209     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1210     del fun, main
   1211   return jaxpr, out_avals, consts

google3/third_party/py/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1186     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1187     in_tracers = map(trace.new_arg, in_avals)
-> 1188     ans = fun.call_wrapped(*in_tracers)
   1189     out_tracers = map(trace.full_raise, ans)
   1190     jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    164 
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:
    168       # Some transformations yield from inside context managers, so we have to

google3/third_party/py/jax/_src/random.py in _shuffle(key, x, axis)
    564   exponent = 3  # see tjablin@'s analysis for explanation of this parameter
    565   uint32max = jnp.iinfo(np.uint32).max
--> 566   num_rounds = int(np.ceil(exponent * np.log(x.size) / np.log(uint32max)))
    567 
    568   for _ in range(num_rounds):

OverflowError: cannot convert float infinity to integer
```


Changed files:
jax/_src/random.py, tests/random_test.py

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 3cccfcbaa..668f6c51f 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -563,7 +563,7 @@ def _shuffle(key, x, axis) -> jnp.ndarray:
   # another analysis (where the keys are generated one bit at a time).
   exponent = 3  # see tjablin@'s analysis for explanation of this parameter
   uint32max = jnp.iinfo(np.uint32).max
-  num_rounds = int(np.ceil(exponent * np.log(x.size) / np.log(uint32max)))
+  num_rounds = int(np.ceil(exponent * np.log(max(1, x.size)) / np.log(uint32max)))
 
   for _ in range(num_rounds):
     key, subkey = split(key)
diff --git a/tests/random_test.py b/tests/random_test.py
index 447940f17..13f3f8576 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -296,7 +296,7 @@ class LaxRandomTest(jtu.JaxTestCase):
       {"testcase_name": "_{}".format(jtu.format_shape_dtype_string(shape, dtype)),
        "dtype": dtype, "shape": shape}
       for dtype in jtu.dtypes.floating + jtu.dtypes.integer
-      for shape in [100, (10, 10), (10, 5, 2)]))
+      for shape in [100, (10, 10), (10, 5, 2), 0, 1, (0, 5), (1, 5)]))
   def testPermutationArray(self, dtype, shape):
     key = random.PRNGKey(0)
     x = jnp.arange(np.prod(shape)).reshape(shape).astype(dtype)
@@ -307,7 +307,8 @@ class LaxRandomTest(jtu.JaxTestCase):
     perm2 = crand(key)
 
     self.assertAllClose(perm1, perm2)
-    self.assertFalse(np.all(perm1 == x))  # seems unlikely!
+    if x.shape[0] > 1:
+      self.assertFalse(np.all(perm1 == x))  # seems unlikely!
     self.assertAllClose(np.sort(perm1.ravel()), x.ravel(), check_dtypes=False)
     self.assertArraysAllClose(
       x, jnp.arange(np.prod(shape)).reshape(shape).astype(dtype))