# Copyright 2026 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""APIs for defining MPMD kernels in Pallas."""

from __future__ import annotations

from collections.abc import Callable, Iterable, Mapping, Sequence
import contextlib
import functools
import itertools as it
from typing import Any, Hashable
from collections.abc import Generator

from jax._src import api
from jax._src import api_util
from jax._src import config
from jax._src import core as jax_core
from jax._src import effects
from jax._src import lax
from jax._src import numpy as jnp
from jax._src import state
from jax._src.state import primitives as state_primitives
from jax._src import flattree as ft
from jax._src import tree_util
from jax._src import util
from jax._src import xla_metadata
from jax._src.mesh import get_abstract_mesh
from jax._src.frozen_dict import FrozenDict
from jax._src.interpreters import batching
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.pallas import core as pallas_core
from jax._src.pallas import pallas_call
from jax._src.state import discharge as state_discharge
from jax._src.typing import Array


def get_super_mesh_shape(
    meshes: Iterable[pallas_core.Mesh],
) -> Mapping[Hashable, int]:
  super_mesh_shape: dict[Hashable, int] = {}
  for mesh in meshes:
    for k, v in mesh.shape.items():
      # An extra check since `check_is_compatible_with` should catch it.
      assert (
          k not in super_mesh_shape or super_mesh_shape[k] == v
      ), f"Conflicting size for axis {k}"
      super_mesh_shape[k] = v
  return super_mesh_shape


@contextlib.contextmanager
def mpmd_map_tracing_context(
  mesh: pallas_core.Mesh,
  other_meshes: tuple[pallas_core.Mesh, ...],
) -> Generator[None]:
  super_mesh_shape = get_super_mesh_shape(other_meshes)
  with (
      mesh.tracing_context(),
      jax_core.extend_axis_env_nd(super_mesh_shape.items()),
      config._check_vma(False),
      xla_metadata.clear_xla_metadata(),
  ):
    yield


mpmd_map_p = jax_core.Primitive("mpmd_map")
mpmd_map_p.multiple_results = True


@mpmd_map_p.def_impl
def _mpmd_map_impl(*args, **params):
  jit_impl = api.jit(functools.partial(mpmd_map_p.bind, **params))
  with config.disable_jit(False):
    return jit_impl(*args)


@mpmd_map_p.def_effectful_abstract_eval
def _mpmd_map_abstract_eval(
    *in_avals,
    jaxprs,
    out_avals,
    input_output_aliases,
    interpret,
    compiler_params,
    meshes,
    **params,
):
  del params, compiler_params  # Unused.

  effs = {*pallas_core.get_interpret_effects(interpret)}
  all_mesh_axis_names = {
      eff.name
      for jaxpr in jaxprs
      for eff in jaxpr.effects
      if isinstance(eff, jax_core.NamedAxisEffect)
  }
  for mesh, jaxpr in zip(meshes, jaxprs):
    for eff in jaxpr.effects:
      if mesh.discharges_effect(eff):
        continue
      if pallas_core.kernel_local_effects.contains(eff):
        continue
      if isinstance(eff, effects.JaxprInputEffect):
        # We emit an effect if we have a Ref input that has been written to in
        # the kernel.
        assert not jaxpr.constvars
        index = jaxpr.invars.index(eff.input)
        if index < len(in_avals) and isinstance(
            in_avals[index], state.AbstractRef
        ):
          effs.add(eff.replace(index))
        continue
      if not isinstance(eff, jax_core.NamedAxisEffect):
        effs.add(eff)
        continue
      if eff.name not in all_mesh_axis_names:
        effs.add(eff)

  # TODO(mattjj,yashkatariya): if we hide vmapped away mesh axes, use this:
  # if not (all(a.sharding.mesh.are_all_axes_manual for a in in_avals) and
  #         all(a.sharding.mesh.are_all_axes_manual for a in out_avals) and
  #         get_abstract_mesh().are_all_axes_manual):
  #   raise ValueError("mpmd_map requires all mesh axes to be Manual, "
  #                    f"got {get_abstract_mesh().axis_types}")

  # NOTE(mattjj,yashkatariya): this doesn't catch auto-mode non-manual axes
  if not (all(p is None for a in in_avals if isinstance(a, jax_core.ShapedArray)
              for p in a.sharding.spec) and
          all(p is None for a in out_avals if isinstance(a, jax_core.ShapedArray)
              for p in a.sharding.spec)):
    raise ValueError(
        "mpmd_map requires all mesh axes to be Manual, got"
        f" {get_abstract_mesh().axis_types}"
    )
  if any(
      isinstance(a, state.AbstractRef)
      and isinstance(a.inner_aval, jax_core.ShapedArray)
      and not a.manual_axis_type.empty
      for a in in_avals
  ):
    raise ValueError(
        "mpmd_map does not support Refs with varying manual axis types:"
        f" {in_avals=}:"
        f" {[None if not isinstance(a, state.AbstractRef) else a.manual_axis_type for a in in_avals]}"
    )

  # TODO(slebedev): Handle pinned buffers as in ``pallas_call``.
  outin_aliases = {
      out_idx: in_idx for in_idx, out_idx in input_output_aliases.items()
  }
  out_avals = [
      in_avals[outin_aliases[out_idx]] if out_idx in outin_aliases else a
      for out_idx, a in enumerate(out_avals)
  ]
  # Avoid returning avals with Pallas memory spaces to the outside world.
  out_avals = tuple(
      a.update(memory_space=jax_core.MemorySpace.Device)
      if isinstance(a, jax_core.ShapedArray)
      and a.memory_space is not None
      and not isinstance(a.memory_space, jax_core.MemorySpace)
      else a
      for a in out_avals
  )
  return out_avals, effs


