{"record":{"id":"572b5562a45d682b","repo":"jax-ml/jax","slug":"cannot-store-scalars-to-vmem","errorCode":null,"errorMessage":"Cannot store scalars to VMEM","messagePattern":"Cannot store scalars to VMEM","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2531,"sourceCode":"    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:\n    raise ValueError(\"Cannot store scalars to VMEM\")\n\n  mem_slice_shape = list(aval_out.shape)\n  for i, a in enumerate(idx.indices):\n    if not isinstance(a, primitives.Slice):\n      mem_slice_shape.insert(i, 1)\n  mem_slice_shape_iter = iter(mem_slice_shape)\n  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:","sourceCodeStart":2513,"sourceCodeEnd":2549,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2513-L2549","documentation":"Raised when storing a scalar (shape ()) value to VMEM. Vector memory ops on TPU are vector ops and cannot write a bare scalar; VMEM stores require at least 1-D values. Scalar stores belong in SMEM.","triggerScenarios":"pl.store to a VMEM ref with val_aval.shape == ().","commonSituations":"Storing a per-block scalar (loss, norm) to an output with scalar block shape in VMEM.","solutions":["Reshape the scalar to (1,) and give the output block shape (1,): store val[None]","Move that output to SMEM and store the scalar there"],"exampleFix":"# before\npl.store(out_ref, scalar_val)  # out block shape () in VMEM\n# after\npl.store(out_ref, scalar_val[None])  # out block shape (1,)","handlingStrategy":"validation","validationCode":"if val.shape == ():\n    val = val[None]  # make it (1,) for VMEM store\npl.store(vmem_ref, val)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never store bare scalars to VMEM; use shape (1,) blocks","Wrap scalar outputs as (1,) at the kernel signature level"],"tags":["jax","pallas","tpu","vmem","scalar-store"],"backgroundTag":"vmem-scalar-store-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}