jax-ml/jax · error · NotImplementedError

Cannot pass the same ref into a mpmd map multiple times

Error message

Cannot pass the same ref into a mpmd map multiple times

What it means

In JAX's Pallas MPMD (multi-program, multi-device) API, the same jax_core.Ref object cannot be passed as an argument to a single mpmd_map call more than once. The wrapper flattens the input PyTree and checks ref identity; a duplicate ref would create ambiguous aliasing between kernel inputs, so it is rejected with NotImplementedError.

Source

Thrown at jax/_src/pallas/mpmd.py:891

  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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove the duplicate: pass each Ref only once per mpmd_map call and adjust the kernel signature to take a single parameter for that buffer
If the kernel logically needs the buffer twice, read/write it through the single parameter inside the kernel body instead of aliasing at the call site
Create a separate, distinct buffer ref if two independent buffers were actually intended

Example fix

# before
mpmd_map(kernel, my_ref, my_ref, grid=grid)
# after
mpmd_map(kernel, my_ref, grid=grid)  # kernel accesses my_ref once; do both uses inside the kernel

Defensive patterns

Strategy: validation

Validate before calling

import jax._src.core as jax_core
from jax.tree_util import tree_flatten

flat, _ = tree_flatten(args)
refs = [a for a in flat if isinstance(a, jax_core.Ref)]
assert len({id(r) for r in refs}) == len(refs), 'duplicate Ref passed to mpmd_map'

Type guard

def has_unique_refs(args) -> bool:
    flat, _ = tree_flatten(args)
    refs = [a for a in flat if isinstance(a, jax_core.Ref)]
    return len({id(r) for r in refs}) == len(refs)

Prevention

Never reuse a Ref object in the mpmd_map argument list; create distinct buffers
Add an assertion on ref identity before calling mpmd_map in test suites

When it happens

Trigger: Calling mpmd.mpmd_map (or the public mpmd_map wrapper) with a kernel that receives the same Ref instance in two positions, e.g. mpmd_map(kernel, ref, ref) or a PyTree that contains the same Ref leaf twice.

Common situations: Reusing an output/input buffer ref for convenience (e.g. passing scratch or a buffer as both 'a' and 'b' arguments), or building an args tuple programmatically that accidentally duplicates a ref object.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/afa3cc1afee45de1. Report an issue: GitHub.