def _mpmd_map_typecheck_rule(
    ctx_factory, *in_atoms, **params
):
  del ctx_factory  # Unused.
  return _mpmd_map_abstract_eval(
      *(x.aval for x in in_atoms), **params
  )


jax_core.custom_typechecks[mpmd_map_p] = _mpmd_map_typecheck_rule


def _default_memory_space(meshes: Sequence[pallas_core.Mesh]):
  defaults = {mesh.default_memory_space for mesh in meshes}
  if len(defaults) != 1:
    raise ValueError(
        "Multiple meshes with different default memory spaces are not"
        " supported."
    )
  return defaults.pop()


def _mpmd_map_discharge_rule(
    ctx: state_discharge.DischargeContext,
    *args: Any,
    jaxprs,
    meshes,
    input_output_aliases,
    out_avals,
    debug,
    interpret,
    compiler_params,
    cost_estimate,
    metadata,
    name,
    external_meshes,
    **_,
):
  write_indices = set()
  for jaxpr in jaxprs:
    assert not jaxpr.constvars
    invar_idx = {v: i for i, v in enumerate(jaxpr.invars)}
    for eff in jaxpr.effects:
      if isinstance(eff, (state.WriteEffect, state.AccumEffect)):
        write_index = invar_idx[eff.input]
        if (
            write_index < len(ctx.in_avals) and
            isinstance(ctx.in_avals[write_index], state.AbstractRef)
        ):
          write_indices.add(write_index)

  default_memory_space = _default_memory_space(meshes)
  in_memory_spaces = [
      pallas_core.get_memory_space_aval(aval) for aval in ctx.in_avals
  ]
  in_memory_spaces = [
      default_memory_space if m is None else m for m in in_memory_spaces
  ]
  args = tuple(
      state_discharge.constrain(
          arg, memory_space, ctx.strip_memory_space,
          neutral_ms=(default_memory_space,),
      )
      if isinstance(aval, state.AbstractRef) else arg
      for arg, aval, memory_space in zip(args, ctx.in_avals, in_memory_spaces)
  )

  write_indices = sorted(write_indices)
  num_in = len(ctx.in_avals)
  num_out_orig = len(out_avals)
  num_out_new = len(write_indices)

  new_jaxprs = []
  all_meshes = (*meshes, *external_meshes)

  def _rewrite_to_include_new_outputs(jaxpr: jax_core.Jaxpr) -> jax_core.Jaxpr:

    def new_body(*args):
      in_refs, orig_out_refs, new_out_refs, scratch_refs = util.split_list(
          args, [num_in, num_out_orig, num_out_new]
      )
      del new_out_refs
      jax_core.eval_jaxpr(
          jaxpr, (), *(in_refs + orig_out_refs + scratch_refs)
      )
      return ()

    all_in_avals = [v.aval for v in jaxpr.invars]
    in_avals_trace, orig_out_avals_trace, scratch_avals_trace = util.split_list(
        all_in_avals, [num_in, num_out_orig]
    )
    new_out_avals_trace = [ctx.in_avals[i] for i in write_indices]
    tracing_avals = (
        in_avals_trace
        + orig_out_avals_trace
        + new_out_avals_trace
        + scratch_avals_trace
    )

    debug_info = api_util.debug_info(
        "mpmd_map_discharge", new_body, tracing_avals, {},
        sourceinfo=jaxpr.debug_info.func_src_info,
    )
    closed_jaxpr, _ = pe.trace_to_jaxpr(
        new_body, ft.flatten_args(*tracing_avals), debug_info)
    return closed_jaxpr

  for mesh, jaxpr in zip(meshes, jaxprs):
    with mpmd_map_tracing_context(mesh, all_meshes):
      new_jaxprs.append(_rewrite_to_include_new_outputs(jaxpr))

  outin_aliases = {o: i for i, o in input_output_aliases.items()}
  orig_out_avals = [
      ctx.in_avals[outin_aliases[i]] if i in outin_aliases else a
      for i, a in enumerate(out_avals)
  ]
  new_out_avals = [ctx.in_avals[i] for i in write_indices]
  updated_out_avals = [
      state_discharge.discharged_aval(
          a, discharge=True, strip_memory_space=ctx.strip_memory_space
      )
      for a in orig_out_avals + new_out_avals
  ]

  new_aliases = dict(input_output_aliases)
  for out_idx, in_idx in enumerate(write_indices):
    new_aliases[in_idx] = num_out_orig + out_idx

  res = mpmd_map_p.bind(
      *args,
      jaxprs=tuple(new_jaxprs),
      meshes=meshes,
      input_output_aliases=FrozenDict(new_aliases),
      out_avals=tuple(updated_out_avals),
      debug=debug,
      interpret=interpret,
      compiler_params=compiler_params,
      cost_estimate=cost_estimate,
      metadata=metadata,
      name=name,
      external_meshes=external_meshes,
  )

  # Split the results into original outputs and updated refs.
  ans, updated_refs = util.split_list(res, [num_out_orig])
  new_invals = [None] * len(ctx.in_avals)
  for out_idx, in_idx in enumerate(write_indices):
    ms = in_memory_spaces[in_idx]
    new_invals[in_idx] = state_discharge.constrain(
        updated_refs[out_idx], ms, ctx.strip_memory_space,
        neutral_ms=(default_memory_space,),
    )
  # constrain aliased outputs
  for in_idx, out_idx in input_output_aliases.items():
    if isinstance(ctx.in_avals[in_idx], state.AbstractRef):
      ms = in_memory_spaces[in_idx]
      ans[out_idx] = state_discharge.constrain(
          ans[out_idx], ms, ctx.strip_memory_space,
          neutral_ms=(default_memory_space,),
      )

  return new_invals, ans


