Repository: jax-ml/jax
Pull request: support custom PRNG implementations
Today jax offers a single splittable PRNG implementation (threefry 2x32) and models its key as a `uint32[2]`. A batch of keys is a `uint32` array whose trailing dimension comprises the length-2 keys.

With this PR, we have two goals:

1. To encapsulate the "raw array" key representation, in order to prevent unsafe or invalid key operations.
2. To enable custom PRNG implementations with custom key representations. We want every function in the `jax.random` to work even when seeding, splitting, and random bit generation is defined by user code.

An added challenge of this upgrade is that the previous "raw array" key model has been around for a long time, and many downstream uses have overfit to it. To varying extents, flax, haiku, trax, and TFP all validate user-provided random keys as `uint32[2]` arrays. We'd like to obviate the need for such checks. Other downstream code creates keys directly with hard-coded values (including a few of our own tests), reshapes key batches (including our `gamma` primitive implementation), or casts other program values to keys. For safety, we would like to disallow such operations.

I'm introducing this change gradually to allow time for downstream fixes. I've tried to do this by recovering existing behavior under a flag that is, for now, on by default.

## design

Again, two goals:

1. to abstract away key representations, and
2. to respect custom PRNGs throughout `jax.random`.

Item 1 is scoped only to the JAX Python level; changing key representation in the Jaxpr IR is not a goal. This PR requires that keys remain backed by `uint32` arrays, now of arbitrary specified shape. We can broaden that later as needed. They will therefore appear as `uint32` arrays in Jaxpr.

A custom PRNG implementation amounts to a key type `K` (determined by a shape) plus a handful of functions operating on such a key:

```haskell
class PRNGImpl where:
  key_shape :: (int, ...)
  seed :: int[] -> K
  fold_in :: K -> int[] -> K
  split[n] :: K -> K[n]
  random_bits[shape, bit_width] :: K -> uint<bit_width>[shape]
```

An example, also used for unit tests introduced in this PR, is the following cartoon PRNG built from two threefry2x32 keys. To define it:

```python
import jax
from jax import prng, numpy as jnp

threefry_seed = jax._src.prng.threefry_seed
threefry_split = jax._src.prng.threefry_split
threefry_random_bits = jax._src.prng.threefry_random_bits
threefry_fold_in = jax._src.prng.threefry_fold_in

def _double_threefry_seed(seed):
  return jnp.vstack([threefry_seed(seed + 1),
                     threefry_seed(seed + 2)])

def _double_threefry_split(key, num):
  split0 = threefry_split(key[0], num)
  split1 = threefry_split(key[1], num)
  merge = jnp.vstack([jnp.expand_dims(split0.T, axis=0),
                      jnp.expand_dims(split1.T, axis=0)])
  return merge.transpose((2, 0, 1))

def _double_threefry_random_bits(key, bit_width, shape):
  bits0 = threefry_random_bits(key[0], bit_width, shape)
  bits1 = threefry_random_bits(key[1], bit_width, shape)
  return bits0 * bits1

def _double_threefry_fold_in(key, data):
  return jnp.vstack([threefry_fold_in(key[0], data),
                     threefry_fold_in(key[1], data)])

double_threefry_prng_impl = prng.PRNGImpl(
    key_shape=(2, 2),
    seed=_double_threefry_seed,
    split=_double_threefry_split,
    random_bits=_double_threefry_random_bits,
    fold_in=_double_threefry_fold_in)
```

And to use it:

```python
>>> seed_value = 73
>>> key = prng.seed_with_impl(double_threefry_prng_impl, seed_value)
>>> key
PRNGKeyArray:
  shape = ()
  impl = PRNGImpl:
           key_shape = (2, 2)
           seed = <function _double_threefry_seed at 0x7f465e44bee0>
           split = <function _double_threefry_split at 0x7f465e44bb80>
           random_bits = <function _double_threefry_random_bits at 0x7f465e44ba60>
           fold_in = <function _double_threefry_fold_in at 0x7f46227e1160>

>>> key, subkey = jax.random.split(key)
>>> jax.random.uniform(subkey)
DeviceArray(0.48681986, dtype=float32)

>>> jax.random.split(key, 10)
PRNGKeyArray:
  shape = (10,)
  impl = PRNGImpl:
           key_shape = (2, 2)
           seed = <function _double_threefry_seed at 0x7f465e44bee0>
           split = <function _double_threefry_split at 0x7f465e44bb80>
           random_bits = <function _double_threefry_random_bits at 0x7f465e44ba60>
           fold_in = <function _double_threefry_fold_in at 0x7f46227e1160>

>>> jax.vmap(lambda i: jax.random.fold_in(key, i))(jnp.arange(3))
PRNGKeyArray:
  shape = (3,)
  impl = PRNGImpl:
           key_shape = (2, 2)
           seed = <function _double_threefry_seed at 0x7f465e44bee0>
           split = <function _double_threefry_split at 0x7f465e44bb80>
           random_bits = <function _double_threefry_random_bits at 0x7f465e44ba60>
           fold_in = <function _double_threefry_fold_in at 0x7f46227e1160>

>>> key + 47
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for +: 'PRNGKeyArray' and 'int'
```

Note that the `shape` field of the `PRNGKeyArray` changes from the scalar `()` to `(10,)` when it is `split` 10 ways and to `(3,)` when it is broadcast by `vmap`.

Internally, this works as follows:

* Users define a PRNG implementation `impl` with a key type `K` like the one above.
* An internal pytree-node class `PRNGKeyArray` wraps the given PRNG implementation `impl` and lifts it to an array-like interface. This has roughly the semantics of an array of `K`-typed elements. Under the hood it is a `uint32` array whose trailing dimensions correspond to `K`. It is array-like only in select ways, for instance supporting indexing and iteration but not numerical operations. Expressions such as `key, subkey = jax.random.split(key)` remain valid but `key + 1` or `jnp.exp(key)` are no longer allowed.
* The `PRNGKeyArray` carries the implementation `impl` around through `jax.random` functions, which then delegate to it for `seed`, `split`, `random_bits`, and `fold_in`.

Note that all new indirection happens at the Python level. This design makes no changes to Jaxpr, although the extra indirection could hopefully help make element-type changes in Jaxpr easier in the future. I discussed this with @LenaMartens and we think it will help clear the path for some Jaxpr-level RNG checking ideas in particular.

I've separately started thinking about whether to generalize the array-like wrapping concept to allow for lightweight custom-element-type arrays more broadly. If RNG keys are one special case, then arrays of `unit`, `float0`, enums, or arbitrary pytrees could be others. I mean "lightweight" both in the sense of "easy to set up" as well as "without necessarily extending Jaxpr."

## temporary limitations

* Both `random.gamma`, and `jax.poisson` are not yet implemented with custom PRNGs, since these have threefry2x32-specific implementations today that need to be generalized. I have them `raise NotImplementedError` when given a custom PRNG for now.

* When fully enabled, this change breaks several of our neighboring libraries. To allow for smoother upgrades, I've retained existing behavior under a flag. Doing so introduces extra code and test complexity in the interim.

* `PRNGKeyArray` may want to support more array-like methods in the future. I implemented an obvious initial handful for now.

----
Fixes #5081.

Relatedly addresses parts of #2294.

