{"record":{"id":"bad5e659facdb1cb","repo":"jax-ml/jax","slug":"can-only-store-scalars-to-smem","errorCode":null,"errorMessage":"Can only store scalars to SMEM","messagePattern":"Can only store scalars to SMEM","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2514,"sourceCode":"  (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:\n    raise ValueError(\"Cannot store scalars to VMEM\")\n","sourceCodeStart":2496,"sourceCodeEnd":2532,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2496-L2532","documentation":"Raised when storing a non-scalar (array-shaped) value to SMEM. TPU scalar memory only holds scalars; vector data must live in VMEM. Storing an array to an SMEM ref is a memory-space/shape mismatch.","triggerScenarios":"pl.store where the ref memory space is SMEM but val_aval.shape is non-empty.","commonSituations":"Wrongly marking a tensor output as SMEM while copying an SMEM scalar-store pattern; only PRNG keys and scalars belong in SMEM.","solutions":["Use VMEM (default) for array-valued refs","Keep SMEM only for scalars and PRNG key material"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"if 'SMEM' in str(spec.memory_space or '') and any(d != 1 for d in (spec.block_shape or ())):\n    raise ValueError('SMEM stores must be scalar; use VMEM for arrays')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Arrays live in VMEM, scalars in SMEM — encode this in spec-builder helpers"],"tags":["jax","pallas","tpu","smem","shape-mismatch"],"backgroundTag":"pallas-smem-scalar-only","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}