state_discharge.register_discharge_rule(mpmd_map_p)(_mpmd_map_discharge_rule)


def _mpmd_map_dce_rule(
    used_outs: list[bool], live_ins: list[bool], eqn: pe.JaxprEqn
) -> tuple[list[bool], pe.JaxprEqn | None]:
  return [True] * len(eqn.invars), eqn

pe.dce_rules[mpmd_map_p] = _mpmd_map_dce_rule


def _mpmd_map_partial_eval_custom(saveable, unks_in, inst_in, eqn):
  new_inst = [
      x
      for x, inst in zip(eqn.invars, inst_in)
      if type(x) is jax_core.Var and not inst
  ]
  num_outs = len(eqn.outvars)

  if not unks_in or any(unks_in):
    return None, eqn, [True] * num_outs, [True] * num_outs, new_inst

  if any(isinstance(v.aval, state.AbstractRef) for v in eqn.invars):
    # Force ``mpmd_map`` to run in both forward and backward passes.
    # Otherwise, partial eval only places the ``mpmd_map`` in the forward
    # pass, leaving the backward pass to read newly-allocated but *unchanged*
    # refs.
    return eqn, eqn, [False] * num_outs, [True] * num_outs, new_inst

  policy = pe.ensure_enum(
      saveable(eqn.primitive, *[x.aval for x in eqn.invars], **eqn.params)
  )
  if isinstance(policy, pe.RecomputeType):
    return eqn, eqn, [False] * num_outs, [True] * num_outs, new_inst
  else:
    return eqn, None, [False] * num_outs, [False] * num_outs, []


pe.partial_eval_jaxpr_custom_rules[mpmd_map_p] = _mpmd_map_partial_eval_custom


def _jaxpr_has_unbatchable_dataflow(
    jaxpr: jax_core.Jaxpr,
    batched_vars: set[jax_core.Var],
) -> bool:
  """Checks if jaxpr needs explicit loop fallback.

  Returns True if:
  1) emit_pipeline's index map uses any batched variable, or
  2) A primitive other than emit_pipeline uses any batched variable.
  """
  if not batched_vars:
    return False
  for eqn in jaxpr.eqns:
    if eqn.primitive.name == "emit_pipeline":
      all_args = eqn.params["args_tree"].unflatten(eqn.invars)
      if any(
          isinstance(v, jax_core.Var) and v in batched_vars
          for v in all_args.all_index_map_consts
      ):
        return True
    else:
      if any(
          isinstance(v, jax_core.Var) and v in batched_vars for v in eqn.invars
      ):
        return True

  return False


def _mpmd_map_batch_with_explicit_loop(
    args,
    dims,
    *,
    jaxprs,
    meshes,
    out_avals,
    input_output_aliases,
    axis_size,
    **params,
):
  """Batches ``mpmd_map`` by looping over the batch dimension.

  This is a fallback for cases in which adding a batch dimension to the kernel
  body is not supported -- currently when a scalar-prefetch (SMEM/index)
  argument is batched. It mirrors ``pallas_call._batch_with_explicit_loop``.
  """
  # Original memory-space constraints of the inputs. Slicing an array along the
  # batch dimension drops any memory-space annotation and needs to be reapplied.
  orig_memory_spaces = [
      getattr(jax_core.typeof(arg), "memory_space", None) for arg in args
  ]

  # Move the batch dimension to axis 0 for all batched args. For non-ref array
  # inputs we also strip any pallas memory-space annotation (e.g. SMEM), since
  # ops like ``dynamic_slice`` used to slice the batch dimension cannot operate
  # on pallas memory spaces outside a kernel. The original constraint is
  # re-applied to each slice in ``loop_body`` below.
  moved_args = []
  for arg, dim, memory_space in zip(args, dims, orig_memory_spaces):
    if dim is None:
      moved_args.append(arg)
      continue
    if memory_space is not None and not isinstance(
        jax_core.typeof(arg), state.AbstractRef
    ):
      arg = state_primitives.with_memory_space_constraint(
          arg, jax_core.MemorySpace.Device
      )
    moved_args.append(batching.moveaxis(arg, dim, 0))

  batched_out_refs = [
      jax_core.new_ref(
          lax.empty((axis_size, *out_aval.shape), out_aval.dtype),
      )
      for out_aval in out_avals
  ]

  def loop_body(i, _):
    sliced_args = []
    for arg, dim, memory_space in zip(moved_args, dims, orig_memory_spaces):
      if dim is None:
        sliced_args.append(arg)
      elif isinstance(arg_aval := jax_core.typeof(arg), state.AbstractRef):
        sliced_args.append(
            jax_core.new_ref(arg[i], memory_space=arg_aval.memory_space)
        )
      else:
        sliced = arg[i]
        # Re-apply the memory-space constraint (e.g. SMEM)
        if memory_space is not None:
          sliced = state_primitives.with_memory_space_constraint(
              sliced, memory_space
          )
        sliced_args.append(sliced)

    outs = mpmd_map_p.bind(
        *sliced_args,
        jaxprs=jaxprs,
        meshes=meshes,
        out_avals=out_avals,
        input_output_aliases=input_output_aliases,
        **params,
    )

    # Write outputs into the batched output refs at position i.
    for batched_ref, out in zip(batched_out_refs, outs):
      batched_ref[i] = out

    # Write back any modified input refs.
    for arg, sliced_arg, dim in zip(moved_args, sliced_args, dims):
      if dim is not None and isinstance(jax_core.typeof(arg), state.AbstractRef):
        arg[i] = jax_core.freeze(sliced_arg)

    return None

  lax.fori_loop(0, axis_size, loop_body, init_val=None)

  batched_outs = [jax_core.freeze(ref) for ref in batched_out_refs]
  return batched_outs, (0,) * len(batched_outs)


