Repository: google/jax
Pull request #22263 (2024-07-03 20:39:24)

Don't wrap singleton ir.Types during HLO lowering.
This is similar to https://github.com/google/jax/pull/22211, but for MLIR types instead of MLIR values.


Patch:
diff --git a/jax/_src/ad_checkpoint.py b/jax/_src/ad_checkpoint.py
index c7c194bff..2184be87d 100644
--- a/jax/_src/ad_checkpoint.py
+++ b/jax/_src/ad_checkpoint.py
@@ -777,10 +777,11 @@ def _remat_lowering(ctx, *args, jaxpr: core.Jaxpr, prevent_cse: bool,
           differentiated=differentiated, policy=policy,
           is_gpu_platform=is_gpu_platform)
 
-    arg_types = map(mlir.aval_to_ir_types, ctx.avals_in)
+    arg_types = map(mlir.aval_to_ir_type, ctx.avals_in)
     flat_args = mlir.flatten_ir_values(args)
     barrier_op = hlo.OptimizationBarrierOp(flat_args)
-    jaxpr_args = mlir.unflatten_ir_values(barrier_op.results, map(len, arg_types))
+    jaxpr_args = mlir.unflatten_ir_values_like_types(
+      barrier_op.results, arg_types)
   else:
     jaxpr_args = args
   outs, tokens_out = mlir.jaxpr_subcomp(
@@ -797,10 +798,10 @@ def _optimization_barrier_abstract_eval(*args):
   return args
 
 def _optimization_barrier_lowering_rule(ctx, *args):
-  barrier_types = map(mlir.aval_to_ir_types, ctx.avals_in)
+  barrier_types = map(mlir.aval_to_ir_type, ctx.avals_in)
   flat_args = mlir.flatten_ir_values(args)
   barrier_op = hlo.OptimizationBarrierOp(flat_args)
-  return mlir.unflatten_ir_values(barrier_op.results, map(len, barrier_types))
+  return mlir.unflatten_ir_values_like_types(barrier_op.results, barrier_types)
 
 def _optimization_barrier(arg):
   flat_args, treedef = tree_flatten(arg)
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
index 799d764ea..2213e1f01 100644
--- a/jax/_src/export/_export.py
+++ b/jax/_src/export/_export.py
@@ -737,7 +737,7 @@ def _wrap_main_func(
     orig_main_name = ir.StringAttr(symbol_table.insert(orig_main)).value
 
     def is_token(typ, attrs):
-      return (typ == mlir.token_type()[0])
+      return (typ == mlir.token_type())
 
     orig_input_types = orig_main.type.inputs  # type: ignore
     arg_attrs = list(ir.ArrayAttr(orig_main.arg_attrs))  # type: ignore
@@ -1329,7 +1329,7 @@ def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
     else:
       current_platform_idx = cast(ir.Value, mlir.ir_constant(np.int32(0)))
     # Compute the rule index based on the current platform
-    i32_type = mlir.aval_to_ir_types(core.ShapedArray((), dtype=np.int32))[0]
+    i32_type = mlir.aval_to_ir_type(core.ShapedArray((), dtype=np.int32))
     if current_platform_idx.type != i32_type:
       current_platform_idx = hlo.ConvertOp(i32_type, current_platform_idx)
     callee_platform_idx = hlo.CaseOp([i32_type],
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index f3e5a3f2a..e50be6592 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -87,16 +87,6 @@ lowerable_effects: effects_lib.EffectTypeSet = effects_lib.lowerable_effects
 
 # IR Helpers
 
-IrValues = Union[ir.Value, tuple[ir.Value, ...]]
-
-def _is_ir_values(x: IrValues) -> bool:
-  """Returns true if `x` is an ir.Value or tuple of ir.Values"""
-  if isinstance(x, ir.Value):
-    return True
-  return (isinstance(x, tuple) and len(x) != 1
-          and all(isinstance(v, ir.Value) for v in x))
-
-
 def dense_int_elements(xs) -> ir.DenseIntElementsAttr:
   return type_cast(ir.DenseIntElementsAttr,
                    ir.DenseIntElementsAttr.get(np.asarray(xs, np.int64)))
@@ -148,6 +138,17 @@ def delegate_lowering(ctx, lowering_fun, *args, **ctx_override_kwargs):
 
 # IR Types
 
+IrTypes = Union[ir.Type, tuple[ir.Type, ...]]
+IrValues = Union[ir.Value, tuple[ir.Value, ...]]
+
+def _is_ir_values(x: IrValues) -> bool:
+  """Returns true if `x` is an ir.Value or tuple of ir.Values"""
+  if isinstance(x, ir.Value):
+    return True
+  return (isinstance(x, tuple) and len(x) != 1
+          and all(isinstance(v, ir.Value) for v in x))
+
+
 # Non-canonicalized dtype to IR type mapping.
 _dtype_to_ir_type : dict[np.dtype, Callable[[], ir.Type]] = {
   np.dtype(dtypes.float0): partial(ir.IntegerType.get_signless, 1),
@@ -200,26 +201,24 @@ def dtype_to_ir_type(dtype: core.bint | np.dtype | np.generic) -> ir.Type:
         f"No dtype_to_ir_type handler for dtype: {dtype}") from err
   return ir_type_factory()
 
-def _array_ir_types(aval: core.ShapedArray | core.DShapedArray
-                    ) -> Sequence[ir.Type]:
+def _array_ir_types(aval: core.ShapedArray | core.DShapedArray) -> ir.Type:
   aval = core.physical_aval(aval)  # type: ignore
   if not core.is_constant_shape(aval.shape):
     return _dynamic_array_ir_types(aval)  # type: ignore
-  return (ir.RankedTensorType.get(aval.shape, dtype_to_ir_type(aval.dtype)),)  # type: ignore
+  return ir.RankedTensorType.get(aval.shape, dtype_to_ir_type(aval.dtype))  # type: ignore
 
-def _dynamic_array_ir_types(aval: core.ShapedArray) -> Sequence[ir.Type]:
+def _dynamic_array_ir_types(aval: core.ShapedArray) -> ir.Type:
   dyn_size = ir.ShapedType.get_dynamic_size()
   shape = [d if type(d) is int else dyn_size for d in aval.shape]
-  return (ir.RankedTensorType.get(shape, dtype_to_ir_type(aval.dtype)),)
+  return ir.RankedTensorType.get(shape, dtype_to_ir_type(aval.dtype))
 
-ir_type_handlers: dict[type[core.AbstractValue],
-                        Callable[[Any], Sequence[ir.Type]]] = {}
+ir_type_handlers: dict[type[core.AbstractValue], Callable[[Any], IrTypes]] = {}
 
-def aval_to_ir_types(aval: core.AbstractValue) -> Sequence[ir.Type]:
+def aval_to_ir_type(aval: core.AbstractValue) -> IrTypes:
   """Converts a JAX aval to zero or more MLIR IR types.
 
   In general, a JAX value may be represented by multiple IR values, so this
-  function returns multiple types."""
+  function may return a tuple of types."""
   try:
     return ir_type_handlers[type(aval)](aval)
   except KeyError as err:
@@ -227,20 +226,9 @@ def aval_to_ir_types(aval: core.AbstractValue) -> Sequence[ir.Type]:
 
 ir_type_handlers[core.ShapedArray] = _array_ir_types
 ir_type_handlers[core.ConcreteArray] = _array_ir_types
-ir_type_handlers[core.AbstractToken] = lambda _: [hlo.TokenType.get()]
+ir_type_handlers[core.AbstractToken] = lambda _: hlo.TokenType.get()
 ir_type_handlers[core.DShapedArray] = _dynamic_array_ir_types
 
-def aval_to_ir_type(aval: core.AbstractValue) -> ir.Type:
-  """Convenience wrapper around aval_to_ir_types for single types.
-
-  For some common cases, e.g. dense arrays, we know JAX values are represented
-  by a single IR value."""
-  types = aval_to_ir_types(aval)
-  if len(types) != 1:
-    raise TypeError(f"aval_to_ir_type called on {aval} which corresponds to "
-                    f"multiple IR types {types}")
-  return types[0]
-
 
 # Constants
 
@@ -750,19 +738,49 @@ def flatten_ir_values(xs: Iterable[IrValues]) -> list[ir.Value]:
 
 _unflatten_done = object()
 
-def unflatten_ir_values(xs: Iterable[ir.Value], ns: Sequence[int]) -> list[IrValues]:
+def _unwrap_singleton_ir_values(x): return x[0] if len(x) == 1 else x
+
+
+def flatten_ir_types(xs: Iterable[IrTypes]) -> list[ir.Type]:
+  """Concatenates/flattens a list of ir.Types or ir.Type sequences."""
+  out = []
+  for x in xs:
+    if isinstance(x, ir.Type):
+      out.append(x)
+    else:
+      out.extend(x)
+  return out
+
+_unflatten_done = object()
+
+def unflatten_ir_types(xs: Iterable[ir.Type], ns: Sequence[int]) -> list[IrTypes]:
   """Splits `xs` into subsequences of lengths `ns`.
 
   Unlike `split_list`, the `sum(ns)` must be equal to `len(xs)`, and if n == 1
-  then values are not wrapped in a singleton list."""
+  then types are not wrapped in a singleton list."""
   xs_iter = iter(xs)
-  unflattened: list[IrValues]
+  unflattened: list[IrTypes]
   unflattened = [next(xs_iter) if n == 1 else tuple(next(xs_iter)
                  for _ in range(n)) for n in ns]
   assert next(xs_iter, _unflatten_done) is _unflatten_done
   return unflattened
 
-def _unwrap_singleton_ir_values(x): return x[0] if len(x) == 1 else x
+def len_ir_types(x: IrTypes) -> int:
+  return 1 if isinstance(x, ir.Type) else len(x)
+
+def unflatten_ir_values_like_types(xs: Iterable[ir.Value],
+                                   ys: Sequence[IrTypes]) -> list[IrValues]:
+  """Splits `xs` into subsequences of lengths `ns`.
+
+  Unlike `split_list`, the `sum(ns)` must be equal to `len(xs)`, and if n == 1
+  then values are not wrapped in a singleton list."""
+  xs_iter = iter(xs)
+  unflattened: list[IrValues]
+  unflattened = [next(xs_iter) if isinstance(y, ir.Type) else
+                 tuple(next(xs_iter) for _ in range(len(y)))
+                 for y in ys]
+  assert next(xs_iter, _unflatten_done) is _unflatten_done
+  return unflattened
 
 
 _module_name_regex = re.compile(r"[^\w.-]")
@@ -1066,12 +1084,8 @@ def _set_up_aliases(input_output_aliases, avals_in, avals_out, donated_args,
   return input_output_aliases, out_donated_args
 
 Token = ir.Value
-
-def token_type() -> Sequence[ir.Type]:
-  return [hlo.TokenType.get()]
-
-def create_token() -> Token:
-  return hlo.create_token()
+token_type = hlo.TokenType.get
+create_token = hlo.create_token
 
 class TokenSet:
   """An immutable container of tokens to be used to lower effectful jaxprs. When lowering
@@ -1172,11 +1186,11 @@ def lower_jaxpr_to_fun(
   # The first dimension variable may be the platform index
   num_dim_vars = len(ctx.shape_poly_state.dim_vars)
   dim_var_avals = [core.ShapedArray((), dtypes.canonicalize_dtype(np.int64))] * num_dim_vars
-  dim_var_types = map(aval_to_ir_types, dim_var_avals)
+  dim_var_types = map(aval_to_ir_type, dim_var_avals)
 
   # Function inputs: *dim_var_values, *tokens, *actual_inputs
-  input_types = map(aval_to_ir_types, jaxpr.in_avals)
-  output_types = map(aval_to_ir_types, jaxpr.out_avals)
+  input_types = map(aval_to_ir_type, jaxpr.in_avals)
+  output_types = map(aval_to_ir_type, jaxpr.out_avals)
   num_tokens = len(effects)
 
   token_types = [token_type() for _ in effects]
@@ -1219,8 +1233,8 @@ def lower_jaxpr_to_fun(
   if xla_donated_args is not None:
     xla_donated_args = [*([False] * (num_dim_vars + num_tokens)), *xla_donated_args]
 
-  flat_input_types = util.flatten(input_types)
-  flat_output_types = util.flatten(output_types)
+  flat_input_types = flatten_ir_types(input_types)
+  flat_output_types = flatten_ir_types(output_types)
   ftype = ir.FunctionType.get(flat_input_types, flat_output_types)
   func_op = func_dialect.FuncOp(name, ftype, ip=ctx.ip)
   func_op.attributes["sym_visibility"] = ir.StringAttr.get(
@@ -1230,29 +1244,29 @@ def lower_jaxpr_to_fun(
   ir_arg_shardings = None
   if arg_shardings is not None:
     ir_arg_shardings = util.flatten(
-        [[_to_physical_op_sharding(a, s)] * len(types)
+        [[_to_physical_op_sharding(a, s)] * len_ir_types(types)
          for a, s, types in zip(input_avals, arg_shardings, input_types)])
 
   ir_arg_memory_kinds = None
   if arg_memory_kinds is not None:
     ir_arg_memory_kinds = util.flatten(
-        [[mk] * len(types) for mk, types in zip(arg_memory_kinds, input_types)])
+        [[mk] * len_ir_types(types) for mk, types in zip(arg_memory_kinds, input_types)])
 
   ir_arg_layouts = None
   if arg_layouts is not None:
     ir_arg_layouts = util.flatten(
-        [[_to_xla_layout(l, a)] * len(types)
+        [[_to_xla_layout(l, a)] * len_ir_types(types)
          for l, a, types in zip(arg_layouts, input_avals, input_types)])
 
   ir_donated_args = None
   if xla_donated_args is not None:
     ir_donated_args = util.flatten(
-        [[is_donated] * len(types) for is_donated, types in zip(xla_donated_args, input_types)])
+        [[is_donated] * len_ir_types(types) for is_donated, types in zip(xla_donated_args, input_types)])
 
   ir_result_shardings = None
   if result_shardings is not None:
     ir_result_shardings = util.flatten(
-        [[_to_physical_op_sharding(a, s)] * len(types)
+        [[_to_physical_op_sharding(a, s)] * len_ir_types(types)
          for a, s, types in zip(output_avals, result_shardings, output_types)])
 
   ir_result_memory_kinds = None
@@ -1264,20 +1278,20 @@ def lower_jaxpr_to_fun(
     for pom, mk, types in zip(propagated_out_mem_kinds, result_memory_kinds,
                               output_types):
       if pom is not None and mk is None:
-        res.append([pom] * len(types))
+        res.append([pom] * len_ir_types(types))
       else:
-        res.append([mk] * len(types))  # type: ignore
+        res.append([mk] * len_ir_types(types))  # type: ignore
       # To add the custom call on the output to signal a transfer, only do it
       # if memory kind comes from out_shardings on `jit` and result_memory_kinds
       # comes from out_shardings on `jit`.
-      custom_call_res.append([mk] * len(types))
+      custom_call_res.append([mk] * len_ir_types(types))
     ir_result_memory_kinds = util.flatten(res)
     custom_call_ir_result_memory_kinds = util.flatten(custom_call_res)
 
   ir_result_layouts = None
   if result_layouts is not None:
     ir_result_layouts = util.flatten(
-        [[_to_xla_layout(l, a)] * len(types)
+        [[_to_xla_layout(l, a)] * len_ir_types(types)
          for l, a, types in zip(result_layouts, output_avals, output_types)])
 
   if (
@@ -1295,7 +1309,7 @@ def lower_jaxpr_to_fun(
         {} for _ in range(len(flat_input_types))]
 
     if replicated_args is not None:
-      replicated_ir_args = [[replicated] * len(types) for replicated, types
+      replicated_ir_args = [[replicated] * len_ir_types(types) for replicated, types
                             in zip(replicated_args, input_types)]
       for attrs, replicated in zip(arg_attrs, util.flatten(replicated_ir_args)):
         if replicated:
@@ -1322,12 +1336,12 @@ def lower_jaxpr_to_fun(
           attrs["jax.buffer_donor"] = ir.BoolAttr.get(True)
 
     if input_output_aliases is not None:
-      output_ids = util.unflatten(list(range(len(flat_output_types))),
-                                  map(len, output_types))
+      output_ids = util.unflatten(
+        list(range(len(flat_output_types))), map(len_ir_types, output_types))
       aliases: list[int | None] = []
       for itypes, alias in zip(input_types, input_output_aliases):
         if alias is None:
-          aliases.extend([None] * len(itypes))
+          aliases.extend([None] * len_ir_types(itypes))
         else:
           aliases.extend(output_ids[alias])
 
@@ -1417,7 +1431,7 @@ def lower_jaxpr_to_fun(
       ]
 
     _, token_args, unflattened_args = util.split_list(
-        unflatten_ir_values(flat_args, map(len, input_types)),
+        unflatten_ir_values_like_types(flat_args, input_types),
         [num_dim_vars, num_tokens])
     tokens_in = TokenSet(zip(effects, token_args))
     args: list[IrValues] = unflattened_args
@@ -1467,7 +1481,9 @@ def wrap_with_memory_kind(
   if aval_out is None:
     result_type = x.type
   else:
-    result_type = aval_to_ir_type(aval_out)
+    typ = aval_to_ir_type(aval_out)
+    assert isinstance(typ, ir.Type), typ
+    result_type = typ
   op = custom_call("annotate_device_placement", result_types=[result_type],
                    operands=[x], has_side_effect=True, api_version=1)
   dict_attr = {"_xla_buffer_placement": ir.StringAttr.get(memory_kind)}
@@ -1493,17 +1509,19 @@ def _emit_lowering_rule_as_fun(lowering_rule,
   """Emits the contents of a lowering rule as a private function."""
   num_dim_vars = len(ctx.module_context.shape_poly_state.dim_vars)
   # TODO(necula) maybe only pass the dim_vars if they are needed?
-  dim_var_types = map(aval_to_ir_types, [core.ShapedArray((), dtypes.canonicalize_dtype(np.int64))] * num_dim_vars)
+  dim_var_types = [
+    aval_to_ir_type(core.ShapedArray((), dtypes.canonicalize_dtype(np.int64)))
+  ] * num_dim_vars
 
-  input_types = map(aval_to_ir_types, ctx.avals_in)
-  output_types = map(aval_to_ir_types, ctx.avals_out)
+  input_types = map(aval_to_ir_type, ctx.avals_in)
+  output_types = map(aval_to_ir_type, ctx.avals_out)
   effs = list(ctx.tokens_in.effects())
   token_types = [token_type() for _ in effs]
   input_types = [*dim_var_types, *token_types, *input_types]
   output_types = [*token_types, *output_types]
 
-  flat_input_types = util.flatten(input_types)
-  flat_output_types = util.flatten(output_types)
+  flat_input_types = flatten_ir_types(input_types)
+  flat_output_types = flatten_ir_types(output_types)
   ftype = ir.FunctionType.get(flat_input_types, flat_output_types)
   assert ctx.primitive is not None
   func_op = func_dialect.FuncOp(ctx.primitive.name, ftype,
@@ -1512,8 +1530,8 @@ def _emit_lowering_rule_as_fun(lowering_rule,
   ctx.module_context.symbol_table.insert(func_op)
   entry_block = func_op.add_entry_block()
   with ir.InsertionPoint(entry_block):
-    unflattened_args = unflatten_ir_values(entry_block.arguments,
-                                           map(len, input_types))
+    unflattened_args = unflatten_ir_values_like_types(
+      entry_block.arguments, input_types)
     dim_var_values, token_args, unflattened_args = util.split_list(unflattened_args, [num_dim_vars, len(ctx.tokens_in)])
     sub_ctx = ctx.replace(tokens_in=TokenSet(zip(effs, token_args)),
                           dim_var_values=dim_var_values)
@@ -1749,7 +1767,7 @@ def lower_per_platform(ctx: LoweringRuleContext,
   # The first dim_var_values is the platform index
   current_platform_idx = ctx.dim_var_values[0]
   # Compute the rule index based on the current platform
-  i32_type = aval_to_ir_types(core.ShapedArray((), dtype=np.int32))[0]
+  i32_type = aval_to_ir_type(core.ShapedArray((), dtype=np.int32))
   if current_platform_idx.type != i32_type:
     current_platform_idx = hlo.convert(i32_type, current_platform_idx)
   rule_idx_op = hlo.CaseOp([i32_type],
@@ -1761,8 +1779,8 @@ def lower_per_platform(ctx: LoweringRuleContext,
       hlo.return_([ir_constant(np.int32(platform_to_kept_rules_idx[p]))])
   ordered_effects = effects_lib.ordered_effects.filter_in(effects)
   rule_out_avals = [core.abstract_token] * len(ordered_effects) + ctx.avals_out
-  output_types = map(aval_to_ir_types, rule_out_avals)
-  case_op = hlo.CaseOp(util.flatten(output_types),
+  output_types = map(aval_to_ir_type, rule_out_avals)
+  case_op = hlo.CaseOp(flatten_ir_types(output_types),
                       index=rule_idx_op,
                       num_branches=len(kept_rules))
   for i, rule in enumerate(kept_rules):
@@ -1788,7 +1806,7 @@ def lower_per_platform(ctx: LoweringRuleContext,
   results = case_op.results
   if ordered_effects:
     tokens, results = util.split_list(
-      unflatten_ir_values(results, map(len, output_types)),
+      unflatten_ir_values_like_types(results, output_types),
       [len(ordered_effects)])
     tokens_out = ctx.tokens_in.update_tokens(TokenSet(zip(ordered_effects,
                                                           tokens)))
@@ -1897,9 +1915,9 @@ def call_lowering(fn_name, name_stack, call_jaxpr, backend,
     call_jaxpr = pe.close_jaxpr(call_jaxpr)
   check_backend_matches(backend, ctx.platforms)
   effects = list(tokens_in.effects())
-  output_types = map(aval_to_ir_types, avals_out)
+  output_types = map(aval_to_ir_type, avals_out)
   output_types = [token_type()] * len(effects) + output_types
-  flat_output_types = util.flatten(output_types)
+  flat_output_types = flatten_ir_types(output_types)
   symbol_name = _lower_jaxpr_to_fun_cached(
       ctx, fn_name, call_jaxpr, effects, name_stack, arg_names=arg_names,
       result_names=result_names).name.value
@@ -1908,7 +1926,7 @@ def call_lowering(fn_name, name_stack, call_jaxpr, backend,
   call = func_dialect.CallOp(flat_output_types,
                              ir.FlatSymbolRefAttr.get(symbol_name),
                              flatten_ir_values(args))
-  out_nodes = unflatten_ir_values(call.results, map(len, output_types))
+  out_nodes = unflatten_ir_values_like_types(call.results, output_types)
   tokens, out_nodes = util.split_list(out_nodes, [len(effects)])
   tokens_out = tokens_in.update_tokens(TokenSet(zip(effects, tokens)))
   return out_nodes, tokens_out
@@ -2186,6 +2204,7 @@ def _wrap_with_spmd_op(name: str,
   else:
     backend_config = ""
   result_type = aval_to_ir_type(aval_out)
+  assert isinstance(result_type, ir.Type), result_type
   out_shape = core.physical_aval(aval_out).shape  # type: ignore
   if core.is_constant_shape(out_shape):
     result_shapes = None
@@ -2225,6 +2244,7 @@ def wrap_with_layout_op(ctx: LoweringRuleContext,
                         layout: DeviceLocalLayout,
                         aval_in: core.AbstractValue):
   result_type = aval_to_ir_type(aval_out)
+  assert isinstance(result_type, ir.Type), result_type
   out_shape = core.physical_aval(aval_out).shape  # type: ignore
   if core.is_constant_shape(out_shape):
     result_shapes = None
@@ -2270,13 +2290,13 @@ def cache_lowering(f):
       func = _emit_lowering_rule_as_fun(partial(f, **params), ctx)
       ctx.module_context.cached_primitive_lowerings[key] = func
 
-    output_types = map(aval_to_ir_types, ctx.avals_out)
+    output_types = map(aval_to_ir_type, ctx.avals_out)
     args = tuple(ctx.dim_var_values) + args
-    flat_output_types = util.flatten(output_types)
+    flat_output_types = flatten_ir_types(output_types)
     call = func_dialect.CallOp(flat_output_types,
                                ir.FlatSymbolRefAttr.get(func.name.value),
                                flatten_ir_values(args))
-    return unflatten_ir_values(call.results, map(len, output_types))
+    return unflatten_ir_values_like_types(call.results, output_types)
   return cached_lowering
 
 
@@ -2381,8 +2401,8 @@ def xla_fallback_lowering(prim: core.Primitive):
     callee_name = merge_mlir_modules(
         module_ctx.module, f"xla_fallback_{prim.name}", xla_module,
         dst_symtab=module_ctx.symbol_table)
-    output_types = map(aval_to_ir_types, ctx.avals_out)
-    flat_output_types = util.flatten(output_types)
+    output_types = map(aval_to_ir_type, ctx.avals_out)
+    flat_output_types = flatten_ir_types(output_types)
     output_type = (ir.TupleType.get_tuple(flat_output_types)
                    if prim.multiple_results else flat_output_types[0])
 
@@ -2394,7 +2414,7 @@ def xla_fallback_lowering(prim: core.Primitive):
     flat_results = [hlo.get_tuple_element(call, i32_attr(i))
                     for i in range(len(flat_output_types))]
 
-    return unflatten_ir_values(flat_results, map(len, output_types))
+    return unflatten_ir_values_like_types(flat_results, output_types)
   return fallback
 
 
@@ -2594,7 +2614,7 @@ def emit_python_callback(
         for result_aval in result_avals]
     return outputs, token, None
 
-  result_types = util.flatten([aval_to_ir_types(aval) for aval in result_avals])
+  result_types = flatten_ir_types([aval_to_ir_type(aval) for aval in result_avals])
   if token:
 
     callback_without_token = _wrapped_callback
@@ -2608,7 +2628,7 @@ def emit_python_callback(
         xla.aval_to_xla_shapes(core.abstract_token)[0], *result_shapes
     ]
     operands = [token, *operands]
-    result_types = [token_type()[0], *result_types]
+    result_types = [token_type(), *result_types]
     operand_mlir_layouts = [_layout_to_mlir_layout(None), *operand_mlir_layouts]
     result_mlir_layouts = [_layout_to_mlir_layout(None), *result_mlir_layouts]
   callback_descriptor, ifrt_callback = (
@@ -2789,7 +2809,7 @@ def reduce_window(
     window_dimensions, window_strides, padding, base_dilation, window_dilation):
   """Builds a ReduceWindowOp, with support for dynamic shapes."""
 
-  scalar_types = [aval_to_ir_type(aval) for aval in init_values_avals]
+  scalar_types = flatten_ir_types([aval_to_ir_type(aval) for aval in init_values_avals])
   if any(not core.is_constant_shape(s)
          for s in [window_dimensions, window_dilation, window_strides, base_dilation, *padding]):
     # d_padding will be an array i32[N, 2] with pad_lo and pad_hi for each
@@ -2800,8 +2820,8 @@ def reduce_window(
       return hlo.reshape(int2d, pads)
     d_padding = hlo.concatenate(list(map(prep_one_pad, padding)), i64_attr(0))
     # Build the reducer
-    reducer_type = ir.FunctionType.get(scalar_types + scalar_types,
-                                       scalar_types)
+    reducer_type = ir.FunctionType.get(
+      scalar_types + scalar_types, scalar_types)
     with ir.InsertionPoint.at_block_begin(ctx.module_context.module.body):
       reducer = func_dialect.FuncOp(reducer_name, reducer_type)
     ctx.module_context.symbol_table.insert(reducer)
@@ -2811,7 +2831,7 @@ def reduce_window(
 
     rw = custom_call(
       "stablehlo.dynamic_reduce_window",
-      result_types=list(map(aval_to_ir_type, out_avals)),
+      result_types=flatten_ir_types(map(aval_to_ir_type, out_avals)),
       operands=[
         *operands, *init_values,
         eval_dynamic_shape_as_tensor(ctx, window_dimensions),
diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
index b78d7da6e..816160680 100644
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -830,8 +830,8 @@ def _cond_lowering(ctx, index, *args, branches):
   tokens_in = ctx.tokens_in.subset(ordered_effects)
   output_token_types = [mlir.token_type() for _ in ordered_effects]
   output_types = [
-      *output_token_types, *map(mlir.aval_to_ir_types, ctx.avals_out)]
-  flat_output_types = util.flatten(output_types)
+      *output_token_types, *map(mlir.aval_to_ir_type, ctx.avals_out)]
+  flat_output_types = mlir.flatten_ir_types(output_types)
 
   # CaseOp takes a single argument 'index' and the corresponding blocks
   # have no arguments; the computation within the block uses implicit
@@ -851,7 +851,8 @@ def _cond_lowering(ctx, index, *args, branches):
       out_vals = [*out_tokens, *out_vals]
       hlo.return_(mlir.flatten_ir_values(out_vals))
 
-  tokens_and_outputs = mlir.unflatten_ir_values(case_op.results, map(len, output_types))
+  tokens_and_outputs = mlir.unflatten_ir_values_like_types(
+    case_op.results, output_types)
   tokens, outputs = util.split_list(tokens_and_outputs, [num_tokens])
   ctx.set_tokens_out(mlir.TokenSet(zip(ordered_effects, tokens)))
   return outputs
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index 7b3e716ef..aa707386c 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -1713,13 +1713,13 @@ def _while_lowering(ctx, *args, cond_jaxpr, body_jaxpr, cond_nconsts,
       return out
     return mlir.lower_fun(fun)(ctx, *args)
 
-  loop_carry_types = _map(mlir.aval_to_ir_types, ctx.avals_in)
+  loop_carry_types = _map(mlir.aval_to_ir_type, ctx.avals_in)
   body_effects = effects.ordered_effects.filter_in(body_jaxpr.effects)
   num_tokens = len(body_effects)
   tokens = [ctx.tokens_in.get(eff) for eff in body_effects]
   token_types = [mlir.token_type() for _ in tokens]
   loop_carry_types = [*token_types, *loop_carry_types]
-  flat_loop_carry_types = util.flatten(loop_carry_types)
+  flat_loop_carry_types = mlir.flatten_ir_types(loop_carry_types)
   args = [*tokens, *args]
 
   flat_args = mlir.flatten_ir_values(args)
@@ -1732,7 +1732,7 @@ def _while_lowering(ctx, *args, cond_jaxpr, body_jaxpr, cond_nconsts,
     flat_cond_args = [
         cond_block.arguments[i] for i in range(len(flat_loop_carry_types))
     ]
-    cond_args = mlir.unflatten_ir_values(flat_cond_args, _map(len, loop_carry_types))
+    cond_args = mlir.unflatten_ir_values_like_types(flat_cond_args, loop_carry_types)
     # Remove tokens from cond args
     cond_args = cond_args[num_tokens:]
     x, _, z = util.split_list(cond_args, [cond_nconsts, body_nconsts])
@@ -1772,7 +1772,7 @@ def _while_lowering(ctx, *args, cond_jaxpr, body_jaxpr, cond_nconsts,
     flat_body_args = [
         body_block.arguments[i] for i in range(len(flat_loop_carry_types))
     ]
-    body_args = mlir.unflatten_ir_values(flat_body_args, _map(len, loop_carry_types))
+    body_args = mlir.unflatten_ir_values_like_types(flat_body_args, loop_carry_types)
     # Tokens are at the front of the args list to the while loop
     token_args, body_args = util.split_list(body_args, [num_tokens])
     tokens_in = mlir.TokenSet(zip(body_effects, token_args))
@@ -1799,7 +1799,7 @@ def _while_lowering(ctx, *args, cond_jaxpr, body_jaxpr, cond_nconsts,
     hlo.return_([*mlir.flatten_ir_values(out_tokens), *mlir.flatten_ir_values(x), *mlir.flatten_ir_values(y),
                   *mlir.flatten_ir_values(new_z)])
 
-  outputs = mlir.unflatten_ir_values(while_op.results, _map(len, loop_carry_types))
+  outputs = mlir.unflatten_ir_values_like_types(while_op.results, loop_carry_types)
   tokens, _, _, z = util.split_list(outputs, [num_tokens, cond_nconsts, body_nconsts])
   if tokens:
     ctx.set_tokens_out(mlir.TokenSet(zip(body_effects, tokens)))
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 62071350d..cdd8c11d1 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4543,8 +4543,8 @@ mlir.lowerable_effects.add_type(InOutFeedEffect)
 
 
 def _infeed_lowering(ctx, token, *, shapes, partitions):
-  output_types = safe_map(mlir.aval_to_ir_types, ctx.avals_out[:-1])
-  flat_output_types = util.flatten(output_types)
+  output_types = safe_map(mlir.aval_to_ir_type, ctx.avals_out[:-1])
+  flat_output_types = mlir.flatten_ir_types(output_types)
   # TODO(phawkins): verify `shapes` have a major-to-minor layout.
   layouts = ir.ArrayAttr.get([
       ir.ArrayAttr.get(
@@ -4561,7 +4561,7 @@ def _infeed_lowering(ctx, token, *, shapes, partitions):
     mlir.set_sharding(infeed, xla.sharding_to_proto(partitions))
   token = infeed.results[-1]
   outs = infeed.results[:-1]
-  return mlir.unflatten_ir_values(outs, safe_map(len, output_types)) + [
+  return mlir.unflatten_ir_values_like_types(outs, output_types) + [
       token,
   ]
 
diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
index f4914981c..206d52ba5 100644
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -2485,11 +2485,11 @@ def _scatter_lower(ctx, operand, indices, updates, *,
       scattered_dims_to_operand_dims=list(dnums.scatter_dims_to_operand_dims),
       index_vector_dim=len(ctx.avals_in[1].shape) - 1,
   )
-  result = mlir.aval_to_ir_types(aval_out)
+  result = mlir.aval_to_ir_type(aval_out)
   operand = [operand]
   updates = [updates]
   op = hlo.ScatterOp(
-      result,
+      (result,),
       operand,
       indices,
       updates,
@@ -2546,7 +2546,7 @@ def _scatter_add_lower_gpu(ctx, operand, indices, updates,
       index_vector_dim=len(ctx.avals_in[1].shape) - 1,
   )
   real_dtype = _real_dtype(aval_out.dtype)
-  operand_type_part = mlir.aval_to_ir_types(
+  operand_type_part = mlir.aval_to_ir_type(
       core.ShapedArray(aval_out.shape, real_dtype))
 
   def _scatter(operand_part, updates_part):
@@ -2554,7 +2554,7 @@ def _scatter_add_lower_gpu(ctx, operand, indices, updates,
     updates_part = [updates_part]
 
     scatter = hlo.ScatterOp(
-        operand_type_part,
+        (operand_type_part,),
         operand_part,
         indices,
         updates_part,
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 992e8d57e..3f27cda44 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -81,7 +81,7 @@ from jax._src.tree_util import (
 from jax._src.util import (
     HashableFunction, safe_map, safe_zip, wraps,
     distributed_debug_log, split_list, weakref_lru_cache,
-    merge_lists, flatten, subs_list, fun_name, fun_qual_name)
+    merge_lists, subs_list, fun_name, fun_qual_name)
 
 map, unsafe_map = safe_map, map
 zip, unsafe_zip = safe_zip, zip
@@ -1907,9 +1907,9 @@ def _pjit_lowering(ctx, *args, name, jaxpr, in_shardings,
                    out_shardings, in_layouts, out_layouts, resource_env,
                    donated_invars, keep_unused, inline):
   effects = list(ctx.tokens_in.effects())
-  output_types = map(mlir.aval_to_ir_types, ctx.avals_out)
+  output_types = map(mlir.aval_to_ir_type, ctx.avals_out)
   output_types = [mlir.token_type()] * len(effects) + output_types
-  flat_output_types = flatten(output_types)
+  flat_output_types = mlir.flatten_ir_types(output_types)
 
   func = _pjit_cached_lower_jaxpr_to_fun(
       ctx, name, jaxpr, tuple(effects), in_shardings,
@@ -1922,7 +1922,7 @@ def _pjit_lowering(ctx, *args, name, jaxpr, in_shardings,
                              ir.FlatSymbolRefAttr.get(func.name.value),
                              mlir.flatten_ir_values(args))
   mlir.wrap_compute_type_in_place(ctx, call)
-  out_nodes = mlir.unflatten_ir_values(call.results, map(len, output_types))
+  out_nodes = mlir.unflatten_ir_values_like_types(call.results, output_types)
   tokens, out_nodes = split_list(out_nodes, [len(effects)])
   tokens_out = ctx.tokens_in.update_tokens(mlir.TokenSet(zip(effects, tokens)))
   ctx.set_tokens_out(tokens_out)
diff --git a/jax/experimental/jax2tf/call_tf.py b/jax/experimental/jax2tf/call_tf.py
index 34599ec57..56cdd1781 100644
--- a/jax/experimental/jax2tf/call_tf.py
+++ b/jax/experimental/jax2tf/call_tf.py
@@ -645,11 +645,11 @@ def emit_tf_embedded_graph_custom_call(
 
   operands = list(operands)
   result_types = list(
-      util.flatten([mlir.aval_to_ir_types(aval) for aval in result_avals])
+      mlir.flatten_ir_types([mlir.aval_to_ir_type(aval) for aval in result_avals])
   )
   if ordered:
     operands.insert(0, ctx.tokens_in.get(call_tf_ordered_effect))
-    result_types.insert(0, mlir.token_type()[0])
+    result_types.insert(0, mlir.token_type())
 
   custom_call = hlo.CustomCallOp(
       result_types,
diff --git a/jax/interpreters/mlir.py b/jax/interpreters/mlir.py
index 15faf0212..90374706b 100644
--- a/jax/interpreters/mlir.py
+++ b/jax/interpreters/mlir.py
@@ -31,7 +31,6 @@ from jax._src.interpreters.mlir import (
   _lowerings as _lowerings,
   _platform_specific_lowerings as _platform_specific_lowerings,
   aval_to_ir_type as aval_to_ir_type,
-  aval_to_ir_types as aval_to_ir_types,
   core_call_lowering as core_call_lowering,
   custom_call as custom_call,
   dense_bool_elements as dense_bool_elements,
@@ -42,7 +41,7 @@ from jax._src.interpreters.mlir import (
   emit_python_callback as emit_python_callback,
   flatten_ir_values as flatten_lowering_ir_args,  # TODO(phawkins): remove me
   flatten_ir_values as flatten_ir_values,
-  unflatten_ir_values as unflatten_ir_values,
+  unflatten_ir_values_like_types as unflatten_ir_values_like_types,
   func_dialect as func_dialect,
   hlo as hlo,
   i32_attr as i32_attr,


Test patch:
diff --git a/tests/jaxpr_effects_test.py b/tests/jaxpr_effects_test.py
index f1a71cb42..931590081 100644
--- a/tests/jaxpr_effects_test.py
+++ b/tests/jaxpr_effects_test.py
@@ -89,12 +89,12 @@ def function_effect_lowering(ctx, *, effect):
     return []
   func = mlir._emit_lowering_rule_as_fun(_f, ctx)
 
-  output_types = map(mlir.aval_to_ir_types, ctx.avals_out)
+  output_types = map(mlir.aval_to_ir_type, ctx.avals_out)
   effs = list(ctx.tokens_in.effects())
   in_tokens = [ctx.tokens_in.get(eff) for eff in effs]
   token_types = [mlir.token_type() for _ in effs]
   output_types = [*token_types, *output_types]
-  flat_output_types = util.flatten(output_types)
+  flat_output_types = mlir.flatten_ir_types(output_types)
   call = mlir.func_dialect.CallOp(flat_output_types,
                                   mlir.ir.FlatSymbolRefAttr.get(func.name.value),
                                   mlir.flatten_ir_values(in_tokens))