Repository: google/jax
Pull request #22989 (2024-08-11 15:11:42)

update README and several docs to use typed RNG keys


Patch:
diff --git a/README.md b/README.md
index b19d7b9ff..8f50aa42a 100644
--- a/README.md
+++ b/README.md
@@ -273,7 +273,7 @@ from jax import random, pmap
 import jax.numpy as jnp
 
 # Create 8 random 5000 x 6000 matrices, one per GPU
-keys = random.split(random.PRNGKey(0), 8)
+keys = random.split(random.key(0), 8)
 mats = pmap(lambda key: random.normal(key, (5000, 6000)))(keys)
 
 # Run a local matmul on each device in parallel (no data transfer)
diff --git a/benchmarks/api_benchmark.py b/benchmarks/api_benchmark.py
index c68dab85d..710ffb6d7 100644
--- a/benchmarks/api_benchmark.py
+++ b/benchmarks/api_benchmark.py
@@ -839,7 +839,7 @@ def serial_dot_products(state):
       out = out + y * x[0]
     return out
 
-  x = jax.random.normal(jax.random.PRNGKey(0), (2, 2))
+  x = jax.random.normal(jax.random.key(0), (2, 2))
   f(x).block_until_ready()  # compile
   while state:
     f(x).block_until_ready()
@@ -929,7 +929,7 @@ def jit_add_chain(state):
   def g(x, y):
     return lax.add(x, y)
 
-  x = jax.random.normal(jax.random.PRNGKey(0), (2, 2))
+  x = jax.random.normal(jax.random.key(0), (2, 2))
   while state:
     @jax.jit
     def f(x):
diff --git a/benchmarks/sparse_benchmark.py b/benchmarks/sparse_benchmark.py
index 65550b9cf..d6328881d 100644
--- a/benchmarks/sparse_benchmark.py
+++ b/benchmarks/sparse_benchmark.py
@@ -109,7 +109,7 @@ def sparse_bcoo_todense_compile(state):
 def _sparse_bcoo_matvec(state, jit: bool = False, compile: bool = False):
   shape = (2000, 2000)
   nse = 10000
-  key = jax.random.PRNGKey(1701)
+  key = jax.random.key(1701)
   mat = sparse.random_bcoo(
       key,
       nse=nse,
diff --git a/cloud_tpu_colabs/JAX_NeurIPS_2020_demo.ipynb b/cloud_tpu_colabs/JAX_NeurIPS_2020_demo.ipynb
index 1278bd01c..279aef3e9 100644
--- a/cloud_tpu_colabs/JAX_NeurIPS_2020_demo.ipynb
+++ b/cloud_tpu_colabs/JAX_NeurIPS_2020_demo.ipynb
@@ -38,7 +38,7 @@
    },
    "outputs": [],
    "source": [
-    "key = random.PRNGKey(0)\n",
+    "key = random.key(0)\n",
     "key, subkey = random.split(key)\n",
     "x = random.normal(key, (5000, 5000))\n",
     "\n",
@@ -189,7 +189,7 @@
    },
    "outputs": [],
    "source": [
-    "key = random.PRNGKey(0)\n",
+    "key = random.key(0)\n",
     "x = random.normal(key, ())\n",
     "\n",
     "print(grad(f)(x))\n",
@@ -261,7 +261,7 @@
    },
    "outputs": [],
    "source": [
-    "key = random.PRNGKey(0)\n",
+    "key = random.key(0)\n",
     "x = random.normal(key, (5000, 5000))"
    ]
   },
diff --git a/cloud_tpu_colabs/JAX_demo.ipynb b/cloud_tpu_colabs/JAX_demo.ipynb
index 6a6993f44..9acb1971c 100644
--- a/cloud_tpu_colabs/JAX_demo.ipynb
+++ b/cloud_tpu_colabs/JAX_demo.ipynb
@@ -27,7 +27,7 @@
     "import jax.numpy as jnp\n",
     "from jax import random\n",
     "\n",
