[MLIR] Fix construction of i1 attributes from numpy arrays and add tests.

This change fixes the construction of MLIR attributes from `numpy` arrays for a Boolean `dtype`. This happened when calling `jax.mlir.ir_attribute` and the problem occurred in `_numpy_array_attribute`: that function *first* bit-packed the Booleans and *then* determined the `dtype` from the result, which is `uint8` at that point rather than `bool`. The fix thus consists of a simple swapping of lines. Note that the virtually equivalent `_numpy_array_constant` already used the right order.

The change extends and improves the tests for the FFI primitive, which uses `jax.mlir.ir_attribute` under the hood. The previous tests use a flawed comparison to verify the round-tripped type: it checked whether the *Python* type of the created MLIR attribute is the expected one. However, that is not enough: two `DenseElementAttr`s with different element types pass that check but really should not. That check would not have caught the bug fixed in this change (plus it didn't test the variant with `bool`s). The new version uses strings instead, which is not only more precise but also more concise in building the expected values. In addition to this improvement, the change adds instances of the corresponding FFI test for essentially all attribute types I could find.

PiperOrigin-RevId: 834770477

Author: Google-ML-Automation

jax/_src/interpreters/mlir.py

@@ -364,11 +364,11 @@ def _numpy_scalar_attribute(val: Any) -> ir.Attribute:
     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)
-  element_type = dtype_to_ir_type(x.dtype)
   return ir.DenseElementsAttr.get(x, type=element_type, shape=shape)  # type: ignore
 
 def _numpy_array_attribute_handler(val: np.ndarray | np.generic) -> ir.Attribute:

tests/ffi_test.py

@@ -28,6 +28,7 @@
 from jax._src import config
 from jax._src import core
 from jax._src import dispatch
+from jax._src import dtypes
 from jax._src import test_util as jtu
 from jax._src.interpreters import mlir
 from jax._src.layout import Layout
@@ -85,29 +86,67 @@ def lowering_rule(ctx, x):
     pattern = rf"result_layouts = \[dense<\[{expected}\]>"
     self.assertRegex(text, pattern)
 
-  @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 test_params(self, param, expected_builder):
+  # Concise helpers to every test instance below in one line.
+  _arr = lambda value, dtype=None: np.array(value, dtype=dtype)
+  _ftens1 = lambda et: f"dense<1.000000e+00> : tensor<{et}>"
+  _itens1 = lambda et: f"dense<1> : tensor<{et}>"
+
+  @parameterized.parameters(
+      (_arr(1, dtypes.int2), _itens1("i2")),
+      (_arr(1, dtypes.int4), _itens1("i4")),
+      (_arr(1, dtypes.uint2), _itens1("ui2")),
+      (_arr(1, dtypes.uint4), _itens1("ui4")),
+      (_arr(1, np.int16), _itens1("i16")),
+      (_arr(1, np.int32), _itens1("i32")),
+      (_arr(1, np.int64), _itens1("i64")),
+      (_arr(1, np.int8), _itens1("i8")),
+      (_arr(1, np.uint16), _itens1("ui16")),
+      (_arr(1, np.uint32), _itens1("ui32")),
+      (_arr(1, np.uint64), _itens1("ui64")),
+      (_arr(1, np.uint8), _itens1("ui8")),
+      (_arr(1.0, dtypes.bfloat16), _ftens1("bf16")),
+      (_arr(1.0, dtypes.float4_e2m1fn), _ftens1("f4E2M1FN")),
+      (_arr(1.0, dtypes.float8_e3m4), _ftens1("f8E3M4")),
+      (_arr(1.0, dtypes.float8_e4m3), _ftens1("f8E4M3")),
+      (_arr(1.0, dtypes.float8_e4m3b11fnuz), _ftens1("f8E4M3B11FNUZ")),
+      (_arr(1.0, dtypes.float8_e4m3fn), _ftens1("f8E4M3FN")),
+      (_arr(1.0, dtypes.float8_e4m3fnuz), _ftens1("f8E4M3FNUZ")),
+      (_arr(1.0, dtypes.float8_e5m2), _ftens1("f8E5M2")),
+      (_arr(1.0, dtypes.float8_e5m2fnuz), _ftens1("f8E5M2FNUZ")),
+      (_arr(1.0, dtypes.float8_e8m0fnu), _ftens1("f8E8M0FNU")),
+      (_arr(1.0, np.bool), "dense<true> : tensor<i1>"),
+      (_arr(1.0, np.float16), _ftens1("f16")),
+      (_arr(1.0, np.float32), _ftens1("f32")),
+      (_arr(1.0, np.float64), _ftens1("f64")),
+      (dtypes.bfloat16(1.0), "1.000000e+00 : bf16"),
+      (np.bool(False), "false"),
+      (np.bool(True), "true"),
+      (np.float16(1.0), "1.000000e+00 : f16"),
+      (np.float32(1.0), "1.000000e+00 : f32"),
+      (np.float64(1.0), "1.000000e+00 : f64"),
+      (np.int16(1), "1 : i16"),
+      (np.int32(1), "1 : i32"),
+      (np.int64(1), "1 : i64"),
+      (np.int8(1), "1 : i8"),
+      (np.uint16(1), "1 : ui16"),
+      (np.uint32(1), "1 : ui32"),
+      (np.uint64(1), "1 : ui64"),
+      (np.uint8(1), "1 : ui8"),
+      (np.zeros((), dtype=dtypes.float0), "dense<false> : tensor<i1>"),
+      ("param", '"param"'),
+  )
+  def test_params(self, param, expected_str):
     def fun(x):
       return jax.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")
     op = self.find_custom_call_in_module(module)
-    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))
+    conf = op.attributes["mhlo.backend_config"]
+    self.assertIsInstance(conf, mlir.ir.DictAttr)
+    self.assertIn("param", conf)
+    self.assertEqual(str(conf["param"]), expected_str)
 
   def test_token(self):
     def fun():