def _mpmd_map_batching_rule(
    axis_data,
    args,
    dims,
    *,
    jaxprs,
    meshes,
    out_avals,
    input_output_aliases,
    **params,
):
  if all(d is None for d in dims):
    out = mpmd_map_p.bind(
        *args,
        jaxprs=jaxprs,
        meshes=meshes,
        out_avals=out_avals,
        input_output_aliases=input_output_aliases,
        **params,
    )
    return out, (None,) * len(out)

  for jaxpr in jaxprs:
    for var, dim in zip(jaxpr.invars[: len(args)], dims):
      if (
          not isinstance(var.aval, state.AbstractRef)
          and dim is not None
      ):
        raise ValueError(
            "Closed-over scalar constants cannot be batched. Pass them as"
            " inputs instead."
        )

  # Preserve the old behavior and avoid batching the body if batching size is 1.
  # This supports existing pl.kernel cases that doesn't use emit_pipeline.
  if axis_data.size == 1:
    squeezed_args = []
    for arg, dim in zip(args, dims):
      if dim is None:
        squeezed_args.append(arg)
      elif isinstance(arg_aval := jax_core.typeof(arg), state.AbstractRef):
        # This is a bit of a hack. We rely on the fact that JAX does not have
        # true mutable refs, and thus it is effectively free to squeeze-copy
        # the underlying array like we do below.
        #
        # TODO(slebedev): Add first class support for ``TransformedRef``s to
        # ``mpmd_map`` and get rid of this.
        squeezed_args.append(
            jax_core.new_ref(
                jnp.squeeze(arg[...], dim),
                memory_space=arg_aval.memory_space,
            )
        )
      else:
        squeezed_args.append(jnp.squeeze(arg, dim))

    outs = mpmd_map_p.bind(
        *squeezed_args,
        jaxprs=jaxprs,
        meshes=meshes,
        out_avals=out_avals,
        input_output_aliases=input_output_aliases,
        **params,
    )

    for arg, squeezed_arg, dim in zip(args, squeezed_args, dims):
      if dim is None:
        continue
      if isinstance(jax_core.typeof(arg), state.AbstractRef):
        arg[...] = jnp.expand_dims(jax_core.freeze(squeezed_arg), dim)
    return [jnp.expand_dims(out, 0) for out in outs], (0,) * len(outs)

  # Check if any batched argument reaches an unbatchable op.
  # If so, fall back to an explicit loop over the batch dimension (mirroring
  # ``pallas_call._batch_with_explicit_loop``).
  needs_explicit_loop = any(
      _jaxpr_has_unbatchable_dataflow(
          jaxpr,
          {
              var
              for var, dim in zip(jaxpr.invars[: len(args)], dims)
              if dim is not None
          },
      )
      for jaxpr in jaxprs
  )
  if needs_explicit_loop:
    return _mpmd_map_batch_with_explicit_loop(
        args,
        dims,
        jaxprs=jaxprs,
        meshes=meshes,
        out_avals=out_avals,
        input_output_aliases=input_output_aliases,
        axis_size=axis_data.size,
        **params,
    )

  # Move the batch dimension to axis 0 for all batched args.
  moved_args = [
      batching.moveaxis(arg, dim, 0) if dim is not None else arg
      for arg, dim in zip(args, dims)
  ]

  # Batch each jaxpr
  batched_jaxprs = []
  num_in, num_out = len(args), len(out_avals)
  all_meshes = (*meshes, *params.get("external_meshes", ()))
  for mesh, jaxpr in zip(meshes, jaxprs):
    in_axes = (
        tuple(0 if d is not None else None for d in dims)
        + (0,) * num_out
        + (None,) * (len(jaxpr.invars) - num_in - num_out)
    )
    with mpmd_map_tracing_context(mesh, all_meshes):
      batched_jaxpr, _ = batching.batch_jaxpr2(jaxpr, axis_data, in_axes)
    batched_jaxprs.append(batched_jaxpr)

  # Update out_avals to include the batch dimension at axis 0.
  out_avals = tree_util.tree_map(
      lambda a: a.update(shape=(axis_data.size, *a.shape)), out_avals
  )

  outs = mpmd_map_p.bind(
      *moved_args,
      jaxprs=tuple(batched_jaxprs),
      meshes=meshes,
      out_avals=out_avals,
      input_output_aliases=input_output_aliases,
      **params,
  )

  return outs, (0,) * len(outs)


batching.fancy_primitive_batchers[mpmd_map_p] = _mpmd_map_batching_rule


def _mpmd_map_is_high(*args, jaxprs, **params):
  del args, params
  return any(jaxpr.is_high for jaxpr in jaxprs)


mpmd_map_p.is_high = _mpmd_map_is_high


