mirror of
https://github.com/ROCm/jax.git
synced 2025-04-25 08:26:07 +00:00
7f06df1ea1
Before this change, the Python PRNG key array was a pytree type
wrapping a `uint32` array. This was a stopgap that misbehaved under
`vmap`, `scan`, and even `jax.tree_map`. For a while, we thought we
might rely on something like the typeclass mechanisms in development
(e.g. `vmappable`) to move away from a pytree.
We're now taking a different approach: introducing key element types
into our IR and other internal machinery. During staging, we map
user-facing PRNG key arrays to abstract arrays such element type.
This leans heavily on our recently-introduced extended element type
capabilities.
As a consequence, `vmap`, `scan`, etc. now work.
A sample of changes made to introduce key-element-type arrays:
* Introduce a new element type (`prng.KeyTy`), with the requisite IR
type mapping and device result handlers, as well as lowering rules
for dtype-polymorphic primitive operations.
* Introduce primitives for basic RNG operations: `random_seed`,
`random_bits`, `random_split`, `random_fold_in`. These primitives
essentially delegate to the underlying PRNG implementation (directly
so in their impl rules, and by translating their staged-out form in
lowering rules).
* Also introduce `random_wrap` and `random_unwrap` for "unsafe"
conversion from/to the base `uint32` array. We need this backwards
compatibility, and it's useful for tests.
* Introduce some `vmap`-based helpers to adapt PRNG impls (which
define basic `random_bits`, `split`, etc. on scalars) to the above
batch-polymorphic primitives. Most of the primitives are vectorized,
but `random_fold_in` is a broadcasting binary op.
* Update the `gamma` primitive rules to account for key-element-type
abstract arrays (nice simplification here).
* Give PRNG implementation short string names ("tags") for IR
pretty-printing.
* Update `lax.stop_gradient` to handle opaque dtypes.
* Fix up loop MLIR lowering, which assumed that shaped arrays of all
dtypes have the same physical shape.
* Add new tests (exercising staging, jaxprs, lowerings, ...)
A sample of changes made to rework Python-level PRNG key arrays:
* Mimic `isinstance(x, KeyArray)` checks on abstract key arrays and
tracers that carry them.
* Patch (only a subset of) standard device array attributes onto PRNG
key arrays.
* Implement various conversion handlers (sharding, constant-creation,
`device_put`).
* Accept PRNG key arrays as input to `lax_numpy.transpose`.
* Update tests and rename some internals.
A sample of extra changes along the way:
* Disallow AD on key-typed arrays in the main API.
* Hoist `random_bits`'s named-shape-handling logic, which used to only
take place in the threefry PRNG's `random_bits` implementation, up
to the new `random_bits` traceable, so that we apply it consistently
across PRNG implementations.
This change leaves some unwanted `lax` and `jax.numpy` operations
superficially available on key arrays during tracing/staging
(e.g. under `jit`), though not outside of it. We ultimately want to
disallow these and raise useful errors, and I'm leaving that for
follow-up work. For now, applying such operations under `jit` may
result in downstream errors in the middle-end instead.
Everything here is still guarded by `config.jax_enable_custom_prng`,
whose default setting hasn't changed (it is off).
181 lines
6.3 KiB
Python
181 lines
6.3 KiB
Python
# Copyright 2018 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.
|
|
|
|
"""Utilities for pseudo-random number generation.
|
|
|
|
The ``jax.random`` package provides a number of routines for deterministic
|
|
generation of sequences of pseudorandom numbers.
|
|
|
|
Basic usage
|
|
-----------
|
|
|
|
>>> seed = 1701
|
|
>>> num_steps = 100
|
|
>>> key = jax.random.PRNGKey(seed)
|
|
>>> for i in range(num_steps):
|
|
... key, subkey = jax.random.split(key)
|
|
... params = compiled_update(subkey, params, next(batches)) # doctest: +SKIP
|
|
|
|
PRNG Keys
|
|
---------
|
|
|
|
Unlike the *stateful* pseudorandom number generators (PRNGs) that users of NumPy and
|
|
SciPy may be accustomed to, JAX random functions all require an explicit PRNG state to
|
|
be passed as a first argument.
|
|
The random state is described by two unsigned 32-bit integers that we call a **key**,
|
|
usually generated by the :py:func:`jax.random.PRNGKey` function::
|
|
|
|
>>> from jax import random
|
|
>>> key = random.PRNGKey(0)
|
|
>>> key
|
|
DeviceArray([0, 0], dtype=uint32)
|
|
|
|
This key can then be used in any of JAX's random number generation routines::
|
|
|
|
>>> random.uniform(key)
|
|
DeviceArray(0.41845703, dtype=float32)
|
|
|
|
Note that using a key does not modify it, so reusing the same key will lead to the same result::
|
|
|
|
>>> random.uniform(key)
|
|
DeviceArray(0.41845703, dtype=float32)
|
|
|
|
If you need a new random number, you can use :meth:`jax.random.split` to generate new subkeys::
|
|
|
|
>>> key, subkey = random.split(key)
|
|
>>> random.uniform(subkey)
|
|
DeviceArray(0.10536897, dtype=float32)
|
|
|
|
Advanced
|
|
--------
|
|
|
|
Design and Context
|
|
==================
|
|
|
|
**TLDR**: JAX PRNG = `Threefry counter PRNG <http://www.thesalmons.org/john/random123/papers/random123sc11.pdf>`_
|
|
+ a functional array-oriented `splitting model <https://dl.acm.org/citation.cfm?id=2503784>`_
|
|
|
|
See `docs/jep/263-prng.md <https://github.com/google/jax/blob/main/docs/jep/263-prng.md>`_
|
|
for more details.
|
|
|
|
To summarize, among other requirements, the JAX PRNG aims to:
|
|
|
|
1. ensure reproducibility,
|
|
2. parallelize well, both in terms of vectorization (generating array values)
|
|
and multi-replica, multi-core computation. In particular it should not use
|
|
sequencing constraints between random function calls.
|
|
|
|
Advanced RNG configuration
|
|
==========================
|
|
|
|
JAX provides several PRNG implementations (controlled by the
|
|
`jax_default_prng_impl` flag).
|
|
|
|
- **default**
|
|
`A counter-based PRNG built around the Threefry hash function <http://www.thesalmons.org/john/random123/papers/random123sc11.pdf>`_.
|
|
- *experimental* A PRNG that thinly wraps the XLA Random Bit Generator (RBG) algorithm. See
|
|
`TF doc <https://www.tensorflow.org/xla/operation_semantics#rngbitgenerator>`_.
|
|
|
|
- "rbg" uses ThreeFry for splitting, and XLA RBG for data generation.
|
|
- "unsafe_rbg" exists only for demonstration purposes, using RBG both for
|
|
splitting (using an untested made up algorithm) and generating.
|
|
|
|
The random streams generated by these experimental implementations haven't
|
|
been subject to any empirical randomness testing (e.g. Big Crush). The
|
|
random bits generated may change between JAX versions.
|
|
|
|
The possible reasons not use the default RNG are:
|
|
|
|
1. it may be slow to compile (specifically for Google Cloud TPUs)
|
|
2. it's slower to execute on TPUs
|
|
|
|
Here is a short summary:
|
|
|
|
.. table::
|
|
:widths: auto
|
|
|
|
================================= ================= === ==========
|
|
Property ThreeFry, default rbg unsafe_rbg
|
|
================================= ================= === ==========
|
|
Fast on TPU ✅ ✅
|
|
always correct w/ scan ✅ ✅
|
|
always correct w/ remat ✅ ✅
|
|
identical across CPU/GPU/TPU ✅ ✅
|
|
identical across JAX/XLA versions ✅
|
|
identical across shardings ✅ ✅ ✅
|
|
================================= ================= === ==========
|
|
|
|
NOTE: RNGs are currently identical across shardings because the random value
|
|
is first materialized replicated on each device and then the slice that each
|
|
device needs is later sliced out.
|
|
"""
|
|
|
|
from jax._src.prng import PRNGKeyArray
|
|
# TODO(frostig): remove this typechecking workaround. Our move away
|
|
# from PRNGKeyArray as a pytree led to Python typechecker breakages in
|
|
# several downstream annotations (e.g. annotations in jax-dependent
|
|
# libraries that are violated by their callers). It may be that the
|
|
# pytree registration decorator invalidated the checks. This will be
|
|
# easier to handle after we always enable_custom_prng.
|
|
import typing
|
|
if typing.TYPE_CHECKING:
|
|
KeyArray = typing.Any
|
|
else:
|
|
# TODO(frostig): replace with KeyArray from jax._src.random once we
|
|
# always enable_custom_prng
|
|
KeyArray = PRNGKeyArray
|
|
|
|
|
|
from jax._src.random import (
|
|
PRNGKey as PRNGKey,
|
|
ball as ball,
|
|
bernoulli as bernoulli,
|
|
beta as beta,
|
|
categorical as categorical,
|
|
cauchy as cauchy,
|
|
choice as choice,
|
|
default_prng_impl as default_prng_impl,
|
|
dirichlet as dirichlet,
|
|
double_sided_maxwell as double_sided_maxwell,
|
|
exponential as exponential,
|
|
fold_in as fold_in,
|
|
gamma as gamma,
|
|
generalized_normal as generalized_normal,
|
|
gumbel as gumbel,
|
|
laplace as laplace,
|
|
logistic as logistic,
|
|
loggamma as loggamma,
|
|
maxwell as maxwell,
|
|
multivariate_normal as multivariate_normal,
|
|
normal as normal,
|
|
orthogonal as orthogonal,
|
|
pareto as pareto,
|
|
permutation as permutation,
|
|
poisson as poisson,
|
|
rademacher as rademacher,
|
|
randint as randint,
|
|
random_gamma_p as random_gamma_p,
|
|
rbg_key as rbg_key,
|
|
shuffle as shuffle,
|
|
split as split,
|
|
t as t,
|
|
threefry_2x32 as threefry_2x32,
|
|
threefry2x32_key as threefry2x32_key,
|
|
threefry2x32_p as threefry2x32_p,
|
|
truncated_normal as truncated_normal,
|
|
uniform as uniform,
|
|
unsafe_rbg_key as unsafe_rbg_key,
|
|
weibull_min as weibull_min,
|
|
)
|