{"record":{"id":"3779ebbc95a19300","repo":"jax-ml/jax","slug":"scatter-only-supports-storing-to-vmem-got-memory","errorCode":null,"errorMessage":"Scatter only supports storing to VMEM, got {memory_space}","messagePattern":"Scatter only supports storing to VMEM, got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":365,"sourceCode":"def _scatter_lowering_rule(\n    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree, add\n):\n  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)\n  ref_aval, *_ = tree.unflatten(ctx.avals_in)\n  if isinstance(ref_aval.memory_space, pallas_core.CoreMemorySpace):\n    if not isinstance(ref_aval.memory_space.mesh, sc_core.VectorSubcoreMesh):\n      raise ValueError(\n          \"Scatter only supports VectorSubcoreMesh, got\"\n          f\" {type(ref_aval.memory_space.mesh)}\"\n      )\n    memory_space = ref_aval.memory_space.memory_space\n  else:\n    memory_space = ref_aval.memory_space\n  if memory_space not in (\n      tpu_core.MemorySpace.VMEM,\n      pallas_core.MemorySpace.DEFAULT,\n  ):\n    raise ValueError(\n        f\"Scatter only supports storing to VMEM, got {memory_space}\"\n    )\n  if transforms:\n    ref_block_shape, *_ = ctx.block_shapes\n    ref, _ = tc_lowering._transform_ref(\n        ref, ref_aval, ref_block_shape, transforms\n    )\n  tpu.vector_store_idx(x, ref, indices, mask=mask, add=add)\n  return ()\n\n\ndef store_scatter(\n    ref: Ref,\n    indices: Sequence[jax.Array],\n    x: jax.Array,\n    *,\n    mask: jax.Array | None = None,\n) -> None:","sourceCodeStart":347,"sourceCodeEnd":383,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L347-L383","documentation":"The SparseCore scatter lowering only stores into VMEM (or DEFAULT). Storing directly to SMEM/CMEM/other memory spaces raises this ValueError during lowering.","triggerScenarios":"Calling store_scatter/addupdate_scatter on a ref allocated in SMEM or another non-VMEM CoreMemorySpace.memory_space.","commonSituations":"Optimizing scratch buffers into SMEM for TC kernels and reusing that code on SC; writing output directly from a non-VMEM intermediate.","solutions":["Perform the scatter into a VMEM buffer, then copy to the target space","Change the ref's memory space to VMEM/DEFAULT","Restructure so only supported spaces are scatter targets"],"exampleFix":"// before\nsc_primitives.store_scatter(smem_ref, idx, x)\n\n// after\nsc_primitives.store_scatter(vmem_ref, idx, x)\n# then copy vmem_ref -> smem_ref if needed","handlingStrategy":"validation","validationCode":"ms = ref_aval.memory_space\nif isinstance(ms, pallas_core.CoreMemorySpace):\n    ms = ms.memory_space\nassert ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT)","typeGuard":"def is_storable_space(ref_aval) -> bool:\n    ms = ref_aval.memory_space\n    if isinstance(ms, pallas_core.CoreMemorySpace):\n        ms = ms.memory_space\n    return ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT)","tryCatchPattern":null,"preventionTips":["Route all scatter writes through VMEM buffers","Keep a VMEM staging buffer per kernel for scatter outputs","Document supported memory spaces at kernel entry"],"tags":["jax","pallas","sparsecore","scatter","memory-space"],"backgroundTag":"jax-pallas-memory-space-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}