def _mpmd_map_to_lojax(
    *hi_args,
    meshes,
    jaxprs,
    external_meshes,
    out_avals,
    input_output_aliases,
    compiler_params,
    interpret,
    debug,
    cost_estimate,
    metadata,
    name,
    **params,
):
  in_avals = [jax_core.typeof(a) for a in hi_args]
  lo_args = [
      lo_val
      for aval, x in zip(in_avals, hi_args)
      for lo_val in aval.lower_val(x)
  ]

  lo_out_avals = [lo_aval for aval in out_avals for lo_aval in aval.lo_ty()]

  all_meshes = (*meshes, *external_meshes)
  lo_jaxprs = []
  for mesh, jaxpr in zip(meshes, jaxprs):
    with mpmd_map_tracing_context(mesh, all_meshes):
      closed_jaxpr = jaxpr
      closed_lo_jaxpr = pe.lower_jaxpr2(closed_jaxpr)
      assert not closed_lo_jaxpr.consts
      lo_jaxprs.append(closed_lo_jaxpr)

  input_index_mapping = pallas_call._get_index_mapping(in_avals)
  output_index_mapping = pallas_call._get_index_mapping(out_avals)
  new_input_output_aliases = {}
  for i, o in input_output_aliases.items():
    assert i in input_index_mapping
    assert o in output_index_mapping
    for i_lo, o_lo in zip(input_index_mapping[i], output_index_mapping[o]):
      new_input_output_aliases[i_lo] = o_lo

  lo_outs = mpmd_map_p.bind(
      *lo_args,
      meshes=meshes,
      jaxprs=tuple(lo_jaxprs),
      external_meshes=external_meshes,
      out_avals=tuple(lo_out_avals),
      input_output_aliases=FrozenDict(new_input_output_aliases),
      compiler_params=compiler_params,
      interpret=interpret,
      debug=debug,
      cost_estimate=cost_estimate,
      metadata=metadata,
      name=name,
      **params,
  )
  return pe.raise_lo_outs(out_avals, lo_outs)


mpmd_map_p.to_lojax = _mpmd_map_to_lojax


def _mpmd_map_tpu_lowering(
    ctx: mlir.LoweringRuleContext,
    *in_nodes,
    jaxprs,
    meshes,
    input_output_aliases,
    debug,
    interpret,
    compiler_params,
    cost_estimate,
    out_avals,
    metadata,
    name,
    external_meshes,
):
  try:
    from jax._src.pallas.mosaic import pallas_call_registration  # pyrefly: ignore[missing-import]
  except ImportError:
    raise pallas_call._unsupported_lowering_error("tpu")
  num_scratch = len(jaxprs[0].invars) - len(in_nodes) - len(ctx.avals_out)
  return pallas_call_registration.mpmd_map_tpu_lowering_rule(
      ctx,
      *in_nodes,
      jaxprs=jaxprs,
      meshes=meshes,
      input_output_aliases=input_output_aliases,
      debug=debug,
      interpret=interpret,
      compiler_params=compiler_params,
      cost_estimate=cost_estimate,
      out_avals=out_avals,
      metadata=metadata,
      name=name,
      external_meshes=external_meshes,
      num_scratch=num_scratch,
  )


def _mpmd_map_fallback_lowering(
    ctx: mlir.LoweringRuleContext,
    *in_nodes,
    meshes,
    jaxprs,
    out_avals,
    input_output_aliases,
    compiler_params,
    interpret,
    debug,
    cost_estimate,
    metadata,
    name,
    external_meshes,
):
  if len(jaxprs) != 1:
    raise NotImplementedError(
        "Lowering multiple mesh/function pairs is not currently supported"
    )
  if external_meshes:
    raise NotImplementedError(
        "External meshes are not currently supported in fallback lowering"
    )
  [jaxpr] = jaxprs
  [mesh] = meshes

  if compiler_params is not None:
    if hasattr(mesh, "dimension_semantics"):
      compiler_params = compiler_params.replace(
          dimension_semantics=mesh.dimension_semantics
      )

  num_scratch = len(jaxpr.invars) - len(in_nodes) - len(out_avals)
  scratch_avals = (
      [v.aval for v in jaxpr.invars[-num_scratch:]] if num_scratch > 0 else []
  )
  scratch_types = tuple(
      pallas_core.MemoryRef(v.inner_aval, v.memory_space) for v in scratch_avals
  )
  grid_spec = pallas_core.GridSpec(
      grid=tuple(mesh.shape.items()),
      in_specs=tuple(
          pallas_core.BlockSpec(
              memory_space=aval.memory_space
              if isinstance(aval, jax_core.ShapedArray)
              and not isinstance(aval.memory_space, jax_core.MemorySpace)
              else mesh.default_memory_space,
          )
          for aval in ctx.avals_in
      ),
      out_specs=tuple(
          pallas_core.BlockSpec(
              memory_space=aval.memory_space
              if isinstance(aval, jax_core.ShapedArray)
              and not isinstance(aval.memory_space, jax_core.MemorySpace)
              else mesh.default_memory_space,
          )
          for aval in out_avals
      ),
      scratch_shapes=scratch_types,
  )

  in_tree = tree_util.tree_structure(in_nodes)
  out_tree = tree_util.tree_structure(out_avals)

  in_origins = [f"arg{i}" for i in range(len(in_nodes))]
  out_origins = [f"out{i}" for i in range(len(out_avals))]

  _, grid_mapping = pallas_core.get_grid_mapping(
      grid_spec,
      [
          v.aval.inner_aval if isinstance(v.aval, state.AbstractRef) else v.aval
          for v in jaxpr.invars[: len(in_nodes)]
      ],
      in_tree,
      in_origins,
      [
          v.aval.inner_aval if isinstance(v.aval, state.AbstractRef) else v.aval
          for v in jaxpr.invars[len(in_nodes) : len(in_nodes) + len(out_avals)]
      ],
      out_tree,
      out_origins,
      debug=debug,
  )

  return pallas_call._pallas_call_lowering(
      ctx,
      *in_nodes,
      jaxpr=jaxpr,
      grid_mapping=grid_mapping,
      mesh=mesh,
      input_output_aliases=tuple(input_output_aliases.items()),
      debug=debug,
      interpret=interpret,
      compiler_params=compiler_params,
      cost_estimate=cost_estimate,
      out_avals=out_avals,
      metadata=metadata,
      name=name,
  )


