# Buggy code
def _threefry2x32_gpu_translation_rule(c, k1, k2, x1, x2):
    shape = lax.broadcast_shapes(
        c.get_shape(k1).dimensions(),
        c.get_shape(k2).dimensions(),
        c.get_shape(x1).dimensions(),
        c.get_shape(x2).dimensions(),
    )
    rank = len(shape)

    def _broadcast(x):
        ndims = c.get_shape(x).rank()
        return xla_client.ops.BroadcastInDim(x, shape, tuple(range(rank - ndims, rank)))

    return cuda_prng.threefry2x32(
        c, (_broadcast(k1), _broadcast(k2)), (_broadcast(x1), _broadcast(x2))
    )


# Traceback (RuntimeError)
---------------------------------------------------------------------------
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

# Fixed code
def _threefry2x32_gpu_translation_rule(c, k1, k2, x1, x2):
    shape = lax.broadcast_shapes(
        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, tuple(range(rank - ndims, rank)))

    return cuda_prng.threefry2x32(
        c, (_broadcast(k1), _broadcast(k2)), (_broadcast(x1), _broadcast(x2))
    )