jax-ml/jax · error · NotImplementedError

Unhandled transforms for multimem_store: {transforms}

Error message

Unhandled transforms for multimem_store: {transforms}

What it means

After lowering resolves the transforms on the local ref for multimem_store, some transforms remain unhandled (e.g. disallowed or peer-related transforms, since allow_peer_refs=False). The Lane-semantics lowering rule only supports fully-consumed transforms; leftover ones abort with this NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:5392

    raise ValueError(
        "JAX device mesh is required by multimem_store, but not defined."
    )
  if set(collective_axes) != set(mesh_info.axis_names):
    raise NotImplementedError(
        "Only collective_axes that include all JAX device mesh"
        f" ({mesh_info.axis_names}) axes are supported, but got"
        f" {collective_axes}"
    )
  if transforms_tree is not None:
    transforms = tree_util.tree_unflatten(transforms_tree, transforms_leaves)
    local_ref_aval = ctx.avals_in[1]
    assert isinstance(local_ref_aval, state_types.AbstractRef)
    transform_avals = transforms_tree.unflatten(ctx.avals_in[2:])
    local_ref, _, transforms = lowering._handle_transforms(
        ctx, local_ref_aval, local_ref, transform_avals, transforms, allow_peer_refs=False
    )
    if transforms:
      raise NotImplementedError(
          f"Unhandled transforms for multimem_store: {transforms}"
      )
  multi_ref = ctx.launch_ctx.to_remote_multicast(local_ref)
  scalar = not ctx.avals_in[0].shape
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    val = lowering._ensure_ir_value(value, ctx.avals_in[0].dtype)
    if scalar:
      with lowering._wrap_in_custom_primitive_if_wg(ctx, [multi_ref.ref, val]) as [multi_ref, val]:
        mgpu_utils.MultimemRef(multi_ref).store(val, indices=[])
        if ctx.module_ctx.auto_barriers:
          mgpu.warpgroup_barrier()
    else:
      mgpu.dialect.vector_store(val, multi_ref.ref, optimized=False, multimem=True)
    return ()

  if scalar:
    multi_ref.store(lowering._ensure_ir_value(value, ctx.avals_in[0].dtype), [])
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a locally-owned ref in this shard instead of a peer/remote ref
2. Update JAX to the latest version, since transform handling coverage changes between releases
3. Avoid composing extra transforms (e.g. slicing/cast helpers) on the ref before multimem_store; apply them to the value instead
4. If peer access is genuinely needed, use the regular store path plus explicit communication instead of multimem_store

try:
    kernel_jit(x)
except NotImplementedError as e:
    if 'multimem_store' in str(e):
        # fall back to regular store + explicit collectives
        run_fallback_kernel(x)

When it happens

Trigger: Passing a ref to multimem_store that carries transforms the lowering cannot resolve with allow_peer_refs=False — e.g. a remote/peer ref produced by another shard, or composite transform combinations left over after _handle_transforms.

Common situations: Using refs obtained from other devices (peer refs) with multimem_store; composing multimem_store with experimental transform APIs; version changes in the transforms handling that leave additional residual transforms.

Related errors

• Unhandled transforms for multimem_load_reduce: {transforms}
• Unsupported core type: {core_type}
• Too many dynamic shapes in the input. Mosaic currently only
• program id was requested but no grid was provided.
• Invalid axis {axis} for num_programs

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