def _mpmd_map_mgpu_lowering(ctx: mlir.LoweringRuleContext, *in_nodes, **params):
  try:
    from jax._src.pallas.mosaic_gpu import pallas_call_registration  # pyrefly: ignore[missing-import]
  except ImportError:
    raise pallas_call._unsupported_lowering_error("cuda")
  return pallas_call_registration.mpmd_map_mgpu_lowering_rule(
      ctx, *in_nodes, **params
  )


@functools.partial(mlir.register_lowering, mpmd_map_p)
def _mpmd_map_lowering(ctx: mlir.LoweringRuleContext, *in_nodes, **params):
  platforms = ctx.module_context.platforms
  if len(platforms) != 1:
    raise NotImplementedError(
        "mpmd_map does not support multi-platform lowering"
    )
  [platform] = platforms
  match platform:
    case "cuda" if not params.get("interpret"):
      return _mpmd_map_mgpu_lowering(ctx, *in_nodes, **params)
    case "cpu" | "cuda" | "rocm":
      return _mpmd_map_fallback_lowering(ctx, *in_nodes, **params)
    case "tpu":
      return _mpmd_map_tpu_lowering(ctx, *in_nodes, **params)
    case _:
      raise ValueError(f"Unsupported platform: {platform}")


def mpmd_map(
    meshes_and_fns: Sequence[tuple[pallas_core.Mesh, Callable[..., None]]],
    /,
    out_types: tree_util.PyTree = (),
    *,
    scratch_types: pallas_core.ScratchShapeTree = (),
    compiler_params: Any | None = None,
    interpret: bool | Any = False,
    debug: bool = False,
    cost_estimate: pallas_core.CostEstimate | None = None,
    name: str | None = None,
    metadata: dict[str, str] | None = None,
) -> Callable[..., Any]:
  interpret = (
      config.pallas_tpu_interpret_mode_context_manager.value or interpret
  )
  return _mpmd_map(
      meshes_and_fns,
      out_types,
      input_output_aliases={},
      scratch_types=scratch_types,
      compiler_params=compiler_params,
      interpret=interpret,
      debug=debug,
      cost_estimate=cost_estimate,
      name=name,
      metadata=metadata,
  )


def _aval_to_ref_aval(
    aval: Any,
    meshes: Sequence[pallas_core.Mesh],
) -> state.AbstractRef | state.TransformedRef:
  match aval:
    case state.AbstractRef():
      return aval
    case state.TransformedRef():
      return aval
    case jax_core.ShapedArray(memory_space=memory_space):
      if memory_space == jax_core.MemorySpace.Device:
        memory_space = _default_memory_space(meshes)
      return state.AbstractRef(aval, memory_space=memory_space)
    case jax_core.AbstractValue():
      return state.AbstractRef(aval, memory_space=None)
    case _ if hasattr(aval, "get_ref_aval"):
      ref_aval = aval.get_ref_aval()
      assert isinstance(ref_aval, state.AbstractRef)
      return ref_aval
    case _:
      raise ValueError(f"Unsupported abstract value type: {type(aval), aval}")


def _error_if_non_ref_consts(consts, debug_info):
  consts_avals = [
      aval
      for c in consts
      if not isinstance(aval := jax_core.typeof(c), state.AbstractRef)
  ]
  non_scalar_consts_avals = [
      aval
      for aval in consts_avals
      if not (isinstance(aval, jax_core.ShapedArray) and not aval.shape)
  ]
  if non_scalar_consts_avals:
    ctx = jax_core.JaxprPpContext()
    pp_consts_avals = ", ".join(
        jax_core.pp_aval(aval, ctx) for aval in non_scalar_consts_avals
    )
    raise ValueError(
        "The kernel function in mpmd_map"
        f" {debug_info.func_src_info} captures non-scalar array constants"
        f" [{pp_consts_avals}]. You can only close over scalars and Refs;"
        " arrays must be passed as explicit inputs."
    )


def _get_unique_consts(
    consts: Sequence[Sequence[Any]],
) -> tuple[list[Array], set[int]]:
  unique_consts = []
  unique_const_ids = set()
  for cs in consts:
    for c in cs:
      if id(c) not in unique_const_ids:
        unique_consts.append(c)
        unique_const_ids.add(id(c))
  return unique_consts, unique_const_ids


