{"record":{"id":"63219a426f6d232e","repo":"jax-ml/jax","slug":"cannot-swap-scalars-to-vmem","errorCode":null,"errorMessage":"Cannot swap scalars to VMEM.","messagePattern":"Cannot swap scalars to VMEM\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2556,"sourceCode":"  mem_slice_shape = [\n      1 if b is pallas_core.squeezed else next(mem_slice_shape_iter)\n      for b in ref_block_shape\n  ]\n  mem_aval = aval_out.update(\n      shape=tuple(mem_slice_shape), sharding=jax_core.get_cur_mesh_sharding()\n  )\n  mem_aval_vec_type = ir.VectorType.get(\n      ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),\n      _dtype_to_ir_type(mem_aval.dtype, is_kernel_boundary=True)\n  )\n  if need_stride:\n    result = tpu.strided_load(mem_aval_vec_type, ref, starts, strides)\n  else:\n    result = vector.load(mem_aval_vec_type, ref, starts)\n  val = _maybe_cast_store_to_memref_type(ctx, val_aval, val)\n  if mem_aval != aval_out:\n    if not aval_out.shape:\n      raise ValueError(\"Cannot swap scalars to VMEM.\")\n    # We are slicing a scalar so provided dummy 1 indices\n    result_vec_type = ir.VectorType.get(\n        ctx.lowering_context.dynamic_shape_replacement_fn(aval_out.shape),\n      _dtype_to_ir_type(aval_out.dtype, is_kernel_boundary=True))\n    result = vector.shape_cast(result_vec_type, result)\n    val_vec_type = ir.VectorType.get(\n        ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),\n      _dtype_to_ir_type(mem_aval.dtype, is_kernel_boundary=True))\n    val = vector.shape_cast(val_vec_type, val)\n    if mask is not None:\n      mask_vec_type = ir.VectorType.get(\n          ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),\n          _dtype_to_ir_type(mask_aval.dtype)\n      )\n      mask = vector.shape_cast(mask_vec_type, mask)\n  result = _maybe_cast_load_to_bool(ctx, val_aval, result)\n\n  if need_stride:","sourceCodeStart":2538,"sourceCodeEnd":2574,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2538-L2574","documentation":"Raised during a VMEM swap when the memory aval is vector-shaped but the output aval is scalar: swapping (read-modify-write returning old value) a scalar in VMEM is unsupported because the returned old value cannot be a bare vector-memory scalar.","triggerScenarios":"pl.swap on a VMEM ref where the stored value/dtype physicalization makes mem_aval != aval_out and aval_out.shape is ().","commonSituations":"Using swap to atomically read-and-write a scalar accumulator in VMEM; extended dtypes (like custom element types) that force physicalized (vector) memory layouts.","solutions":["Give the swapped ref shape (1,) so the old value is a 1-element vector","Avoid swap for scalars: use separate load then store"],"exampleFix":"# before\nold = pl.swap(ref, scalar)  # ref block shape ()\n# after\nold = pl.swap(ref, scalar[None])[0]  # ref block shape (1,)","handlingStrategy":"validation","validationCode":"if val.shape == ():\n    val, ref_block = val[None], (1,)  # swap with vector shape\nold = pl.swap(ref_block_ref, val)[0]","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid pl.swap on scalar blocks; use (1,) shaped blocks"],"tags":["jax","pallas","tpu","swap","vmem","scalar"],"backgroundTag":"vmem-scalar-store-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}