{"record":{"id":"afa3cc1afee45de1","repo":"jax-ml/jax","slug":"cannot-pass-the-same-ref-into-a-mpmd-map-multiple","errorCode":null,"errorMessage":"Cannot pass the same ref into a mpmd map multiple times","messagePattern":"Cannot pass the same ref into a mpmd map multiple times","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mpmd.py","lineNumber":891,"sourceCode":"  flat_out_types_with_paths, out_tree = tree_util.tree_flatten_with_path(\n      out_types\n  )\n  out_paths, flat_out_types = util.unzip2(flat_out_types_with_paths)\n  # TODO(sharadmv): Use out_paths for debugging info.\n  del out_paths\n  flat_out_avals = tuple(\n      map(pallas_core._convert_out_shape_to_aval, flat_out_types)\n  )\n\n  def wrapper(*args):\n    flat_args_ft = ft.flatten(args)\n    flat_args, in_tree = flat_args_ft.vals, flat_args_ft.tree\n\n    seen_ref_ids = set()\n    for arg in flat_args:\n      if isinstance(arg, jax_core.Ref):\n        if id(arg) in seen_ref_ids:\n          raise NotImplementedError(\n              \"Cannot pass the same ref into a mpmd map multiple times\"\n          )\n        seen_ref_ids.add(id(arg))\n    # TODO(sharadmv): Use in_paths for debugging info.\n    flat_avals = tuple(map(jax_core.typeof, flat_args))\n\n    external_meshes = []\n    meshes = tuple(mesh for mesh, _ in meshes_and_fns)\n\n    flat_scratch_types, scratch_tree = tree_util.tree_flatten(scratch_types)\n    if len(meshes_and_fns) > 1:\n      # TODO(rdyro): For MPMD with more than one mesh, come up with a better\n      # solution for how to enforce core_type presence in scratch_shape.\n      # TODO(rdyro): Check if we need to have a similar check for in-kernel\n      # allocations (e.g., run_scoped, empty_ref) or can we assume the\n      # core_type is inherited from the caller (we then need the core_type in\n      # the caller context during tracing).\n      # TODO(rdyro): Also check inputs and outputs for core type.","sourceCodeStart":873,"sourceCodeEnd":909,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mpmd.py#L873-L909","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nmpmd_map(kernel, my_ref, my_ref, grid=grid)\n# after\nmpmd_map(kernel, my_ref, grid=grid)  # kernel accesses my_ref once; do both uses inside the kernel","handlingStrategy":"validation","validationCode":"import jax._src.core as jax_core\nfrom jax.tree_util import tree_flatten\n\nflat, _ = tree_flatten(args)\nrefs = [a for a in flat if isinstance(a, jax_core.Ref)]\nassert len({id(r) for r in refs}) == len(refs), 'duplicate Ref passed to mpmd_map'","typeGuard":"def has_unique_refs(args) -> bool:\n    flat, _ = tree_flatten(args)\n    refs = [a for a in flat if isinstance(a, jax_core.Ref)]\n    return len({id(r) for r in refs}) == len(refs)","tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","pallas","mpmd","ref-aliasing","notimplementederror"],"backgroundTag":"duplicate-ref-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}