def _dedup_consts_and_unify_jaxpr_signatures(
    jaxprs: Sequence[jax_core.Jaxpr],
    consts_per_fn: Sequence[Sequence[Any]],
    flat_args: Sequence[Any],
    unflat_in_avals: Sequence[jax_core.AbstractValue],
    unflat_out_avals: Sequence[jax_core.AbstractValue],
    flat_kernel_avals: Sequence[jax_core.AbstractValue],
    meshes: Sequence[pallas_core.Mesh],
    all_meshes: tuple[pallas_core.Mesh, ...],
) -> tuple[list[jax_core.Jaxpr], list[Array]]:
  # Example:
  #   c1, c2, c3 are closed-over refs.
  #   fn1 closes over [c1, c2] -> traced jaxpr1 has constvars for [c1, c2]
  #   fn2 closes over [c2, c3] -> traced jaxpr2 has constvars for [c2, c3]
  #
  #   `_dedup_consts_and_unify_jaxpr_signatures` will:
  #     1. Deduplicate constants to `unique_consts` = [c1, c2, c3].
  #     2. Rewrite jaxprs to take all `unique_consts` as explicit inputs instead
  #        of constvars:
  #        new_jaxpr1: (in_args, c1, c2, c3, out_args, scratch_args) -> ()
  #        new_jaxpr2: (in_args, c1, c2, c3, out_args, scratch_args) -> ()
  #     3. Return the new jaxprs (with empty constvars) and `unique_consts`.

  unique_consts, const_ids = _get_unique_consts(consts_per_fn)

  arg_ids = {id(arg) for arg in flat_args}
  if any(const_id in arg_ids for const_id in const_ids):
    raise NotImplementedError(
        "Closed-over ref aliases with a passed-in ref is not supported."
    )

  unique_const_avals = [jax_core.typeof(c) for c in unique_consts]

  num_inputs = len(tree_util.tree_leaves(unflat_in_avals))
  num_outputs = len(tree_util.tree_leaves(unflat_out_avals))

  in_avals_flat, out_avals_flat, scratch_avals_flat = util.split_list(
      flat_kernel_avals, [num_inputs, num_outputs]
  )

  def make_rewritten_body(original_jaxpr, original_consts):
    def _rewritten_body(*args):
      in_args, unique_const_args, out_args, scratch_args = util.split_list(
          args, [num_inputs, len(unique_consts), num_outputs]
      )

      # Extract only the consts used by this jaxpr.
      c_map = {id(uc): arg for uc, arg in zip(unique_consts, unique_const_args)}
      mapped_consts = [c_map[id(c)] for c in original_consts]

      eval_args = in_args + out_args + scratch_args
      jax_core.eval_jaxpr(original_jaxpr, mapped_consts, *eval_args)
      return []

    return _rewritten_body

  new_jaxprs = []
  tracing_avals = (
      in_avals_flat + unique_const_avals + out_avals_flat + scratch_avals_flat
  )
  for mesh, jaxpr, consts in zip(meshes, jaxprs, consts_per_fn):
    debug_info = api_util.debug_info(
        "mpmd_map_closed_over",
        make_rewritten_body(jaxpr, consts), tracing_avals, {},
        sourceinfo=jaxpr.debug_info.func_src_info,
    )
    fun_to_trace = make_rewritten_body(jaxpr, consts)
    with mpmd_map_tracing_context(mesh, all_meshes):
      jaxpr, _ = pe.trace_to_jaxpr(
          fun_to_trace, ft.flatten_args(*tracing_avals),
          debug_info)
    assert not jaxpr.consts, jaxpr.consts
    new_jaxprs.append(jaxpr)
  return new_jaxprs, unique_consts