Problem statement:
Allow pluggable PRNG implementations
Currently we use ThreeFry2x32 as our PRNG, but it has an uncomfortably small key space (2^64).

An easy way to improve this would be to use the 4x32 variant of ThreeFry instead which has a key space of size 2^128.

Hints from issue discussion:
This sounds fun...

For backward compatibility, do you think it might make sense to implement this beside the current approach (e.g. with `jax.random.PRNGKey4()` or something similar?)
Or maybe better, a `keysize` keyword to `PRNGKey` that defaults to 2, and can optionally be set to 4, with the intention of eventually changing the default?
I was thinking that perhaps we would do something like `PRNGKey(..., algorithm=THREEFRY4x32)` to allow for the possiblity we might want to support other PRNGs.
https://github.com/google/jax/issues/2294 might be related?  It would allow you to make this change more transparently since no code would assume a key size; it would be hidden behind the opaque object.
Generalizing this issue: we should allow pluggable PRNG bit generators (e.g., Philox on GPU).

Commit message: factor PRNG routines from random module to prng

Patch to review:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
new file mode 100644
index 000000000000..0b3bcfc19920
--- /dev/null
+++ b/jax/_src/prng.py
@@ -0,0 +1,341 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import partial
+
+import numpy as np
+
+from jax import lax
+from jax import core
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src.api import jit, vmap
+from jax.lib import xla_bridge
+from jax.lib import xla_client
+from jax.lib import cuda_prng
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax._src.util import prod
+
+
+UINT_DTYPES = {
+    8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}  # type: ignore[has-type]
+
+
+def PRNGKey(seed: int) -> jnp.ndarray:
+  """Create a pseudo-random number generator (PRNG) key given an integer seed.
+
+  Args:
+    seed: a 64- or 32-bit integer used as the value of the key.
+
+  Returns:
+    A PRNG key, which is modeled as an array of shape (2,) and dtype uint32. The
+    key is constructed from a 64-bit seed by effectively bit-casting to a pair
+    of uint32 values (or from a 32-bit seed by first padding out with zeros).
+  """
+  # Avoid overflowerror in X32 mode by first converting ints to int64.
+  # This breaks JIT invariance of PRNGKey for large ints, but supports the
+  # common use-case of instantiating PRNGKey with Python hashes in X32 mode.
+  if isinstance(seed, int):
+    seed_arr = jnp.asarray(np.int64(seed))
+  else:
+    seed_arr = jnp.asarray(seed)
+  if seed_arr.shape:
+    raise TypeError(f"PRNGKey seed must be a scalar; got {seed!r}.")
+  if not np.issubdtype(seed_arr.dtype, np.integer):
+    raise TypeError(f"PRNGKey seed must be an integer; got {seed!r}")
+
+  convert = lambda k: lax.reshape(lax.convert_element_type(k, np.uint32), [1])
+  k1 = convert(lax.shift_right_logical(seed_arr, lax._const(seed_arr, 32)))
+  k2 = convert(jnp.bitwise_and(seed_arr, np.uint32(0xFFFFFFFF)))
+  return lax.concatenate([k1, k2], 0)
+
+
+def _is_prng_key(key: jnp.ndarray) -> bool:
+  try:
+    return key.shape == (2,) and key.dtype == np.uint32
+  except AttributeError:
+    return False
+
+
+def _make_rotate_left(dtype):
+  if not jnp.issubdtype(dtype, np.integer):
+    raise TypeError("_rotate_left only accepts integer dtypes.")
+  nbits = np.array(jnp.iinfo(dtype).bits, dtype)
+
+  def _rotate_left(x, d):
+    if lax.dtype(d) != dtype:
+      d = lax.convert_element_type(d, dtype)
+    if lax.dtype(x) != dtype:
+      x = lax.convert_element_type(x, dtype)
+    return lax.shift_left(x, d) | lax.shift_right_logical(x, nbits - d)
+  return _rotate_left
+
+
+def _bit_stats(bits):
+  """This is a debugging function to compute the statistics of bit fields."""
+  return np.array([list(map(int, np.binary_repr(x, 64))) for x in bits]).mean(0)
+
+
+### hash function and split
+
+def _threefry2x32_abstract_eval(*args):
+  if any(a.dtype != jnp.uint32 for a in args):
+    raise TypeError("Arguments to threefry2x32 must have uint32 type, got {}"
+                    .format(args))
+  if all(isinstance(arg, core.ShapedArray) for arg in args):
+    shape = lax._broadcasting_shape_rule(*args)
+    named_shape = core.join_named_shapes(*(a.named_shape for a in args))
+    aval = core.ShapedArray(shape, jnp.dtype(jnp.uint32), named_shape=named_shape)
+  else:
+    aval = core.UnshapedArray(jnp.dtype(jnp.uint32))
+  return (aval,) * 2
+
+
+rotate_left = _make_rotate_left(np.uint32)
+
+
+def apply_round(v, rot):
+  v = v[:]
+  v[0] = v[0] + v[1]
+  v[1] = rotate_left(v[1], rot)
+  v[1] = v[0] ^ v[1]
+  return v
+
+
+def rotate_list(xs):
+  return xs[1:] + xs[:1]
+
+
+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] + 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):
+  """Apply the Threefry 2x32 hash.
+
+  Args:
+    keypair: a pair of 32bit unsigned integers used for the key.
+    count: an array of dtype uint32 used for the counts.
+
+  Returns:
+    An array of dtype uint32 with the same shape as `count`.
+  """
+  x = [x1, x2]
+
+  rotations = [np.array([13, 15, 26, 6], dtype=np.uint32),
+               np.array([17, 29, 16, 24], dtype=np.uint32)]
+  ks = [key1, key2, key1 ^ key2 ^ np.uint32(0x1BD11BDA)]
+
+  x[0] = x[0] + ks[0]
+  x[1] = x[1] + ks[1]
+
+  if use_rolled_loops:
+    x, _, _ = lax.fori_loop(0, 5, rolled_loop_step, (x, rotate_list(ks), rotations))
+
+  else:
+    for r in rotations[0]:
+      x = apply_round(x, r)
+    x[0] = x[0] + ks[1]
+    x[1] = x[1] + ks[2] + np.uint32(1)
+
+    for r in rotations[1]:
+      x = apply_round(x, r)
+    x[0] = x[0] + ks[2]
+    x[1] = x[1] + ks[0] + np.uint32(2)
+
+    for r in rotations[0]:
+      x = apply_round(x, r)
+    x[0] = x[0] + ks[0]
+    x[1] = x[1] + ks[1] + np.uint32(3)
+
+    for r in rotations[1]:
+      x = apply_round(x, r)
+    x[0] = x[0] + ks[1]
+    x[1] = x[1] + ks[2] + np.uint32(4)
+
+    for r in rotations[0]:
+      x = apply_round(x, r)
+    x[0] = x[0] + ks[2]
+    x[1] = x[1] + ks[0] + np.uint32(5)
+
+  return tuple(x)
+
+
+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)))
+
+
+threefry2x32_p = core.Primitive("threefry2x32")
+threefry2x32_p.multiple_results = True
+threefry2x32_p.def_impl(partial(xla.apply_primitive, threefry2x32_p))
+threefry2x32_p.def_abstract_eval(_threefry2x32_abstract_eval)
+batching.defbroadcasting(threefry2x32_p)
+xla.translations_with_avals[threefry2x32_p] = xla.lower_fun(
+    partial(_threefry2x32_lowering, use_rolled_loops=False),
+    multiple_results=True, with_avals=True)
+xla.backend_specific_translations['cpu'][threefry2x32_p] = xla.lower_fun(
+    partial(_threefry2x32_lowering, use_rolled_loops=True),
+    multiple_results=True)
+if cuda_prng:
+  xla.backend_specific_translations['gpu'][threefry2x32_p] = \
+      _threefry2x32_gpu_translation_rule
+
+
+@jit
+def threefry_2x32(keypair, count):
+  """Apply the Threefry 2x32 hash.
+
+  Args:
+    keypair: a pair of 32bit unsigned integers used for the key.
+    count: an array of dtype uint32 used for the counts.
+
+  Returns:
+    An array of dtype uint32 with the same shape as `count`.
+  """
+  key1, key2 = keypair
+  if not lax.dtype(key1) == lax.dtype(key2) == lax.dtype(count) == np.uint32:
+    msg = "threefry_2x32 requires uint32 arguments, got {}"
+    raise TypeError(msg.format([lax.dtype(x) for x in [key1, key2, count]]))
+
+  try:
+    odd_size = count.size % 2
+  except core.InconclusiveDimensionOperation as e:
+    msg = ("jax.random functions have limited support for shape polymorphism. "
+           "In particular, the product of the known dimensions must be even.")
+    raise core.InconclusiveDimensionOperation(msg) from e
+
+  if odd_size:
+    x = list(jnp.split(jnp.concatenate([count.ravel(), np.uint32([0])]), 2))
+  else:
+    x = list(jnp.split(count.ravel(), 2))
+
+  x = threefry2x32_p.bind(key1, key2, x[0], x[1])
+  out = jnp.concatenate(x)
+  assert out.dtype == np.uint32
+  return lax.reshape(out[:-1] if odd_size else out, count.shape)
+
+
+def split(key: jnp.ndarray, num: int = 2) -> jnp.ndarray:
+  """Splits a PRNG key into `num` new keys by adding a leading axis.
+
+  Args:
+    key: a PRNGKey (an array with shape (2,) and dtype uint32).
+    num: optional, a positive integer indicating the number of keys to produce
+      (default 2).
+
+  Returns:
+    An array with shape (num, 2) and dtype uint32 representing `num` new keys.
+  """
+  return _split(key, int(num))  # type: ignore
+
+
+@partial(jit, static_argnums=(1,))
+def _split(key, num) -> jnp.ndarray:
+  counts = lax.iota(np.uint32, num * 2)
+  return lax.reshape(threefry_2x32(key, counts), (num, 2))
+
+
+def fold_in(key: jnp.ndarray, data: int) -> jnp.ndarray:
+  """Folds in data to a PRNG key to form a new PRNG key.
+
+  Args:
+    key: a PRNGKey (an array with shape (2,) and dtype uint32).
+    data: a 32bit integer representing data to be folded in to the key.
+
+  Returns:
+    A new PRNGKey that is a deterministic function of the inputs and is
+    statistically safe for producing a stream of new pseudo-random values.
+  """
+  return _fold_in(key, jnp.uint32(data))
+
+
+@jit
+def _fold_in(key, data):
+  return threefry_2x32(key, PRNGKey(data))
+
+
+@partial(jit, static_argnums=(1, 2))
+def _random_bits(key, bit_width, shape):
+  """Sample uniform random bits of given width and shape using PRNG key."""
+  if not _is_prng_key(key):
+    raise TypeError("_random_bits got invalid prng key.")
+  if bit_width not in (8, 16, 32, 64):
+    raise TypeError("requires 8-, 16-, 32- or 64-bit field width.")
+  shape = core.as_named_shape(shape)
+  for name, size in shape.named_items:
+    real_size = lax.psum(1, name)
+    if real_size != size:
+      raise ValueError(f"The shape of axis {name} was specified as {size}, "
+                       f"but it really is {real_size}")
+    axis_index = lax.axis_index(name)
+    key = fold_in(key, axis_index)
+  size = prod(shape.positional)
+  # Compute ceil(bit_width * size / 32) in a way that is friendly to shape
+  # polymorphism
+  max_count, r = divmod(bit_width * size, 32)
+  if r > 0:
+    max_count += 1
+
+  if core.is_constant_dim(max_count):
+    nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max)
+  else:
+    nblocks, rem = 0, max_count
+
+  if not nblocks:
+    bits = threefry_2x32(key, lax.iota(np.uint32, rem))
+  else:
+    keys = split(key, nblocks + 1)
+    subkeys, last_key = keys[:-1], keys[-1]
+    blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
+    last = threefry_2x32(last_key, lax.iota(np.uint32, rem))
+    bits = lax.concatenate([blocks.ravel(), last], 0)
+
+  dtype = UINT_DTYPES[bit_width]
+  if bit_width == 64:
+    bits = [lax.convert_element_type(x, dtype) for x in jnp.split(bits, 2)]
+    bits = lax.shift_left(bits[0], dtype(32)) | bits[1]
+  elif bit_width in [8, 16]:
+    # this is essentially bits.view(dtype)[:size]
+    bits = lax.bitwise_and(
+      np.uint32(np.iinfo(dtype).max),
+      lax.shift_right_logical(
+        lax.broadcast(bits, (1,)),
+        lax.mul(
+          np.uint32(bit_width),
+          lax.broadcasted_iota(np.uint32, (32 // bit_width, 1), 0)
+        )
+      )
+    )
+    bits = lax.reshape(bits, (np.uint32(max_count * 32 // bit_width),), (1, 0))
+    bits = lax.convert_element_type(bits, dtype)[:size]
+  return lax.reshape(bits, shape)
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 6302a81d4c49..81ecd9f93841 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -26,9 +26,8 @@
 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, _check_arraylike
+from jax._src.prng import PRNGKey, fold_in, _random_bits, split, UINT_DTYPES
 from jax.lib import xla_bridge
-from jax.lib import xla_client
-from jax.lib import cuda_prng
 from jax.numpy.linalg import cholesky, svd, eigh
 from jax.interpreters import ad
 from jax.interpreters import batching
@@ -45,309 +44,6 @@
 DTypeLikeFloat = Any
 
 
-_UINT_DTYPES = {8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}
-
-
-def PRNGKey(seed: int) -> jnp.ndarray:
-  """Create a pseudo-random number generator (PRNG) key given an integer seed.
-
-  Args:
-    seed: a 64- or 32-bit integer used as the value of the key.
-
-  Returns:
-    A PRNG key, which is modeled as an array of shape (2,) and dtype uint32. The
-    key is constructed from a 64-bit seed by effectively bit-casting to a pair
-    of uint32 values (or from a 32-bit seed by first padding out with zeros).
-  """
-  # Avoid overflowerror in X32 mode by first converting ints to int64.
-  # This breaks JIT invariance of PRNGKey for large ints, but supports the
-  # common use-case of instantiating PRNGKey with Python hashes in X32 mode.
-  if isinstance(seed, int):
-    seed_arr = jnp.asarray(np.int64(seed))
-  else:
-    seed_arr = jnp.asarray(seed)
-  if seed_arr.shape:
-    raise TypeError(f"PRNGKey seed must be a scalar; got {seed!r}.")
-  if not np.issubdtype(seed_arr.dtype, np.integer):
-    raise TypeError(f"PRNGKey seed must be an integer; got {seed!r}")
-
-  convert = lambda k: lax.reshape(lax.convert_element_type(k, np.uint32), [1])
-  k1 = convert(lax.shift_right_logical(seed_arr, lax._const(seed_arr, 32)))
-  k2 = convert(jnp.bitwise_and(seed_arr, np.uint32(0xFFFFFFFF)))
-  return lax.concatenate([k1, k2], 0)
-
-def _is_prng_key(key: jnp.ndarray) -> bool:
-  try:
-    return key.shape == (2,) and key.dtype == np.uint32
-  except AttributeError:
-    return False
-
-
-### utilities
-
-
-def _make_rotate_left(dtype):
-  if not jnp.issubdtype(dtype, np.integer):
-    raise TypeError("_rotate_left only accepts integer dtypes.")
-  nbits = np.array(jnp.iinfo(dtype).bits, dtype)
-
-  def _rotate_left(x, d):
-    if lax.dtype(d) != dtype:
-      d = lax.convert_element_type(d, dtype)
-    if lax.dtype(x) != dtype:
-      x = lax.convert_element_type(x, dtype)
-    return lax.shift_left(x, d) | lax.shift_right_logical(x, nbits - d)
-  return _rotate_left
-
-
-def _bit_stats(bits):
-  """This is a debugging function to compute the statistics of bit fields."""
-  return np.array([list(map(int, np.binary_repr(x, 64))) for x in bits]).mean(0)
-
-
-### hash function and split
-
-def _threefry2x32_abstract_eval(*args):
-  if any(a.dtype != jnp.uint32 for a in args):
-    raise TypeError("Arguments to threefry2x32 must have uint32 type, got {}"
-                    .format(args))
-  if all(isinstance(arg, core.ShapedArray) for arg in args):
-    shape = lax._broadcasting_shape_rule(*args)
-    named_shape = core.join_named_shapes(*(a.named_shape for a in args))
-    aval = core.ShapedArray(shape, jnp.dtype(jnp.uint32), named_shape=named_shape)
-  else:
-    aval = core.UnshapedArray(jnp.dtype(jnp.uint32))
-  return (aval,) * 2
-
-rotate_left = _make_rotate_left(np.uint32)
-
-def apply_round(v, rot):
-  v = v[:]
-  v[0] = v[0] + v[1]
-  v[1] = rotate_left(v[1], rot)
-  v[1] = v[0] ^ v[1]
-  return v
-
-def rotate_list(xs):
-  return xs[1:] + xs[:1]
-
-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] + 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):
-  """Apply the Threefry 2x32 hash.
-
-  Args:
-    keypair: a pair of 32bit unsigned integers used for the key.
-    count: an array of dtype uint32 used for the counts.
-
-  Returns:
-    An array of dtype uint32 with the same shape as `count`.
-  """
-  x = [x1, x2]
-
-  rotations = [np.array([13, 15, 26, 6], dtype=np.uint32),
-               np.array([17, 29, 16, 24], dtype=np.uint32)]
-  ks = [key1, key2, key1 ^ key2 ^ np.uint32(0x1BD11BDA)]
-
-  x[0] = x[0] + ks[0]
-  x[1] = x[1] + ks[1]
-
-  if use_rolled_loops:
-    x, _, _ = lax.fori_loop(0, 5, rolled_loop_step, (x, rotate_list(ks), rotations))
-
-  else:
-    for r in rotations[0]:
-      x = apply_round(x, r)
-    x[0] = x[0] + ks[1]
-    x[1] = x[1] + ks[2] + np.uint32(1)
-
-    for r in rotations[1]:
-      x = apply_round(x, r)
-    x[0] = x[0] + ks[2]
-    x[1] = x[1] + ks[0] + np.uint32(2)
-
-    for r in rotations[0]:
-      x = apply_round(x, r)
-    x[0] = x[0] + ks[0]
-    x[1] = x[1] + ks[1] + np.uint32(3)
-
-    for r in rotations[1]:
-      x = apply_round(x, r)
-    x[0] = x[0] + ks[1]
-    x[1] = x[1] + ks[2] + np.uint32(4)
-
-    for r in rotations[0]:
-      x = apply_round(x, r)
-    x[0] = x[0] + ks[2]
-    x[1] = x[1] + ks[0] + np.uint32(5)
-
-  return tuple(x)
-
-
-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)))
-
-threefry2x32_p = core.Primitive("threefry2x32")
-threefry2x32_p.multiple_results = True
-threefry2x32_p.def_impl(partial(xla.apply_primitive, threefry2x32_p))
-threefry2x32_p.def_abstract_eval(_threefry2x32_abstract_eval)
-batching.defbroadcasting(threefry2x32_p)
-xla.translations_with_avals[threefry2x32_p] = xla.lower_fun(
-    partial(_threefry2x32_lowering, use_rolled_loops=False),
-    multiple_results=True, with_avals=True)
-xla.backend_specific_translations['cpu'][threefry2x32_p] = xla.lower_fun(
-    partial(_threefry2x32_lowering, use_rolled_loops=True),
-    multiple_results=True)
-if cuda_prng:
-  xla.backend_specific_translations['gpu'][threefry2x32_p] = \
-      _threefry2x32_gpu_translation_rule
-
-@jit
-def threefry_2x32(keypair, count):
-  """Apply the Threefry 2x32 hash.
-
-  Args:
-    keypair: a pair of 32bit unsigned integers used for the key.
-    count: an array of dtype uint32 used for the counts.
-
-  Returns:
-    An array of dtype uint32 with the same shape as `count`.
-  """
-  key1, key2 = keypair
-  if not lax.dtype(key1) == lax.dtype(key2) == lax.dtype(count) == np.uint32:
-    msg = "threefry_2x32 requires uint32 arguments, got {}"
-    raise TypeError(msg.format([lax.dtype(x) for x in [key1, key2, count]]))
-
-  try:
-    odd_size = count.size % 2
-  except core.InconclusiveDimensionOperation as e:
-    msg = ("jax.random functions have limited support for shape polymorphism. "
-           "In particular, the product of the known dimensions must be even.")
-    raise core.InconclusiveDimensionOperation(msg) from e
-
-  if odd_size:
-    x = list(jnp.split(jnp.concatenate([count.ravel(), np.uint32([0])]), 2))
-  else:
-    x = list(jnp.split(count.ravel(), 2))
-
-  x = threefry2x32_p.bind(key1, key2, x[0], x[1])
-  out = jnp.concatenate(x)
-  assert out.dtype == np.uint32
-  return lax.reshape(out[:-1] if odd_size else out, count.shape)
-
-
-def split(key: jnp.ndarray, num: int = 2) -> jnp.ndarray:
-  """Splits a PRNG key into `num` new keys by adding a leading axis.
-
-  Args:
-    key: a PRNGKey (an array with shape (2,) and dtype uint32).
-    num: optional, a positive integer indicating the number of keys to produce
-      (default 2).
-
-  Returns:
-    An array with shape (num, 2) and dtype uint32 representing `num` new keys.
-  """
-  return _split(key, int(num))  # type: ignore
-
-@partial(jit, static_argnums=(1,))
-def _split(key, num) -> jnp.ndarray:
-  counts = lax.iota(np.uint32, num * 2)
-  return lax.reshape(threefry_2x32(key, counts), (num, 2))
-
-
-def fold_in(key: jnp.ndarray, data: int) -> jnp.ndarray:
-  """Folds in data to a PRNG key to form a new PRNG key.
-
-  Args:
-    key: a PRNGKey (an array with shape (2,) and dtype uint32).
-    data: a 32bit integer representing data to be folded in to the key.
-
-  Returns:
-    A new PRNGKey that is a deterministic function of the inputs and is
-    statistically safe for producing a stream of new pseudo-random values.
-  """
-  return _fold_in(key, jnp.uint32(data))
-
-@jit
-def _fold_in(key, data):
-  return threefry_2x32(key, PRNGKey(data))
-
-
-@partial(jit, static_argnums=(1, 2))
-def _random_bits(key, bit_width, shape):
-  """Sample uniform random bits of given width and shape using PRNG key."""
-  if not _is_prng_key(key):
-    raise TypeError("_random_bits got invalid prng key.")
-  if bit_width not in (8, 16, 32, 64):
-    raise TypeError("requires 8-, 16-, 32- or 64-bit field width.")
-  shape = core.as_named_shape(shape)
-  for name, size in shape.named_items:
-    real_size = lax.psum(1, name)
-    if real_size != size:
-      raise ValueError(f"The shape of axis {name} was specified as {size}, "
-                       f"but it really is {real_size}")
-    axis_index = lax.axis_index(name)
-    key = fold_in(key, axis_index)
-  size = prod(shape.positional)
-  # Compute ceil(bit_width * size / 32) in a way that is friendly to shape
-  # polymorphism
-  max_count, r = divmod(bit_width * size, 32)
-  if r > 0:
-    max_count += 1
-
-  if core.is_constant_dim(max_count):
-    nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max)
-  else:
-    nblocks, rem = 0, max_count
-
-  if not nblocks:
-    bits = threefry_2x32(key, lax.iota(np.uint32, rem))
-  else:
-    keys = split(key, nblocks + 1)
-    subkeys, last_key = keys[:-1], keys[-1]
-    blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
-    last = threefry_2x32(last_key, lax.iota(np.uint32, rem))
-    bits = lax.concatenate([blocks.ravel(), last], 0)
-
-  dtype = _UINT_DTYPES[bit_width]
-  if bit_width == 64:
-    bits = [lax.convert_element_type(x, dtype) for x in jnp.split(bits, 2)]
-    bits = lax.shift_left(bits[0], dtype(32)) | bits[1]
-  elif bit_width in [8, 16]:
-    # this is essentially bits.view(dtype)[:size]
-    bits = lax.bitwise_and(
-      np.uint32(np.iinfo(dtype).max),
-      lax.shift_right_logical(
-        lax.broadcast(bits, (1,)),
-        lax.mul(
-          np.uint32(bit_width),
-          lax.broadcasted_iota(np.uint32, (32 // bit_width, 1), 0)
-        )
-      )
-    )
-    bits = lax.reshape(bits, (np.uint32(max_count * 32 // bit_width),), (1, 0))
-    bits = lax.convert_element_type(bits, dtype)[:size]
-  return lax.reshape(bits, shape)
-
-
 ### random samplers
 
 
@@ -414,7 +110,7 @@ def _uniform(key, shape, dtype, minval, maxval) -> jnp.ndarray:
   # equivalent float representations, which might not be true on all platforms.
   float_bits = lax.bitwise_or(
       lax.shift_right_logical(bits, np.array(nbits - nmant, lax.dtype(bits))),
-      np.array(1., dtype).view(_UINT_DTYPES[nbits]))
+      np.array(1., dtype).view(UINT_DTYPES[nbits]))
   floats = lax.bitcast_convert_type(float_bits, dtype) - np.array(1., dtype)
   return lax.max(
       minval,
@@ -480,7 +176,7 @@ def _randint(key, shape, minval, maxval, dtype):
   rbits = lambda key: _random_bits(key, nbits, shape)
   higher_bits, lower_bits = rbits(k1), rbits(k2)
 
-  unsigned_dtype = _UINT_DTYPES[nbits]
+  unsigned_dtype = UINT_DTYPES[nbits]
   span = lax.convert_element_type(maxval - minval, unsigned_dtype)
 
   # Ensure that span=1 when maxval <= minval, so minval is always returned;
diff --git a/jax/random.py b/jax/random.py
index 54c1d9acb050..6dda271c6db0 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -78,8 +78,15 @@
 
 # flake8: noqa: F401
 
-from jax._src.random import (
+from jax._src.prng import (
   PRNGKey,
+  fold_in,
+  split,
+  threefry2x32_p,
+  threefry_2x32,
+)
+
+from jax._src.random import (
   bernoulli,
   beta,
   categorical,
@@ -88,7 +95,6 @@
   dirichlet,
   double_sided_maxwell,
   exponential,
-  fold_in,
   gamma,
   gumbel,
   laplace,
@@ -103,10 +109,7 @@
   randint,
   random_gamma_p,
   shuffle,
-  split,
   t,
-  threefry2x32_p,
-  threefry_2x32,
   truncated_normal,
   uniform,
   weibull_min,


Code review comments:
diff_hunk: @@ -0,0 +1,475 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import partial
+from typing import Callable, Iterator, NamedTuple
+import warnings
+
+import numpy as np
+
+from jax import lax
+from jax import core
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src.api import jit, vmap
+from jax.config import config
+from jax.lib import xla_bridge
+from jax.lib import xla_client
+from jax.lib import cuda_prng
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax._src.pprint_util import pp, vcat
+from jax._src.util import prod
+
+
+UINT_DTYPES = {
+    8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}  # type: ignore[has-type]
+
+# -- PRNG implementation interface --
+
+class PRNGImpl(NamedTuple):
+  key_shape: core.Shape
+  seed: Callable
+  split: Callable
+  random_bits: Callable
+  fold_in: Callable
+
+  def pprint(self):
+    return (pp(self.__class__.__name__) >> pp(':')) + vcat([
+        pp(k) >> pp(' = ') >> pp(v) for k, v in self._asdict().items()
+    ]).indent(2)
+
+
+# -- PRNG key arrays --
+
+def _is_prng_key_data(impl, keys: jnp.ndarray) -> bool:
+  ndim = len(impl.key_shape)
+  try:
+    return (keys.ndim >= 1 and
+            keys.shape[-ndim:] == impl.key_shape and
+            keys.dtype == np.uint32)
+  except AttributeError:
+    return False
+
+@tree_util.register_pytree_node_class
+class PRNGKeyArray:
+  """An array whose elements are PRNG keys.
+
+  This class lifts the definition of a PRNG, provided in the form of a
+  ``PRNGImpl``, into an array-like pytree class. Instances of this
+  class behave like an array whose base elements are keys, hiding the
+  fact that keys are typically arrays (of ``uint32`` dtype) themselves.
+
+  PRNGKeyArrays are also restricted relative to JAX arrays in that
+  they do not expose arithmetic operations. They instead expose
+  wrapper methods around the PRNG implementation functions (``split``,
+  ``random_bits``, ``fold_in``).
+  """
+
+  impl: PRNGImpl
+  keys: jnp.ndarray
+
+  def __init__(self, impl, key_data: jnp.ndarray):
+    # key_data might be a placeholder python `object` or `bool`
+    # instead of a jnp.ndarray due to tree_unflatten
+    if (type(key_data) not in [object, bool] and
+        not _is_prng_key_data(impl, key_data)):
+      raise TypeError(
+          f'Invalid PRNG key data {key_data} for PRNG implementation {impl}')
+    self.impl = impl
+    self.keys = key_data
+
+  def tree_flatten(self):
+    return (self.keys,), self.impl
+
+  @classmethod
+  def tree_unflatten(cls, impl, keys):
+    keys, = keys
+    return cls(impl, keys)
+
+  @property
+  def dtype(self):
+    # TODO(frostig): remove after deprecation window
+    if config.jax_enable_custom_prng:
+      raise AttributeError("'PRNGKeyArray' has no attribute 'dtype'")
path: jax/_src/prng.py
text: @user1:
You might have already explained this to me, but why do we want to remove the ability to do `key.dtype`? Seems like in the current implementation this can only be `unit32`, regardless of the `PRNGImpl`.

@author:
I'm not sure what to return as the `dtype`? If the purpose is to be an array of opaque keys as elements, then `uint32` seems incorrect (that's the dtype of the data buffer representation of an individual key). I figured that `dtype` is one of the ways in which this array type is not a numpy or jax device array. Perhaps this is too cautious, and if so I'd be happy to revisit it.

diff_hunk: @@ -44,308 +45,102 @@
 DTypeLikeInt = Any
 DTypeLikeFloat = Any
 
+# TODO(frostig): simplify once we always enable_custom_prng
+KeyArray = Union[Array, prng.PRNGKeyArray]
 
-_UINT_DTYPES = {8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}
-
-
-def PRNGKey(seed: int) -> jnp.ndarray:
-  """Create a pseudo-random number generator (PRNG) key given an integer seed.
-
-  Args:
-    seed: a 64- or 32-bit integer used as the value of the key.
-
-  Returns:
-    A PRNG key, which is modeled as an array of shape (2,) and dtype uint32. The
-    key is constructed from a 64-bit seed by effectively bit-casting to a pair
-    of uint32 values (or from a 32-bit seed by first padding out with zeros).
-  """
-  # Avoid overflowerror in X32 mode by first converting ints to int64.
-  # This breaks JIT invariance of PRNGKey for large ints, but supports the
-  # common use-case of instantiating PRNGKey with Python hashes in X32 mode.
-  if isinstance(seed, int):
-    seed_arr = jnp.asarray(np.int64(seed))
-  else:
-    seed_arr = jnp.asarray(seed)
-  if seed_arr.shape:
-    raise TypeError(f"PRNGKey seed must be a scalar; got {seed!r}.")
-  if not np.issubdtype(seed_arr.dtype, np.integer):
-    raise TypeError(f"PRNGKey seed must be an integer; got {seed!r}")
-
-  convert = lambda k: lax.reshape(lax.convert_element_type(k, np.uint32), [1])
-  k1 = convert(lax.shift_right_logical(seed_arr, lax._const(seed_arr, 32)))
-  k2 = convert(jnp.bitwise_and(seed_arr, np.uint32(0xFFFFFFFF)))
-  return lax.concatenate([k1, k2], 0)
-
-def _is_prng_key(key: jnp.ndarray) -> bool:
-  try:
-    return key.shape == (2,) and key.dtype == np.uint32
-  except AttributeError:
-    return False
+UINT_DTYPES = prng.UINT_DTYPES
 
 
 ### utilities
 
 
-def _make_rotate_left(dtype):
-  if not jnp.issubdtype(dtype, np.integer):
-    raise TypeError("_rotate_left only accepts integer dtypes.")
-  nbits = np.array(jnp.iinfo(dtype).bits, dtype)
-
-  def _rotate_left(x, d):
-    if lax.dtype(d) != dtype:
-      d = lax.convert_element_type(d, dtype)
-    if lax.dtype(x) != dtype:
-      x = lax.convert_element_type(x, dtype)
-    return lax.shift_left(x, d) | lax.shift_right_logical(x, nbits - d)
-  return _rotate_left
-
-
-def _bit_stats(bits):
-  """This is a debugging function to compute the statistics of bit fields."""
-  return np.array([list(map(int, np.binary_repr(x, 64))) for x in bits]).mean(0)
-
-
-### hash function and split
+# 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 _check_prng_key(key):
+  # TODO(frostig): remove once we always enable_custom_prng
+  if type(key) is prng.PRNGKeyArray:
+    return key, False
+  elif _arraylike(key):
+    if config.jax_enable_custom_prng:
+      warnings.warn(
+          'Raw arrays as random keys to jax.random functions are deprecated. '
path: jax/_src/random.py
text: @user1:
Will this warning not trigger for every user if the default is `custom_prng=False`? Maybe there should be a hint on how to resolve the warning if that is the case?

More generally, do we expect users to turn on this `custom_prng` flag themselves? As I understand it, it's only a small group of users who rely on the keys being `jnp` arrays, and for most users this will be a no-op change, so I'm not sure they should be aware of the flag. I might be wrong though!

diff_hunk: @@ -0,0 +1,475 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import partial
+from typing import Callable, Iterator, NamedTuple
+import warnings
+
+import numpy as np
+
+from jax import lax
+from jax import core
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src.api import jit, vmap
+from jax.config import config
+from jax.lib import xla_bridge
+from jax.lib import xla_client
+from jax.lib import cuda_prng
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax._src.pprint_util import pp, vcat
+from jax._src.util import prod
+
+
+UINT_DTYPES = {
+    8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}  # type: ignore[has-type]
+
+# -- PRNG implementation interface --
+
+class PRNGImpl(NamedTuple):
+  key_shape: core.Shape
+  seed: Callable
+  split: Callable
+  random_bits: Callable
+  fold_in: Callable
path: jax/_src/prng.py
text: @user1:
can you more explicitly type the function signatures? Or document them here?

@author:
Good call! Done.

diff_hunk: @@ -0,0 +1,471 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import partial
+from typing import Callable, Iterator, NamedTuple
+import warnings
+
+import numpy as np
+
+from jax import lax
+from jax import core
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src.api import jit, vmap
+from jax.lib import xla_bridge
+from jax.lib import xla_client
+from jax.lib import cuda_prng
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax._src.pprint_util import pp, vcat
+from jax._src.util import prod
+
+
+UINT_DTYPES = {
+    8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}  # type: ignore[has-type]
+
+# -- PRNG implementation interface --
+
+class PRNGImpl(NamedTuple):
+  key_shape: core.Shape
+  seed: Callable
+  split: Callable
+  random_bits: Callable
+  fold_in: Callable
+
+  def pprint(self):
+    return (pp(self.__class__.__name__) >> pp(':')) + vcat([
+        pp(k) >> pp(' = ') >> pp(v) for k, v in self._asdict().items()
+    ]).indent(2)
+
+
+# -- PRNG key arrays --
+
+def _is_prng_key_data(impl, keys: jnp.ndarray) -> bool:
+  ndim = len(impl.key_shape)
+  try:
+    return (keys.ndim >= 1 and
+            keys.shape[-ndim:] == impl.key_shape and
+            keys.dtype == np.uint32)
+  except AttributeError:
+    return False
+
+@tree_util.register_pytree_node_class
+class PRNGKeyArray:
+  """An array whose elements are PRNG keys.
+
+  This class lifts the definition of a PRNG, provided in the form of a
+  ``PRNGImpl``, into an array-like pytree class. Instances of this
+  class behave like an array whose base elements are keys, hiding the
+  fact that keys are typically arrays (of ``uint32`` dtype) themselves.
+
+  PRNGKeyArrays are also restricted relative to JAX arrays in that
+  they do not expose arithmetic operations. They instead expose
+  wrapper methods around the PRNG implementation functions (``split``,
+  ``random_bits``, ``fold_in``).
+  """
+
+  impl: PRNGImpl
path: jax/_src/prng.py
text: @author:
It's a natural question! I kicked around an idea like this earlier on.

I'm choosing to wrap-and-delegate because it directly expresses what I'm trying to enforce. `PRNGKeyArray` is a class made to adapt a PRNG implementation, which is expressed on an individual key, into an "array" of keys. It's an internal definition, not part of the API, and we have no need in jax for a hierarchy of behaviors around it.

If you're writing custom PRNGs from outside of jax, you're welcome to supply the PRNG interface however you like. It can be a member of whatever fun type hierarchy you've set up in your code, or just a plain record of four functions. You'll notice it need not even be a proper `prng.PRNGImpl`, but rather anything that duck-types as one. Hopefully that's plenty of flexibility, less for us to document, and less for anyone to learn.

To your specific points:

> * simplify your flatten and unflatten methods by removing `self.impl`,
> * simplify your implementation by removing `PRNGImpl`,

These seem like minimal savings, so I'm not concerned.

> * simplify type annotations,

How so?

> * make it easier to have a hierarchy of behavior if the need arises,

I don't foresee such a need. We can always revisit if/when it comes up.

> * and make it easier for users to define their own `PRNGKeyArray`.

I don't want them to!

What do you think? As someone who might use this, do these choices get in your way at all?

@user1:
Thanks for the exhaustive reply.

> `PRNGKeyArray` is a class made to adapt a PRNG implementation, which is expressed on an individual key, into an "array" of keys. It's an internal definition, not part of the API, and we have no need in jax for a hierarchy of behaviors around it.

`PRNGKeyArray` is a member of the union `KeyArray`, which is passed to the public functions in `random.py`.  So, `PRNGKeyArray` is part of the API, isn't it?  How should a user annotate the generator that's produced by `seed_with_impl` and that's passed around to user functions?

About `PRNGImpl`, my point is that it could just be folded into PRNGKeyArray since it doesn't appear that you use it on its own anywhere except to produce a `PRNGKeyArray`.  I do understand your point that you want to treat `PRNGKeyArray` as a wrapper, but it only ever wraps one thing, and the thing it wraps is only ever used within the wrapper.

> we have no need in jax for a hierarchy of behaviors around it.

Fair enough.

> simplify type annotations,
> How so?

The functions in `PRNGImpl` are annotated as `Callable`. It would be nice to just define them as abstract methods so that they have full signatures.  I guess I don't see the point of using function-valued members that you set instead of abstract methods that you implement in subclasses.  Are you avoiding inheritance for some reason?

> make it easier for users to define their own PRNGKeyArray.
> I don't want them to!

One possible user-defined PRNGKeyArray subclasss  you might find is a hypothetical Haiku `HaikuKeyArray`.  It would simplify Haiku code to not have to extract keys from a `PRNGKeySequence`, and instead use the `HaikuKeyArray`, which would be passable to all the JAX random functions.  It would be mutable, and it would hold a `PRNGKeyArray` internally, which it would split, just like the `PRNGKeySequence` currently does.  Why make it more difficult to implement such a subclass?

> do these choices get in your way at all?

They do not!  My real motivation is to throw away my `Generator` class [here](https://github.com/NeilGirdhar/tjax/blob/master/tjax/_src/generator.py#L19) with a superior solution from the JAX codebase :smile: I just thought it would be nicer to have one class rather than two, and to use abstract methods rather than function-valued members.  In my code, I pass around `Generator` everywhere.  I had hoped for a similar solution like `JaxGenerator` or something like that :smile: 

Anyway, I think I've expressed my points.  Thanks for the reply.  I'm sure whatever you choose will work great.

diff_hunk: @@ -0,0 +1,475 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import partial
+from typing import Callable, Iterator, NamedTuple
+import warnings
+
+import numpy as np
+
+from jax import lax
+from jax import core
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src.api import jit, vmap
+from jax.config import config
+from jax.lib import xla_bridge
+from jax.lib import xla_client
+from jax.lib import cuda_prng
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax._src.pprint_util import pp, vcat
+from jax._src.util import prod
+
+
+UINT_DTYPES = {
+    8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}  # type: ignore[has-type]
+
+# -- PRNG implementation interface --
+
+class PRNGImpl(NamedTuple):
+  key_shape: core.Shape
+  seed: Callable
+  split: Callable
+  random_bits: Callable
+  fold_in: Callable
+
+  def pprint(self):
+    return (pp(self.__class__.__name__) >> pp(':')) + vcat([
+        pp(k) >> pp(' = ') >> pp(v) for k, v in self._asdict().items()
+    ]).indent(2)
+
+
+# -- PRNG key arrays --
+
+def _is_prng_key_data(impl, keys: jnp.ndarray) -> bool:
+  ndim = len(impl.key_shape)
+  try:
+    return (keys.ndim >= 1 and
+            keys.shape[-ndim:] == impl.key_shape and
+            keys.dtype == np.uint32)
+  except AttributeError:
+    return False
+
+@tree_util.register_pytree_node_class
+class PRNGKeyArray:
+  """An array whose elements are PRNG keys.
+
+  This class lifts the definition of a PRNG, provided in the form of a
+  ``PRNGImpl``, into an array-like pytree class. Instances of this
+  class behave like an array whose base elements are keys, hiding the
+  fact that keys are typically arrays (of ``uint32`` dtype) themselves.
+
+  PRNGKeyArrays are also restricted relative to JAX arrays in that
+  they do not expose arithmetic operations. They instead expose
+  wrapper methods around the PRNG implementation functions (``split``,
+  ``random_bits``, ``fold_in``).
+  """
+
+  impl: PRNGImpl
+  keys: jnp.ndarray
+
+  def __init__(self, impl, key_data: jnp.ndarray):
+    # key_data might be a placeholder python `object` or `bool`
+    # instead of a jnp.ndarray due to tree_unflatten
+    if (type(key_data) not in [object, bool] and
+        not _is_prng_key_data(impl, key_data)):
+      raise TypeError(
+          f'Invalid PRNG key data {key_data} for PRNG implementation {impl}')
+    self.impl = impl
+    self.keys = key_data
+
+  def tree_flatten(self):
+    return (self.keys,), self.impl
+
+  @classmethod
+  def tree_unflatten(cls, impl, keys):
+    keys, = keys
+    return cls(impl, keys)
+
+  @property
+  def dtype(self):
+    # TODO(frostig): remove after deprecation window
+    if config.jax_enable_custom_prng:
+      raise AttributeError("'PRNGKeyArray' has no attribute 'dtype'")
+    else:
+      warnings.warn(
+          'deprecated `dtype` attribute of PRNG key arrays', FutureWarning)
+      return np.uint32
+
+  def shape(self):
+    # TODO(frostig): simplify once we always enable_custom_prng
+    if config.jax_enable_custom_prng:
+      return self._shape
+    else:
+      warnings.warn(
+          'deprecated `shape` attribute of PRNG key arrays. In a future version '
+          'of JAX this attribute will be removed or its value may change.',
+          FutureWarning)
+      return self.keys.shape
+
+  @property
+  def _shape(self):
+    base_ndim = len(self.impl.key_shape)
+    return self.keys.shape[:-base_ndim]
path: jax/_src/prng.py
text: @user1:
Should we assert that `self._shape` is always equal to `self.impl.key_shape`? Do we ever expect them to be different?

@author:
The two are different: `_shape` should be the shape of the encompassing "array of keys," whereas `impl.key_shape` is the shape of an individual key. The underlying buffer holding all the key data is an array of shape `(*self._shape, *self.impl.key_shape)`.

diff_hunk: @@ -0,0 +1,475 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from functools import partial
+from typing import Callable, Iterator, NamedTuple
+import warnings
+
+import numpy as np
+
+from jax import lax
+from jax import core
+from jax import numpy as jnp
+from jax import tree_util
+from jax._src.api import jit, vmap
+from jax.config import config
+from jax.lib import xla_bridge
+from jax.lib import xla_client
+from jax.lib import cuda_prng
+from jax.interpreters import batching
+from jax.interpreters import xla
+from jax._src.pprint_util import pp, vcat
+from jax._src.util import prod
+
+
+UINT_DTYPES = {
+    8: jnp.uint8, 16: jnp.uint16, 32: jnp.uint32, 64: jnp.uint64}  # type: ignore[has-type]
+
+# -- PRNG implementation interface --
+
+class PRNGImpl(NamedTuple):
+  key_shape: core.Shape
+  seed: Callable
+  split: Callable
+  random_bits: Callable
+  fold_in: Callable
+
+  def pprint(self):
+    return (pp(self.__class__.__name__) >> pp(':')) + vcat([
+        pp(k) >> pp(' = ') >> pp(v) for k, v in self._asdict().items()
+    ]).indent(2)
+
+
+# -- PRNG key arrays --
+
+def _is_prng_key_data(impl, keys: jnp.ndarray) -> bool:
+  ndim = len(impl.key_shape)
+  try:
+    return (keys.ndim >= 1 and
+            keys.shape[-ndim:] == impl.key_shape and
+            keys.dtype == np.uint32)
+  except AttributeError:
+    return False
+
+@tree_util.register_pytree_node_class
+class PRNGKeyArray:
+  """An array whose elements are PRNG keys.
+
+  This class lifts the definition of a PRNG, provided in the form of a
+  ``PRNGImpl``, into an array-like pytree class. Instances of this
+  class behave like an array whose base elements are keys, hiding the
+  fact that keys are typically arrays (of ``uint32`` dtype) themselves.
+
+  PRNGKeyArrays are also restricted relative to JAX arrays in that
+  they do not expose arithmetic operations. They instead expose
+  wrapper methods around the PRNG implementation functions (``split``,
+  ``random_bits``, ``fold_in``).
+  """
+
+  impl: PRNGImpl
+  keys: jnp.ndarray
+
+  def __init__(self, impl, key_data: jnp.ndarray):
+    # key_data might be a placeholder python `object` or `bool`
+    # instead of a jnp.ndarray due to tree_unflatten
+    if (type(key_data) not in [object, bool] and
+        not _is_prng_key_data(impl, key_data)):
+      raise TypeError(
+          f'Invalid PRNG key data {key_data} for PRNG implementation {impl}')
+    self.impl = impl
+    self.keys = key_data
+
+  def tree_flatten(self):
+    return (self.keys,), self.impl
+
+  @classmethod
+  def tree_unflatten(cls, impl, keys):
+    keys, = keys
+    return cls(impl, keys)
+
+  @property
+  def dtype(self):
+    # TODO(frostig): remove after deprecation window
+    if config.jax_enable_custom_prng:
+      raise AttributeError("'PRNGKeyArray' has no attribute 'dtype'")
+    else:
+      warnings.warn(
+          'deprecated `dtype` attribute of PRNG key arrays', FutureWarning)
+      return np.uint32
+
+  def shape(self):
+    # TODO(frostig): simplify once we always enable_custom_prng
+    if config.jax_enable_custom_prng:
+      return self._shape
+    else:
+      warnings.warn(
+          'deprecated `shape` attribute of PRNG key arrays. In a future version '
+          'of JAX this attribute will be removed or its value may change.',
+          FutureWarning)
+      return self.keys.shape
+
+  @property
+  def _shape(self):
+    base_ndim = len(self.impl.key_shape)
+    return self.keys.shape[:-base_ndim]
+
+  def _is_scalar(self):
+    base_ndim = len(self.impl.key_shape)
+    return self.keys.ndim == base_ndim
+
+  def __len__(self):
+    if self._is_scalar():
+      raise TypeError('len() of unsized object')
+    return len(self.keys)
+
+  def __iter__(self) -> Iterator['PRNGKeyArray']:
+    if self._is_scalar():
+      raise TypeError('iteration over a 0-d single PRNG key')
+    return (PRNGKeyArray(self.impl, k) for k in iter(self.keys))
+
+  def __getitem__(self, idx) -> 'PRNGKeyArray':
+    if not isinstance(idx, tuple):
+      idx = (idx,)
+    if any(type(i) is not int for i in idx):
+      raise NotImplementedError(
+          'PRNGKeyArray only supports indexing with integer indices. '
+          f'Cannot index at {idx}')
+    base_ndim = len(self.impl.key_shape)
+    ndim = self.keys.ndim - base_ndim
+    if len(idx) > ndim:
+      raise IndexError(
+          f'too many indices for PRNGKeyArray: array is {ndim}-dimensional '
+          f'but {len(idx)} were indexed')
+    return PRNGKeyArray(self.impl, self.keys[idx])
+
+  def _fold_in(self, data: int) -> 'PRNGKeyArray':
+    return PRNGKeyArray(self.impl, self.impl.fold_in(self.keys, data))
+
+  def _random_bits(self, bit_width, shape) -> jnp.ndarray:
+    return self.impl.random_bits(self.keys, bit_width, shape)
+
+  def _split(self, num: int) -> 'PRNGKeyArray':
+    return PRNGKeyArray(self.impl, self.impl.split(self.keys, num))
+
+  def __repr__(self):
+    arr_shape = self._shape
+    pp_keys = pp('shape = ') >> pp(arr_shape)
+    if isinstance(self.impl, PRNGImpl):
+      pp_impl = self.impl.pprint()
+    else:
+      pp_impl = pp(self.impl)
path: jax/_src/prng.py
text: @user1:
Isn't `self.impl` always a `PRNGImpl`?