jax-ml/jax · error · NotImplementedError

Scalars are not supported in async_store_smem

Error message

Scalars are not supported in async_store_smem

What it means

_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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:649

      ctx,
      ref_aval,
      ref,
      ref_transform_avals,
      ref_transforms,
      handle_transposes=True,
  )

  base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
  if base_index is not None:
    barrier = barrier[base_index]

  cluster_idx_val = lowering._as_index(cluster_idx)
  gpu_cluster_dim = lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, cluster_dim)

  shape = ctx.avals_in[0].shape
  dtype = ctx.avals_in[0].dtype
  if not shape:
    raise NotImplementedError("Scalars are not supported in async_store_smem")

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    if remaining_ref_transforms:
      raise ValueError(f"Unexpected unhandled transforms: {remaining_ref_transforms}")
    assert isinstance(barrier, mgpu.DialectBarrierRef)
    cluster_idx_i32 = arith_dialect.index_cast(
        ir.IntegerType.get_signless(32), cluster_idx_val
    )
    atomic_type = None
    if atomic is not None:
      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))
    mgpu.dialect.async_store_smem(
        src,
        ref_smem,
        barrier.as_barrier_memref(),
        gpu_cluster_dim.value,
        cluster_idx_i32,
        atomic_type=atomic_type,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Store scalars with a regular store: smem[...] = value
2. Reshape the scalar to at least 1-D (value[None]) and size the SMEM ref accordingly
3. Use atomic or sync operations appropriate for scalars

Example fix

# before
async_store_smem(smem_scalar, scalar_val, barrier)
# after
smem_scalar[()] = scalar_val  # plain synchronous store

assert len(value.shape) > 0, 'async_store_smem does not support scalars'

When it happens

Trigger: Calling async_store_smem with a scalar value or a scalar (size-0 ndim) SMEM ref target.

Common situations: 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.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b81ea5e8cac6fb8f. Report an issue: GitHub.