def _mpmd_map(
    meshes_and_fns: Sequence[tuple[pallas_core.Mesh, Callable[..., None]]],
    /,
    out_types: tree_util.PyTree = (),
    *,
    input_output_aliases: Mapping[int, int] = {},
    scratch_types: pallas_core.ScratchShapeTree = (),
    compiler_params: Any | None = None,
    interpret: bool | Any = False,
    debug: bool = False,
    cost_estimate: pallas_core.CostEstimate | None = None,
    name: str | None = None,
    metadata: dict[str, str] | None = None,
) -> Callable[..., Any]:
  """Like ``pallas_call``, but MPMD and without pipelining."""
  if not meshes_and_fns:
    raise ValueError("At least one mesh/function pair is required")

  is_output_sequence = isinstance(out_types, Sequence)

  flat_out_types_with_paths, out_tree = tree_util.tree_flatten_with_path(
      out_types
  )
  out_paths, flat_out_types = util.unzip2(flat_out_types_with_paths)
  # TODO(sharadmv): Use out_paths for debugging info.
  del out_paths
  flat_out_avals = tuple(
      map(pallas_core._convert_out_shape_to_aval, flat_out_types)
  )

  def wrapper(*args):
    flat_args_ft = ft.flatten(args)
    flat_args, in_tree = flat_args_ft.vals, flat_args_ft.tree

    seen_ref_ids = set()
    for arg in flat_args:
      if isinstance(arg, jax_core.Ref):
        if id(arg) in seen_ref_ids:
          raise NotImplementedError(
              "Cannot pass the same ref into a mpmd map multiple times"
          )
        seen_ref_ids.add(id(arg))
    # TODO(sharadmv): Use in_paths for debugging info.
    flat_avals = tuple(map(jax_core.typeof, flat_args))

    external_meshes = []
    meshes = tuple(mesh for mesh, _ in meshes_and_fns)

    flat_scratch_types, scratch_tree = tree_util.tree_flatten(scratch_types)
    if len(meshes_and_fns) > 1:
      # TODO(rdyro): For MPMD with more than one mesh, come up with a better
      # solution for how to enforce core_type presence in scratch_shape.
      # TODO(rdyro): Check if we need to have a similar check for in-kernel
      # allocations (e.g., run_scoped, empty_ref) or can we assume the
      # core_type is inherited from the caller (we then need the core_type in
      # the caller context during tracing).
      # TODO(rdyro): Also check inputs and outputs for core type.
      for scratch_type in flat_scratch_types:
        from jax._src.pallas.mosaic import core as tpu_core  # pyrefly: ignore[missing-import]

        if not isinstance(
            scratch_type.memory_space, pallas_core.CoreMemorySpace
        ) and scratch_type.memory_space not in (
            tpu_core.MemorySpace.HBM,
            tpu_core.MemorySpace.VMEM_SHARED,
        ):
          raise NotImplementedError(
              "MPMD map with more than one mesh requires scratch_type to have"
              f" a `core_type` specified, but {scratch_type=} is missing it."
          )

    # Kernels may have Refs that belong to external meshes (usually for
    # async kernels). For example, the SC ScalarSubcore may have a Reference
    # to a TC semaphore that it is signaling. There is no explicit TC mesh as
    # part of the user-provided meshes, and are instead snuck in via the aval.
    for aval in [*flat_avals, *flat_out_avals, *flat_scratch_types]:
      if (
          isinstance(aval, jax_core.ShapedArray)
          and isinstance(aval.memory_space, pallas_core.CoreMemorySpace)
          and aval.memory_space.mesh not in it.chain(meshes, external_meshes)
      ):
        external_meshes.append(aval.memory_space.mesh)

    all_meshes = (*meshes, *external_meshes)
    # Check that meshes are compatible with each other (e.g, have a consistent
    # core axis name in the sparsecore).
    for i, mesh in enumerate(all_meshes):
      for other_mesh in list(all_meshes)[i + 1 :]:
        mesh.check_is_compatible_with(other_mesh)

    unflat_in_avals = in_tree.unflatten(flat_avals)
    unflat_out_avals = out_tree.unflatten(flat_out_avals)
    unflat_scratch_types = scratch_tree.unflatten(flat_scratch_types)
    kernel_arg_avals = list(unflat_in_avals)
    if is_output_sequence:
      kernel_arg_avals.extend(unflat_out_avals)
    else:
      kernel_arg_avals.append(unflat_out_avals)
    if isinstance(unflat_scratch_types, Mapping):
      kernel_kwarg_avals = unflat_scratch_types
    else:
      kernel_arg_avals.extend(unflat_scratch_types)
      kernel_kwarg_avals = {}

    # Wrap kernel_arg_avals and kernel_kwarg_avals into Refs except
    # TransformedRefs to allowed their transform data be passed as scalar args.
    unflat_kernel_avals = tree_util.tree_map(
        functools.partial(_aval_to_ref_aval, meshes=meshes),
        (kernel_arg_avals, kernel_kwarg_avals),
        is_leaf=lambda x: isinstance(x, state.TransformedRef),
    )
    in_avals_ft = ft.flatten(unflat_kernel_avals)
    flat_kernel_avals = list(in_avals_ft.vals)

    jaxprs: list[jax_core.Jaxpr] = []
    consts_per_fn = []
    debug_infos = [api_util.debug_info("mpmd_map", fn, kernel_arg_avals, kernel_kwarg_avals)
                   for _, fn in meshes_and_fns]
    if name is not None:
      debug_infos = [di.replace_func_name(name) for di in debug_infos]
    # If names are non-distinct (e.g. because user passed multiple functions
    # with the same name, or because of the name= arg handled just above),
    # uniquify them with the core type.
    if len({di.func_name for di in debug_infos}) != len(debug_infos):
      debug_infos = [di.replace_func_name(f"{di.func_name}__{mesh.core_type}")
                     for di, mesh in zip(debug_infos, meshes)]
    for (mesh, fn), debug_info in zip(meshes_and_fns, debug_infos):
      with mpmd_map_tracing_context(mesh, all_meshes):
        jaxpr, out_avals = pe.trace_to_jaxpr(
            fn, in_avals_ft, debug_info
        )
      fun_out_tree = out_avals.tree
      if fun_out_tree != tree_util.tree_structure(None):
        raise ValueError(
            f"The kernel function in mpmd_map {debug_info.func_src_info}"
            f" should return None. It returns a PyTree: {fun_out_tree}."
        )
      if jaxpr.consts:
        _error_if_non_ref_consts(jaxpr.consts, debug_info)
      jaxprs.append(jaxpr)
      consts_per_fn.append(jaxpr.consts)

    if any(consts_per_fn):
      # If we close over any constants in the kernel functions, we need to
      # deduplicate them and then unify the jaxpr signatures.
      jaxprs, consts = _dedup_consts_and_unify_jaxpr_signatures(
          jaxprs,
          consts_per_fn,
          flat_args,
          unflat_in_avals,
          unflat_out_avals,
          flat_kernel_avals,
          meshes,
          all_meshes,
      )
    else:
      consts: list[Array] = []

    if debug:
      for mesh, jaxpr in zip(meshes, jaxprs):
        print(f"jaxpr for {mesh.core_type}")
        print(jaxpr)

    # TODO(slebedev): The named scope should not be necessary here.
    ctx = (
        api.named_scope(name) if name is not None else contextlib.nullcontext()
    )
    with ctx:
      flat_outs = mpmd_map_p.bind(
          *flat_args,
          *consts,
          meshes=tuple(meshes),
          jaxprs=tuple(jaxprs),
          external_meshes=tuple(external_meshes),
          out_avals=flat_out_avals,
          input_output_aliases=FrozenDict(input_output_aliases),
          compiler_params=compiler_params,
          interpret=interpret,
          debug=debug,
          cost_estimate=cost_estimate,
          metadata=FrozenDict(metadata) if metadata is not None else None,
          name=name,
      )
    return out_tree.unflatten(flat_outs)

  return wrapper