{"record":{"id":"f983e06a6725b55a","repo":"jax-ml/jax","slug":"smem-store-does-not-support-masks","errorCode":null,"errorMessage":"SMEM store does not support masks","messagePattern":"SMEM store does not support masks","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2512,"sourceCode":"  is_smem_store = memory_space == \"#tpu.memory_space<smem>\"\n  is_vmem_store = memory_space == \"#tpu.memory_space<vmem>\"\n  (aval_out,) = ctx.avals_out\n  if not isinstance(val, ir.Value):\n    val = ir_constant(val, mlir_type=_dtype_to_ir_type(val_aval.dtype))\n  if not is_smem_store and not ref_block_shape:\n    raise NotImplementedError(\n        \"Indexing into a ()-shaped Ref not yet supported on TPU.\")\n\n  starts, _, strides, _, _ = _indexer_to_start_size_stride(\n      idx,\n      ref_block_shape,\n      cast_to_index=True,\n  )\n  need_stride = not all((s is None or s == 1) for s in strides)\n\n  if is_smem_store:\n    if mask is not None:\n      raise ValueError(\"SMEM store does not support masks\")\n    if val_aval.shape:\n      raise ValueError(\"Can only store scalars to SMEM\")\n    result = memref.load(ref, starts)\n    result = _maybe_cast_load_to_bool(ctx, val_aval, result)\n    val = _maybe_cast_store_to_memref_type(ctx, val_aval, val)\n    memref.store(val, ref, starts)\n    return result\n\n  if not is_vmem_store:\n    extra = \"\"\n    if memory_space == \"#tpu.memory_space<any>\":\n      extra = \" ANY memory space can only be accessed using async_copy.\"\n    raise ValueError(\n        \"Loads and stores are only allowed on VMEM and SMEM references.\" + extra\n    )\n\n  # handling VMEM store below\n  if not val_aval.shape:","sourceCodeStart":2494,"sourceCodeEnd":2530,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2494-L2530","documentation":"Raised when a store to SMEM (scalar memory) supplies a mask. SMEM stores compile to scalar memref.store ops which have no masking support on TPU.","triggerScenarios":"pl.store(ref, val, mask=m) where ref's memory space is SMEM.","commonSituations":"Reusing a generic masked-store helper for both VMEM tensors and SMEM scalars; storing a scalar with a scalar predicate mask.","solutions":["Remove the mask for SMEM stores; guard with control flow or clamp the index instead","If conditional scalar writes are needed, compute the value to write (e.g. select old vs new) and store unconditionally"],"exampleFix":"// before\npl.store(smem_ref, val, mask=cond)\n// after\nval = jnp.where(cond, val, pl.load(smem_ref))\npl.store(smem_ref, val)","handlingStrategy":"fallback","validationCode":"def smem_store(ref, val, cond):\n    if cond is None:\n        pl.store(ref, val)\n    else:\n        pl.store(ref, jnp.where(cond, val, pl.load(ref)))","typeGuard":"def is_smem_ref(ref) -> bool:\n    return str(getattr(ref.aval, 'memory_space', '')).endswith('smem')","tryCatchPattern":null,"preventionTips":["Never pass masks to SMEM stores","Implement conditional scalar updates via load + where + store"],"tags":["jax","pallas","tpu","smem","masked-store"],"backgroundTag":"smem-store-no-mask","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}