-    "key = random.PRNGKey(0)"
+    "key = random.key(0)"
    ]
   },
   {
@@ -194,7 +194,7 @@
    },
    "outputs": [],
    "source": [
-    "key = random.PRNGKey(0)\n",
+    "key = random.key(0)\n",
     "x = random.normal(key, ())\n",
     "\n",
     "print(grad(f)(x))\n",
@@ -246,7 +246,7 @@
     "\n",
     "layer_sizes = [5, 2, 3]\n",
     "\n",
-    "key = random.PRNGKey(0)\n",
+    "key = random.key(0)\n",
     "key, *keys = random.split(key, len(layer_sizes))\n",
     "params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))\n",
     "\n",
@@ -351,7 +351,7 @@
    },
    "outputs": [],
    "source": [
-    "key = random.PRNGKey(0)\n",
+    "key = random.key(0)\n",
     "x = random.normal(key, (5000, 5000))"
    ]
   },
@@ -754,7 +754,7 @@
    },
    "outputs": [],
    "source": [
-    "keys = random.split(random.PRNGKey(0), 8)\n",
+    "keys = random.split(random.key(0), 8)\n",
     "mats = pmap(lambda key: random.normal(key, (5000, 5000)))(keys)\n",
     "result = pmap(jnp.dot)(mats, mats)\n",
     "print(pmap(jnp.mean)(result))"
diff --git a/cloud_tpu_colabs/Lorentz_ODE_Solver.ipynb b/cloud_tpu_colabs/Lorentz_ODE_Solver.ipynb
index 1777d3d1e..84abf8658 100644
--- a/cloud_tpu_colabs/Lorentz_ODE_Solver.ipynb
+++ b/cloud_tpu_colabs/Lorentz_ODE_Solver.ipynb
@@ -366,7 +366,7 @@
         "\n",
         "# set some initial conditions for each replicate\n",
         "ys = jnp.zeros((N_dev, N, 3))\n",
-        "state0 = jr.uniform(jr.PRNGKey(1), \n",
+        "state0 = jr.uniform(jr.key(1), \n",
         "                    minval=-1., maxval=1.,\n",
         "                    shape=(N_dev, 3))\n",
         "state0 = state0 * jnp.array([18,18,1]) + jnp.array((0.,0.,10.))\n",
diff --git a/cloud_tpu_colabs/Pmap_Cookbook.ipynb b/cloud_tpu_colabs/Pmap_Cookbook.ipynb
index 4f4ba8c16..981f0a9e8 100644
--- a/cloud_tpu_colabs/Pmap_Cookbook.ipynb
+++ b/cloud_tpu_colabs/Pmap_Cookbook.ipynb
@@ -263,7 +263,7 @@
         "from jax import random\n",
         "\n",
         "# create 8 random keys\n",
-        "keys = random.split(random.PRNGKey(0), 8)\n",
+        "keys = random.split(random.key(0), 8)\n",
         "# create a 5000 x 6000 matrix on each device by mapping over keys\n",
         "mats = pmap(lambda key: random.normal(key, (5000, 6000)))(keys)\n",
         "# the stack of matrices is represented logically as a single array\n",
diff --git a/docs/Custom_Operation_for_GPUs.md b/docs/Custom_Operation_for_GPUs.md
index 28a817154..8490bd489 100644
--- a/docs/Custom_Operation_for_GPUs.md
+++ b/docs/Custom_Operation_for_GPUs.md
@@ -306,7 +306,7 @@ per_core_batch_size=4
 seq_len=512
 emb_dim=512
 x = jax.random.normal(
-    jax.random.PRNGKey(0),
+    jax.random.key(0),
     shape=(jax.local_device_count() * per_core_batch_size, seq_len, emb_dim),
     dtype=jnp.bfloat16,
 )
diff --git a/docs/Custom_Operation_for_GPUs.py b/docs/Custom_Operation_for_GPUs.py
index 4c0b4b6f7..31a00c490 100644
--- a/docs/Custom_Operation_for_GPUs.py
+++ b/docs/Custom_Operation_for_GPUs.py
@@ -479,7 +479,7 @@ seq_len = 512
 emb_dim = 512
 assert jax.local_device_count() > 1, "Only 1 GPU, the example work, but it is this really what you want?"
 x = jax.random.normal(
-    jax.random.PRNGKey(0),
+    jax.random.key(0),
     shape=(jax.local_device_count() * per_core_batch_size, seq_len, emb_dim),
     dtype=jnp.float16,
 )
diff --git a/docs/_tutorials/advanced-autodiff.md b/docs/_tutorials/advanced-autodiff.md
index da95f96d8..d58a45d1d 100644
--- a/docs/_tutorials/advanced-autodiff.md
+++ b/docs/_tutorials/advanced-autodiff.md
@@ -590,7 +590,7 @@ def vmap_mjp(f, x, M):
     outs, = vmap(vjp_fun)(M)
     return outs
 
-key = random.PRNGKey(0)
+key = random.key(0)
 num_covecs = 128
 U = random.normal(key, (num_covecs,) + y.shape)
 
@@ -714,7 +714,7 @@ Here's a check:
 
 ```{code-cell}
 def check(seed):
-  key = random.PRNGKey(seed)
+  key = random.key(seed)
 
   # random coeffs for u and v
   key, subkey = random.split(key)
@@ -768,7 +768,7 @@ Here's a check of the VJP rules:
 
 ```{code-cell}
 def check(seed):
-  key = random.PRNGKey(seed)
+  key = random.key(seed)
 
   # random coeffs for u and v
   key, subkey = random.split(key)
diff --git a/docs/notebooks/shard_map.ipynb b/docs/notebooks/shard_map.ipynb
index 919690d23..6fbe67f4d 100644
--- a/docs/notebooks/shard_map.ipynb
+++ b/docs/notebooks/shard_map.ipynb
@@ -1509,7 +1509,7 @@
     "layer_sizes = [784, 128, 128, 128, 128, 128, 8]\n",
     "batch_size = 32\n",
     "\n",
-    "params, batch = init(jax.random.PRNGKey(0), layer_sizes, batch_size)"
+    "params, batch = init(jax.random.key(0), layer_sizes, batch_size)"
    ]
   },
   {
diff --git a/docs/notebooks/shard_map.md b/docs/notebooks/shard_map.md
index 6f9dfbb65..389b62c8a 100644
--- a/docs/notebooks/shard_map.md
+++ b/docs/notebooks/shard_map.md
@@ -1055,7 +1055,7 @@ def init(key, layer_sizes, batch_size):
 layer_sizes = [784, 128, 128, 128, 128, 128, 8]
 batch_size = 32
 
-params, batch = init(jax.random.PRNGKey(0), layer_sizes, batch_size)
+params, batch = init(jax.random.key(0), layer_sizes, batch_size)
 ```
 
 Compare these examples with the purely [automatic partitioning examples in the
diff --git a/jax/_src/extend/ffi.py b/jax/_src/extend/ffi.py
index 66af4f331..df1c09eff 100644
--- a/jax/_src/extend/ffi.py
+++ b/jax/_src/extend/ffi.py
@@ -22,6 +22,7 @@ from typing import Any
 
 from jax._src import core
 from jax._src import dispatch
+from jax._src import dtypes
 from jax._src import util
 from jax._src.callback import _check_shape_dtype, callback_batching_rule
 from jax._src.interpreters import ad
@@ -31,6 +32,7 @@ from jax._src.lib import jaxlib
 from jax._src.lib import xla_client
 from jax._src.lib.mlir import ir
 from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray
+import numpy as np
 
 map, unsafe_map = util.safe_map, map
 
@@ -135,7 +137,7 @@ def ffi_lowering(
     kwargs = dict(lowering_args)
     kwargs.setdefault("api_version", 4)
     kwargs["backend_config"] = dict(
-      backend_config or {}, **{k: mlir.ir_attribute(v) for k, v in params.items()})
+      backend_config or {}, **{k: _ir_attribute(v) for k, v in params.items()})
     if "result_types" not in kwargs:
       kwargs["result_types"] = [mlir.aval_to_ir_type(aval) for aval in ctx.avals_out]
     if operand_layouts is None:
@@ -152,6 +154,28 @@ def _default_layouts(shapes: Iterable[Sequence[DimSize]]) -> list[list[DimSize]]
   return [list(reversed(range(len(shape)))) for shape in shapes]
 
 
+def _ir_attribute(obj: Any) -> ir.Attribute:
+  # TODO(dfm): Similar functions exist in Pallas and Mosaic GPU. Perhaps these
+  # could be consolidated into mlir or similar.
+  if isinstance(obj, str):
+    return ir.StringAttr.get(obj)
+  elif isinstance(obj, bool):
+    return ir.BoolAttr.get(obj)
+  elif isinstance(obj, int):
+    return mlir.i64_attr(obj)
+  elif isinstance(obj, float):
+    return ir.FloatAttr.get_f64(obj)
+  elif hasattr(obj, "dtype"):
+    if not (dtypes.is_python_scalar(obj) or np.isscalar(obj)):
+      raise TypeError("Only scalar attributes are supported")
+    mlir_type = mlir.dtype_to_ir_type(obj.dtype)
+    if isinstance(mlir_type, ir.IntegerType):
+      return ir.IntegerAttr.get(mlir_type, obj)
+    elif isinstance(mlir_type, ir.FloatType):
+      return ir.FloatAttr.get(mlir_type, obj)
+  raise TypeError(f"Unsupported attribute type: {type(obj)}")
+
+
 def ffi_call(
     target_name: str,
     result_shape_dtypes: DuckTypedArray | Sequence[DuckTypedArray],
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 4d15e803a..84b1824bd 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -277,7 +277,12 @@ def ir_constant(val: Any) -> IrValues:
   raise TypeError(f"No constant handler for type: {type(val)}")
 
 def _numpy_array_constant(x: np.ndarray | np.generic) -> IrValues:
-  attr = _numpy_array_attribute(x)
+  element_type = dtype_to_ir_type(x.dtype)
+  shape = x.shape
+  if x.dtype == np.bool_:
+    x = np.packbits(x, bitorder='little')  # type: ignore
+  x = np.ascontiguousarray(x)
+  attr = ir.DenseElementsAttr.get(x, type=element_type, shape=shape)  # type: ignore
   return hlo.constant(attr)
 
 
@@ -339,87 +344,6 @@ def _token_constant_handler(val):
   return hlo.create_token()
 register_constant_handler(core.Token, _token_constant_handler)
 
-# Attributes
-
-AttributeHandler = Callable[[Any], ir.Attribute]
-_attribute_handlers: dict[type[Any], AttributeHandler] = {}
-
-def register_attribute_handler(type_: type[Any], handler_fun: AttributeHandler):
-  _attribute_handlers[type_] = handler_fun
-
-def get_attribute_handler(type_: type[Any]) -> AttributeHandler:
-  return _attribute_handlers[type_]
-
-def _numpy_scalar_attribute(val: Any) -> ir.Attribute:
-  mlir_type = dtype_to_ir_type(val.dtype)
-  if isinstance(mlir_type, ir.IntegerType):
-    return ir.IntegerAttr.get(mlir_type, val)
-  elif isinstance(mlir_type, ir.FloatType):
-    return ir.FloatAttr.get(mlir_type, val)
-  else:
-    raise TypeError(f"Unsupported scalar attribute type: {type(val)}")
-
-def _numpy_array_attribute(x: np.ndarray | np.generic) -> ir.Attribute:
-  element_type = dtype_to_ir_type(x.dtype)
-  shape = x.shape
-  if x.dtype == np.bool_:
-    x = np.packbits(x, bitorder='little')  # type: ignore
-  x = np.ascontiguousarray(x)
-  return ir.DenseElementsAttr.get(x, type=element_type, shape=shape)  # type: ignore
-
-def _numpy_array_attribute_handler(val: np.ndarray | np.generic) -> ir.Attribute:
-  if 0 in val.strides and val.size > 0:
-    raise ValueError(
-        "NumPy arrays with zero strides are not supported as MLIR attributes")
-  if val.dtype == dtypes.float0:
-    val = np.zeros(val.shape, dtype=np.bool_)
-  if dtypes.is_python_scalar(val) or np.isscalar(val):
-    return _numpy_scalar_attribute(val)
-  else:
-    return _numpy_array_attribute(val)
-
-register_attribute_handler(np.ndarray, _numpy_array_attribute_handler)
-
-for _scalar_type in [np.int8, np.int16, np.int32, np.int64,
-                     np.uint8, np.uint16, np.uint32, np.uint64,
-                     np.float16, np.float32, np.float64,
-                     np.complex64, np.complex128,
-                     np.bool_, np.longlong, dtypes.bfloat16]:
-  register_attribute_handler(_scalar_type, _numpy_array_attribute_handler)  # type: ignore
-
-def _python_scalar_attribute_handler(dtype, val):
-  return _numpy_scalar_attribute(np.array(val, dtype))
-
-for ptype, dtype in dtypes.python_scalar_dtypes.items():
-  register_attribute_handler(
-      ptype, partial(_python_scalar_attribute_handler, dtype))
-
-register_attribute_handler(str, ir.StringAttr.get)
-register_attribute_handler(bytes, ir.StringAttr.get)
-
-def _dict_attribute_handler(val: dict[str, Any]) -> ir.Attribute:
-  return ir.DictAttr.get({k: ir_attribute(v) for k, v in val.items()})
-
-register_attribute_handler(dict, _dict_attribute_handler)
-
-def _sequence_attribute_handler(val: Sequence[Any]) -> ir.Attribute:
-  return ir.ArrayAttr.get([ir_attribute(v) for v in val])
-
-register_attribute_handler(list, _sequence_attribute_handler)
-register_attribute_handler(tuple, _sequence_attribute_handler)
-
-def ir_attribute(val: Any) -> ir.Attribute:
-  """Convert a Python value to an MLIR attribute."""
-  for t in type(val).__mro__:
-    handler = _attribute_handlers.get(t)
-    if handler:
-      out = handler(val)
-      assert isinstance(out, ir.Attribute), (type(val), out)
-      return out
-  if hasattr(val, '__jax_array__'):
-    return ir_attribute(val.__jax_array__())
-  raise TypeError(f"No attribute handler defined for type: {type(val)}")
-
 # Source locations
 
 def get_canonical_source_file(file_name: str, caches: TracebackCaches) -> str:
diff --git a/jax/_src/nn/initializers.py b/jax/_src/nn/initializers.py
index 7d228e4be..eb1bb1609 100644
--- a/jax/_src/nn/initializers.py
+++ b/jax/_src/nn/initializers.py
@@ -184,7 +184,7 @@ def truncated_normal(stddev: RealNumeric = 1e-2,
 
   >>> import jax, jax.numpy as jnp
   >>> initializer = jax.nn.initializers.truncated_normal(5.0)
-  >>> initializer(jax.random.PRNGKey(42), (2, 3), jnp.float32)  # doctest: +SKIP
+  >>> initializer(jax.random.key(42), (2, 3), jnp.float32)  # doctest: +SKIP
   Array([[ 2.9047365,  5.2338114,  5.29852  ],
          [-3.836303 , -4.192359 ,  0.6022964]], dtype=float32)
   """
diff --git a/jax/interpreters/mlir.py b/jax/interpreters/mlir.py
index 78b070614..edfc56ddd 100644
--- a/jax/interpreters/mlir.py
+++ b/jax/interpreters/mlir.py
@@ -49,7 +49,6 @@ from jax._src.interpreters.mlir import (
   i32_attr as i32_attr,
   i64_attr as i64_attr,
   ir as ir,
-  ir_attribute as ir_attribute,
   ir_constant as ir_constant,
   ir_type_handlers as ir_type_handlers,
   jaxpr_subcomp as jaxpr_subcomp,


Test patch:
diff --git a/tests/extend_test.py b/tests/extend_test.py
index 098d52b10..45c689239 100644
--- a/tests/extend_test.py
+++ b/tests/extend_test.py
@@ -31,6 +31,8 @@ from jax._src import prng
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax._src.interpreters import mlir
+from jax._src.lib.mlir import ir
+from jax._src.extend import ffi
 
 jax.config.parse_flags_with_absl()
 
@@ -101,6 +103,15 @@ class FfiTest(jtu.JaxTestCase):
     for header in ["c_api.h", "api.h", "ffi.h"]:
       self.assertTrue(os.path.exists(os.path.join(base_dir, header)))
 
+  @parameterized.parameters(
+      [True, int(1), float(5.0),
+       np.int32(-5), np.float32(0.5)])
+  def testIrAttribute(self, value):
+    with mlir.make_ir_context(), ir.Location.unknown():
+      const = mlir.ir_constant(value)
+      attr = ffi._ir_attribute(value)
+      assert const.type.element_type == attr.type
+
   @parameterized.parameters([True, 1, 5.0, "param", np.float32(0.5)])
   def testParams(self, param):
     prim = core.Primitive("test_ffi")