Repository: google/jax
Issue: random.uniform of 0 sized arrays fails on GPU when JIT is disabled

When JIT is enabled, calling `jax.random.uniform` with a shape that has one or more 0-sized dimensions returns an empty array of the requested dimensions.

When JIT is disabled, calling `jax.random.uniform` with the same shape fails with a "CUDA operation failed" error.

Simple repro:
```py
import jax

key = jax.numpy.array([1, 1], jax.numpy.uint32)
shape = (1, 0)

jax.random.uniform(key, shape)    # ok

with jax.disable_jit():
  jax.random.uniform(key, shape)  # CUDA operation failed: invalid configuration argument
```

Full callstack:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-2-1d1849d1c41c> in <module>()
      5 jax.random.uniform(key, shape)    # ok
      6 with jax.disable_jit():
----> 7   jax.random.uniform(key, shape)  # CUDA operation failed: invalid configuration argument

9 frames
google3/third_party/py/jax/random.py in uniform(key, shape, dtype, minval, maxval)
    380   dtype = dtypes.canonicalize_dtype(dtype)
    381   shape = abstract_arrays.canonicalize_shape(shape)
--> 382   return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
    383 
    384 @partial(jit, static_argnums=(1, 2))

google3/third_party/py/jax/api.py in f_jitted(*args, **kwargs)
    190   def f_jitted(*args, **kwargs):
    191     if _jit_is_disabled():
--> 192       return fun(*args, **kwargs)
    193     if max(static_argnums + donate_argnums, default=-1) >= len(args):
    194       msg = ("jitted function has static_argnums={}, donate_argnums={} but "

google3/third_party/py/jax/random.py in _uniform(key, shape, dtype, minval, maxval)
    399     raise TypeError("uniform only accepts 32- or 64-bit dtypes.")
    400 
--> 401   bits = _random_bits(key, nbits, shape)
    402 
    403   # The strategy here is to randomize only the mantissa bits with an exponent of

google3/third_party/py/jax/random.py in _random_bits(key, bit_width, shape)
    312   nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max)
    313   if not nblocks:
--> 314     bits = threefry_2x32(key, lax.iota(np.uint32, rem))
    315   else:
    316     *subkeys, last_key = split(key, nblocks + 1)

google3/third_party/py/jax/api.py in f_jitted(*args, **kwargs)
    190   def f_jitted(*args, **kwargs):
    191     if _jit_is_disabled():
--> 192       return fun(*args, **kwargs)
    193     if max(static_argnums + donate_argnums, default=-1) >= len(args):
    194       msg = ("jitted function has static_argnums={}, donate_argnums={} but "

google3/third_party/py/jax/random.py in threefry_2x32(keypair, count)
    258     x = list(jnp.split(count.ravel(), 2))
    259 
--> 260   x = threefry2x32_p.bind(key1, key2, x[0], x[1])
    261   out = jnp.concatenate(x)
    262   assert out.dtype == np.uint32

google3/third_party/py/jax/core.py in bind(self, *args, **params)
    261     top_trace = find_top_trace(args)
    262     tracers = map(top_trace.full_raise, args)
--> 263     out = top_trace.process_primitive(self, tracers, params)
    264     return map(full_lower, out) if self.multiple_results else full_lower(out)
    265 

google3/third_party/py/jax/core.py in process_primitive(self, primitive, tracers, params)
    570 
    571   def process_primitive(self, primitive, tracers, params):
--> 572     return primitive.impl(*tracers, **params)
    573 
    574   def process_call(self, primitive, f, tracers, params):

google3/third_party/py/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    232   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
    233   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
--> 234   return compiled_fun(*args)
    235 
    236 

google3/third_party/py/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, result_handler, *args)
    347   device, = compiled.local_devices()
    348   input_bufs = list(it.chain.from_iterable(device_put(x, device) for x in args if x is not token))
--> 349   out_bufs = compiled.execute(input_bufs)
    350   if FLAGS.jax_debug_nans: check_nans(prim, out_bufs)
    351   return result_handler(*out_bufs)

RuntimeError: CUDA operation failed: invalid configuration argument
```

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

diff --git a/jax/random.py b/jax/random.py
index 044261ee8..7571bf61d 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -213,6 +213,10 @@ def _threefry2x32_gpu_translation_rule(c, k1, k2, x1, x2):
       c.get_shape(k1).dimensions(), c.get_shape(k2).dimensions(),
       c.get_shape(x1).dimensions(), c.get_shape(x2).dimensions())
   rank = len(shape)
+  if 0 in shape:
+    zeros = xla_client.ops.Broadcast(
+        xla_bridge.constant(c, np.array(0, np.uint32)), shape)
+    return xla_client.ops.Tuple(c, [zeros, zeros])
   def _broadcast(x):
     ndims = c.get_shape(x).rank()
     return xla_client.ops.BroadcastInDim(x, shape,
diff --git a/tests/random_test.py b/tests/random_test.py
index 9c0d8581e..0ba2c3228 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -120,6 +120,14 @@ class LaxRandomTest(jtu.JaxTestCase):
     np.testing.assert_equal(result[:n], np.full((n,), 0xc4923a9c, dtype=np.uint32))
     np.testing.assert_equal(result[n:], np.full((n,), 0x483df7a0, dtype=np.uint32))
 
+  def testThreefry2x32Empty(self):
+    # Regression test for an op-by-op crash for empty arrays in CUDA mode.
+    with api.disable_jit():
+      result = random.threefry_2x32(
+        (np.uint32(0x13198a2e), np.uint32(0x03707344)),
+        jnp.ones((10, 0,), jnp.uint32))
+    np.testing.assert_equal(result, np.zeros((10, 0,), dtype=np.uint32))
+
   def testRngRandomBitsViewProperty(self):
     # TODO: add 64-bit if it ever supports this property.
     # TODO: will this property hold across endian-ness?