{"record":{"id":"b81ea5e8cac6fb8f","repo":"jax-ml/jax","slug":"scalars-are-not-supported-in-async-store-smem","errorCode":null,"errorMessage":"Scalars are not supported in async_store_smem","messagePattern":"Scalars are not supported in async_store_smem","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":649,"sourceCode":"      ctx,\n      ref_aval,\n      ref,\n      ref_transform_avals,\n      ref_transforms,\n      handle_transposes=True,\n  )\n\n  base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)\n  if base_index is not None:\n    barrier = barrier[base_index]\n\n  cluster_idx_val = lowering._as_index(cluster_idx)\n  gpu_cluster_dim = lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, cluster_dim)\n\n  shape = ctx.avals_in[0].shape\n  dtype = ctx.avals_in[0].dtype\n  if not shape:\n    raise NotImplementedError(\"Scalars are not supported in async_store_smem\")\n\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n    if remaining_ref_transforms:\n      raise ValueError(f\"Unexpected unhandled transforms: {remaining_ref_transforms}\")\n    assert isinstance(barrier, mgpu.DialectBarrierRef)\n    cluster_idx_i32 = arith_dialect.index_cast(\n        ir.IntegerType.get_signless(32), cluster_idx_val\n    )\n    atomic_type = None\n    if atomic is not None:\n      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))\n    mgpu.dialect.async_store_smem(\n        src,\n        ref_smem,\n        barrier.as_barrier_memref(),\n        gpu_cluster_dim.value,\n        cluster_idx_i32,\n        atomic_type=atomic_type,","sourceCodeStart":631,"sourceCodeEnd":667,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L631-L667","documentation":"_async_store_smem_lowering inspects ctx.avals_in[0].shape; if the value being stored is a scalar (empty shape), it raises NotImplementedError because the TMA/async path operates on blocks of memory, not single elements.","triggerScenarios":"Calling async_store_smem with a scalar value or a scalar (size-0 ndim) SMEM ref target.","commonSituations":"Storing per-iteration scalars (e.g. loop counters, reduction accumulators) through the async path instead of plain stores; accidentally squeezing a (1,)-shaped buffer to a scalar.","solutions":["Store scalars with a regular store: smem[...] = value","Reshape the scalar to at least 1-D (value[None]) and size the SMEM ref accordingly","Use atomic or sync operations appropriate for scalars"],"exampleFix":"# before\nasync_store_smem(smem_scalar, scalar_val, barrier)\n# after\nsmem_scalar[()] = scalar_val  # plain synchronous store","handlingStrategy":"validation","validationCode":"assert len(value.shape) > 0, 'async_store_smem does not support scalars'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Reserve async_store_smem for block data; use plain stores for scalars"],"tags":["jax","pallas","mosaic-gpu","scalar","async-store","not-implemented"],"backgroundTag":"unsupported-operation-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}