Repository: google/jax
Pull request #22988 (2024-08-11 15:06:32)

indent code consistently in auto-parallelization and shard_map tutorials


Patch:
diff --git a/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb b/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb
index 09fbbd3a7..3d8c5b020 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 43c14bc41..cb5d4602c 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 919690d23..b1792a6c0 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 6f9dfbb65..7469bf32b 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/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")