Repository: google/jax

Issue:
RBG-based PRNG doesn't work with bf16
bf16 is the norm for TPU training, and RBG-based PRNG is useful to speed up random number generation for large models. It seems that our RBG-based PRNG doesn't yet support bf16. Would be nice to enable this.

See repro

```python
import jax
import jax.numpy as jnp
 
print('-- default software PRNG')
print(jax.random.uniform(jax.random.PRNGKey(1), (2, 2), jnp.bfloat16))
 
print('-- RBG-based PRNG')
 
from jax import prng
jax.config.update('jax_enable_custom_prng', True)
 
def rbg_key(seed: int):
  return prng.seed_with_impl(prng.rbg_prng_impl, seed)
 
jax.random.PRNGKey = rbg_key

print(jax.random.uniform(jax.random.PRNGKey(1), (2, 2), jnp.bfloat16))
```

Default PRNG works
```
-- default software PRNG
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[[0.53125 0.453125]
 [0.5 0.984375]]
```

RBG-based says unsupported

```
-- RBG-based PRNG
Traceback (most recent call last):
  File "/usr/local/google/home/zhangqiaorjc/Downloads/t.py", line 16, in <module>
    print(jax.random.uniform(jax.random.PRNGKey(1), (2, 2), jnp.bfloat16))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/random.py", line 239, in uniform
    return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/random.py", line 258, in _uniform
    bits = _random_bits(key, nbits, shape)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/random.py", line 88, in _random_bits
    return key._random_bits(bit_width, shape)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/prng.py", line 184, in _random_bits
    return self.impl.random_bits(self._keys, bit_width, shape)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/prng.py", line 540, in _rbg_random_bits
    _, bits = lax.rng_bit_generator(key, shape, dtype=UINT_DTYPES[bit_width])
jax._src.source_info_util.JaxStackTraceBeforeTransformation: RuntimeError: INVALID_ARGUMENT: Unsupported shape for RngBitGenerator: U16

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/google/home/zhangqiaorjc/Downloads/t.py", line 16, in <module>
    print(jax.random.uniform(jax.random.PRNGKey(1), (2, 2), jnp.bfloat16))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/random.py", line 239, in uniform
    return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/api.py", line 429, in cache_miss
    donated_invars=donated_invars, inline=inline)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 1671, in bind
    return call_bind(self, fun, *args, **params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 1683, in call_bind
    outs = top_trace.process_call(primitive, fun, tracers, params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/core.py", line 596, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dispatch.py", line 143, in _xla_call_impl
    *unsafe_map(arg_spec, args))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/linear_util.py", line 272, in memoized_fun
    ans = call(fun, *args)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dispatch.py", line 170, in _xla_callable_uncached
    *arg_specs).compile().unsafe_call
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/dispatch.py", line 265, in lower_xla_callable
    arg_partitions=None, out_partitions=None)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/interpreters/xla.py", line 772, in lower_jaxpr_to_xla_module
    out_nodes = jaxpr_subcomp(ctx, jaxpr.jaxpr, xla_consts, *xla_args)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/interpreters/xla.py", line 600, in jaxpr_subcomp
    *in_nodes, **eqn.params)
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/lax/lax.py", line 4160, in _rng_bit_generator_translation_rule
    c, xops.RngBitGenerator(algorithm, key, xla_shape))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/interpreters/xla.py", line 611, in xla_destructure
    num_elements = len(c.get_shape(ans).tuple_shapes())
jax._src.traceback_util.UnfilteredStackTrace: RuntimeError: INVALID_ARGUMENT: Unsupported shape for RngBitGenerator: U16

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/google/home/zhangqiaorjc/Downloads/t.py", line 16, in <module>
    print(jax.random.uniform(jax.random.PRNGKey(1), (2, 2), jnp.bfloat16))
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/_src/random.py", line 239, in uniform
    return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
  File "/usr/local/google/home/zhangqiaorjc/repos/jax/jax/interpreters/xla.py", line 611, in xla_destructure
    num_elements = len(c.get_shape(ans).tuple_shapes())
RuntimeError: INVALID_ARGUMENT: Unsupported shape for RngBitGenerator: U16```


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4150,9 +4150,17 @@ def _rng_bit_generator_translation_rule(
   # sidestep issues with the jax_enable_x64=False configuration. As a result, we
   # need to convert u32[4] -> u64[2] here in the translation rule. However, we
   # also polymorphically allow a u64[2] for backward compatibility.
+  #
+  # Separately, xops.RngBitGenerator doesn't support generating u8 or
+  # u16, so we request u32 and truncate in that case.
   assert ((key_shape == (4,) and key_dtype == np.dtype('uint32')) or
           (key_shape == (2,) and key_dtype == np.dtype('uint64'))), (key_shape, key_dtype)
-  xla_shape = xc.Shape.array_shape(np.dtype(dtype), shape)
+  dtype = np.dtype(dtype)
+  if dtype == np.dtype('uint32') or dtype == np.dtype('uint64'):
+    rbg_dtype = dtype
+  else:
+    rbg_dtype = np.dtype('uint32')
+  xla_shape = xc.Shape.array_shape(rbg_dtype, shape)
   if key_dtype == np.dtype('uint32'):
     u64_etype = xla.dtype_to_primitive_type(np.dtype('uint64'))
     key = xops.BitcastConvertType(xops.Reshape(key, (2, 2)), u64_etype)
@@ -4161,6 +4169,9 @@ def _rng_bit_generator_translation_rule(
   if key_dtype == np.dtype('uint32'):
     u32_etype = xla.dtype_to_primitive_type(np.dtype('uint32'))
     out_key = xops.Reshape(xops.BitcastConvertType(out_key, u32_etype), (4,))
+  if rbg_dtype != dtype:
+    out_vals = xops.ConvertElementType(
+      out_vals, xla.dtype_to_primitive_type(dtype))
   return [out_key, out_vals]
 
 
@@ -4199,6 +4210,10 @@ def rng_bit_generator(key, shape, dtype=np.uint32,
   friendly API.
   """
   shape = jax.core.canonicalize_shape(shape)
+  dtype = dtypes.canonicalize_dtype(dtype)
+  if np.dtype(dtype) not in {np.dtype('uint8'), np.dtype('uint16'),
+                             np.dtype('uint32'), np.dtype('uint64')}:
+    raise TypeError(f'rng_bit_generator: unsupported dtype {dtype}')
   return tuple(
       rng_bit_generator_p.bind(
           key, shape=shape, dtype=dtype, algorithm=algorithm))


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2501,6 +2501,18 @@ def fn(k):
     self.assertArraysEqual(out[0], out_jit[0])
     self.assertArraysEqual(out[1], out_jit[1])
 
+  def testRngBitGenerator2(self):
+    def f(key):
+      return lax.rng_bit_generator(key, shape=(5, 7))
+
+    key = np.array((1, 2, 3, 4)).astype(np.uint32)
+    out1 = f(key)
+    out2 = jax.jit(f)(key)
+    self.assertEqual(out1[0].shape, (4,))
+    self.assertEqual(out1[1].shape, (5, 7))
+    self.assertArraysEqual(out1[0], out2[0])
+    self.assertArraysEqual(out1[1], out2[1])
+
   @jtu.skip_on_devices("tpu")
   def testRngBitGeneratorReturnedKey(self):
     # This test ensures that the key bit-packing/unpacking operations used in
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -51,6 +51,11 @@ def _prng_key_as_array(key):
   return key.unsafe_raw_array() if config.jax_enable_custom_prng else key
 
 
+PRNG_IMPLS = [('threefry2x32', prng.threefry_prng_impl),
+              ('rbg', prng.rbg_prng_impl),
+              ('unsafe_rbg', prng.unsafe_rbg_prng_impl)]
+
+
 @jtu.with_config(jax_numpy_rank_promotion="raise")
 class PrngTest(jtu.JaxTestCase):
 
@@ -128,6 +133,35 @@ def testRngRandomBits(self):
       expected64 = np.array([676898860, 3164047411, 4010691890], dtype=np.uint32)
     self.assertArraysEqual(bits64, expected64)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_" + name, "prng_name": name}
+      for name, _ in PRNG_IMPLS))
+  def testRngRandomBitsShapeDtype(self, prng_name):
+    # Like testRngRandomBits, but only meant to exercise random_bits
+    # on every PRNG implementation. Instead of values, only checks
+    # that shapes/dtypes are as expected.
+
+    with jax.default_prng_impl(prng_name):
+      key = random.PRNGKey(1701)
+
+      bits8 = jax._src.random._random_bits(key, 8, (3,))
+      self.assertEqual(bits8.shape, (3,))
+      self.assertEqual(bits8.dtype, np.dtype('uint8'))
+
+      bits16 = jax._src.random._random_bits(key, 16, (3,))
+      self.assertEqual(bits16.shape, (3,))
+      self.assertEqual(bits16.dtype, np.dtype('uint16'))
+
+      bits32 = jax._src.random._random_bits(key, 32, (3,))
+      self.assertEqual(bits32.shape, (3,))
+      self.assertEqual(bits32.dtype, np.dtype('uint32'))
+
+      with jtu.ignore_warning(category=UserWarning, message="Explicitly requested dtype.*"):
+        bits64 = jax._src.random._random_bits(key, 64, (3,))
+      expected_dtype = np.dtype('uint64' if config.x64_enabled else 'uint32')
+      self.assertEqual(bits64.shape, (3,))
+      self.assertEqual(bits64.dtype, expected_dtype)
+
   def testRngRandomBitsViewProperty(self):
     # TODO: add 64-bit if it ever supports this property.
     # TODO: will this property hold across endian-ness?
@@ -209,9 +243,7 @@ def test_prng_seeds_and_keys(self, seed, type, jit, key):
   def test_default_prng_selection(self):
     if not config.jax_enable_custom_prng:
       self.skipTest("test requires config.jax_enable_custom_prng")
-    for name, impl in [('threefry2x32', prng.threefry_prng_impl),
-                       ('rbg', prng.rbg_prng_impl),
-                       ('unsafe_rbg', prng.unsafe_rbg_prng_impl)]:
+    for name, impl in PRNG_IMPLS:
       with jax.default_prng_impl(name):
         self.assertIs(random.default_prng_impl(), impl)
         key = random.PRNGKey(42)
@@ -223,9 +255,7 @@ def test_default_prng_selection(self):
   def test_default_prng_selection_without_custom_prng_mode(self):
     if config.jax_enable_custom_prng:
       self.skipTest("test requires that config.jax_enable_custom_prng is False")
-    for name, impl in [('threefry2x32', prng.threefry_prng_impl),
-                       ('rbg', prng.rbg_prng_impl),
-                       ('unsafe_rbg', prng.unsafe_rbg_prng_impl)]:
+    for name, impl in PRNG_IMPLS:
       with jax.default_prng_impl(name):
         self.assertIs(random.default_prng_impl(), impl)
         key = random.PRNGKey(42)
@@ -1334,15 +1364,6 @@ def test_ravel(self):
     keys = jnp.ravel(keys)
     self.assertEqual(keys.shape, (4,))
 
-def _sampler_unimplemented_with_rbg(*args, **kwargs):
-  # TODO(mattjj): enable these tests if/when RngBitGenerator supports them
-  raise SkipTest('8- and 16-bit types not supported with RBG PRNG')
-
-for attr in dir(LaxRandomWithRBGPRNGTest):
-  if 'int8' in attr or 'int16' in attr or 'float16' in attr:
-    setattr(LaxRandomWithRBGPRNGTest, attr, _sampler_unimplemented_with_rbg)
-    setattr(LaxRandomWithUnsafeRBGPRNGTest, attr, _sampler_unimplemented_with_rbg)
-
 def _sampler_unimplemented_with_custom_prng(*args, **kwargs):
   raise SkipTest('sampler only implemented for default RNG')