Repository: google/jax
Pull request #22991 (2024-08-11 19:44:34)

emphasize typed over legacy RNG keys in `random` module docs
Update both docstrings and move the `PRNGKey` function listing lower in the API reference.


Patch:
diff --git a/docs/jax.random.rst b/docs/jax.random.rst
index 9d6369d2d..6c5427c05 100644
--- a/docs/jax.random.rst
+++ b/docs/jax.random.rst
@@ -12,13 +12,13 @@ Key Creation & Manipulation
 .. autosummary::
   :toctree: _autosummary
 
-  PRNGKey
   key
   key_data
   wrap_key_data
   fold_in
   split
   clone
+  PRNGKey
 
 Random Samplers
 ~~~~~~~~~~~~~~~
diff --git a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
index 3d8c5b020..09fbbd3a7 100644
--- a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
+++ b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
@@ -1749,20 +1749,20 @@
    "outputs": [],
    "source": [
     "def init_layer(key, n_in, n_out):\n",
-    "  k1, k2 = jax.random.split(key)\n",
-    "  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)\n",
-    "  b = jax.random.normal(k2, (n_out,))\n",
-    "  return W, b\n",
+    "    k1, k2 = jax.random.split(key)\n",
+    "    W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)\n",
+    "    b = jax.random.normal(k2, (n_out,))\n",
+    "    return W, b\n",
     "\n",
     "def init_model(key, layer_sizes, batch_size):\n",
-    "  key, *keys = jax.random.split(key, len(layer_sizes))\n",
-    "  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))\n",
+    "    key, *keys = jax.random.split(key, len(layer_sizes))\n",
+    "    params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))\n",
     "\n",
-    "  key, *keys = jax.random.split(key, 3)\n",
-    "  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))\n",
-    "  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))\n",
+    "    key, *keys = jax.random.split(key, 3)\n",
+    "    inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))\n",
+    "    targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))\n",
     "\n",
-    "  return params, (inputs, targets)\n",
+    "    return params, (inputs, targets)\n",
     "\n",
     "layer_sizes = [784, 8192, 8192, 8192, 10]\n",
     "batch_size = 8192\n",
diff --git a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
index cb5d4602c..43c14bc41 100644
--- a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
+++ b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.md
@@ -631,20 +631,20 @@ gradfun = jax.jit(jax.grad(loss))
 :id: R0x62AIa3vGU
 
 def init_layer(key, n_in, n_out):
-  k1, k2 = jax.random.split(key)
-  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)
-  b = jax.random.normal(k2, (n_out,))
-  return W, b
+    k1, k2 = jax.random.split(key)
+    W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)
+    b = jax.random.normal(k2, (n_out,))
+    return W, b
 
 def init_model(key, layer_sizes, batch_size):
-  key, *keys = jax.random.split(key, len(layer_sizes))
-  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))
+    key, *keys = jax.random.split(key, len(layer_sizes))
+    params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))
 
-  key, *keys = jax.random.split(key, 3)
-  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))
-  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))
+    key, *keys = jax.random.split(key, 3)
+    inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))
+    targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))
 
-  return params, (inputs, targets)
+    return params, (inputs, targets)
 
 layer_sizes = [784, 8192, 8192, 8192, 10]
 batch_size = 8192
