{"record":{"id":"0788608de93c7ca6","repo":"jax-ml/jax","slug":"unhandled-transforms-for-multimem-store-transfor","errorCode":null,"errorMessage":"Unhandled transforms for multimem_store: {transforms}","messagePattern":"Unhandled transforms for multimem_store: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":5392,"sourceCode":"    raise ValueError(\n        \"JAX device mesh is required by multimem_store, but not defined.\"\n    )\n  if set(collective_axes) != set(mesh_info.axis_names):\n    raise NotImplementedError(\n        \"Only collective_axes that include all JAX device mesh\"\n        f\" ({mesh_info.axis_names}) axes are supported, but got\"\n        f\" {collective_axes}\"\n    )\n  if transforms_tree is not None:\n    transforms = tree_util.tree_unflatten(transforms_tree, transforms_leaves)\n    local_ref_aval = ctx.avals_in[1]\n    assert isinstance(local_ref_aval, state_types.AbstractRef)\n    transform_avals = transforms_tree.unflatten(ctx.avals_in[2:])\n    local_ref, _, transforms = lowering._handle_transforms(\n        ctx, local_ref_aval, local_ref, transform_avals, transforms, allow_peer_refs=False\n    )\n    if transforms:\n      raise NotImplementedError(\n          f\"Unhandled transforms for multimem_store: {transforms}\"\n      )\n  multi_ref = ctx.launch_ctx.to_remote_multicast(local_ref)\n  scalar = not ctx.avals_in[0].shape\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n    val = lowering._ensure_ir_value(value, ctx.avals_in[0].dtype)\n    if scalar:\n      with lowering._wrap_in_custom_primitive_if_wg(ctx, [multi_ref.ref, val]) as [multi_ref, val]:\n        mgpu_utils.MultimemRef(multi_ref).store(val, indices=[])\n        if ctx.module_ctx.auto_barriers:\n          mgpu.warpgroup_barrier()\n    else:\n      mgpu.dialect.vector_store(val, multi_ref.ref, optimized=False, multimem=True)\n    return ()\n\n  if scalar:\n    multi_ref.store(lowering._ensure_ir_value(value, ctx.avals_in[0].dtype), [])\n  else:","sourceCodeStart":5374,"sourceCodeEnd":5410,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L5374-L5410","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use a locally-owned ref in this shard instead of a peer/remote ref","Update JAX to the latest version, since transform handling coverage changes between releases","Avoid composing extra transforms (e.g. slicing/cast helpers) on the ref before multimem_store; apply them to the value instead","If peer access is genuinely needed, use the regular store path plus explicit communication instead of multimem_store"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    kernel_jit(x)\nexcept NotImplementedError as e:\n    if 'multimem_store' in str(e):\n        # fall back to regular store + explicit collectives\n        run_fallback_kernel(x)","preventionTips":["Only pass locally allocated refs to multimem ops","Keep a non-multimem fallback kernel for CI on unsupported configurations"],"tags":["jax","pallas","transforms","multimem","lowering"],"backgroundTag":"unhandled-lowering-transforms","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}