diff --git a/docs/notebooks/shard_map.ipynb b/docs/notebooks/shard_map.ipynb
index b1792a6c0..919690d23 100644
--- a/docs/notebooks/shard_map.ipynb
+++ b/docs/notebooks/shard_map.ipynb
@@ -1483,20 +1483,20 @@
    "outputs": [],
    "source": [
     "def init_layer(key, n_in, n_out):\n",
-    "  k1, k2 = jax.random.split(key)\n",
-    "  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)\n",
-    "  b = jax.random.normal(k2, (n_out,))\n",
-    "  return W, b\n",
+    "    k1, k2 = jax.random.split(key)\n",
+    "    W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)\n",
+    "    b = jax.random.normal(k2, (n_out,))\n",
+    "    return W, b\n",
     "\n",
     "def init(key, layer_sizes, batch_size):\n",
-    "  key, *keys = jax.random.split(key, len(layer_sizes))\n",
-    "  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))\n",
+    "    key, *keys = jax.random.split(key, len(layer_sizes))\n",
+    "    params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))\n",
     "\n",
-    "  key, *keys = jax.random.split(key, 3)\n",
-    "  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))\n",
-    "  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))\n",
+    "    key, *keys = jax.random.split(key, 3)\n",
+    "    inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))\n",
+    "    targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))\n",
     "\n",
-    "  return params, (inputs, targets)"
+    "    return params, (inputs, targets)"
    ]
   },
   {
diff --git a/docs/notebooks/shard_map.md b/docs/notebooks/shard_map.md
index 7469bf32b..6f9dfbb65 100644
--- a/docs/notebooks/shard_map.md
+++ b/docs/notebooks/shard_map.md
@@ -1035,20 +1035,20 @@ def loss(params, batch):
 
 ```{code-cell}
 def init_layer(key, n_in, n_out):
-  k1, k2 = jax.random.split(key)
-  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)
-  b = jax.random.normal(k2, (n_out,))
-  return W, b
+    k1, k2 = jax.random.split(key)
+    W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)
+    b = jax.random.normal(k2, (n_out,))
+    return W, b
 
 def init(key, layer_sizes, batch_size):
-  key, *keys = jax.random.split(key, len(layer_sizes))
-  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))
+    key, *keys = jax.random.split(key, len(layer_sizes))
+    params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))
 
-  key, *keys = jax.random.split(key, 3)
-  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))
-  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))
+    key, *keys = jax.random.split(key, 3)
+    inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))
+    targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))
 
-  return params, (inputs, targets)
+    return params, (inputs, targets)
 ```
 
 ```{code-cell}
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/random.py b/jax/_src/random.py
index 113bcc450..6105d56f9 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -197,9 +197,10 @@ def key(seed: int | ArrayLike, *,
         impl: PRNGSpecDesc | None = None) -> KeyArray:
   """Create a pseudo-random number generator (PRNG) key given an integer seed.
 
-  The result is a scalar array with a key that indicates the default PRNG
-  implementation, as determined by the optional ``impl`` argument or,
-  otherwise, by the ``jax_default_prng_impl`` config flag.
+  The result is a scalar array containing a key, whose dtype indicates
+  the default PRNG implementation, as determined by the optional
+  ``impl`` argument or, otherwise, by the ``jax_default_prng_impl``
+  config flag at the time when this function is called.
 
   Args:
     seed: a 64- or 32-bit integer used as the value of the key.
@@ -214,11 +215,20 @@ def key(seed: int | ArrayLike, *,
 
 def PRNGKey(seed: int | ArrayLike, *,
             impl: PRNGSpecDesc | None = None) -> KeyArray:
-  """Create a pseudo-random number generator (PRNG) key given an integer seed.
-
-  The resulting key carries the default PRNG implementation, as
-  determined by the optional ``impl`` argument or, otherwise, by the
-  ``jax_default_prng_impl`` config flag.
+  """Create a legacy PRNG key given an integer seed.
+
+  This function produces old-style legacy PRNG keys, which are arrays
+  of dtype ``uint32``. For more, see the note in the `PRNG keys
+  <https://jax.readthedocs.io/en/latest/jax.random.html#prng-keys>`_
+  section. When possible, :func:`jax.random.key` is recommended for
+  use instead.
+
+  The resulting key does not carry a PRNG implementation. The returned
+  key matches the implementation given by the optional ``impl``
+  argument or, otherwise, determined by the ``jax_default_prng_impl``
+  config flag. Callers must ensure that same implementation is set as
+  the default when passing this key as an argument to other functions
+  (such as ``jax.random.split`` and ``jax.random.normal``).
 
   Args:
     seed: a 64- or 32-bit integer used as the value of the key.
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 d43063c7a..45c689239 100644
--- a/tests/extend_test.py
+++ b/tests/extend_test.py
@@ -25,11 +25,14 @@ import jax.numpy as jnp
 
 from jax._src import abstract_arrays
 from jax._src import api
+from jax._src import core
 from jax._src import linear_util
 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()
 
@@ -100,34 +103,25 @@ 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, mlir.ir.BoolAttr.get),
-      (1, mlir.i64_attr),
-      (5.0, lambda x: mlir.ir.FloatAttr.get(mlir.ir.F64Type.get(), x)),
-      ("param", mlir.ir.StringAttr.get),
-      (np.float32(0.5),
-       lambda x: mlir.ir.FloatAttr.get(mlir.ir.F32Type.get(), x)),
-  ])
-  def testParams(self, param, expected_builder):
-    def fun(x):
-      return jex.ffi.ffi_call("test_ffi", x, x, param=param)
-
-    # Here we inspect the lowered IR to test that the parameter has been
-    # serialized with the appropriate type.
-    module = jax.jit(fun).lower(0.5).compiler_ir("stablehlo")
-    for func in module.body.operations:
-      for block in func.body.blocks:
-        for op in block.operations:
-          if op.OPERATION_NAME == "stablehlo.custom_call":
-            config = op.attributes["mhlo.backend_config"]
-            self.assertIsInstance(config, mlir.ir.DictAttr)
-            self.assertIn("param", config)
-            with mlir.make_ir_context(), mlir.ir.Location.unknown():
-              expected = expected_builder(param)
-            self.assertEqual(type(config["param"]), type(expected))
-            self.assertTrue(expected.type.isinstance(config["param"].type))
-            return
-    self.fail("No custom_call found in the lowered IR")
+  @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")
+    prim.def_abstract_eval(lambda *args, **kwargs: args[0])
+    mlir.register_lowering(prim, jex.ffi.ffi_lowering("test_ffi"))
+
+    # TODO(dfm): Currently testing that lowering works with different types of
+    # parameters, but we should probably actually check the emitted HLO.
+    func = jax.jit(lambda *args: prim.bind(*args, param=param))
+    func.lower(jnp.linspace(0, 5, 10))
 
   @jtu.sample_product(
     shape=[(1,), (4,), (5,)],