jax-ml/jax · error · NotImplementedError

Unsupported element type for atomic stores: {element_type}

Error message

Unsupported element type for atomic stores: {element_type}

What it means

store_tiled_async's atomic path only knows how to lower f32, 32-bit integers, f16 and bf16; any other mlir_dtype (f64, 64-bit ints, i8, etc.) reaches the final else and is rejected with the element type echoed in the message.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:3916

      if atomic in ("and", "or", "xor"):
        ptx_type = "b32"
      else:
        ptx_type = "s32" if self.is_signed else "u32"
    elif isinstance(element_type, (ir.F16Type, ir.BF16Type)):
      if cluster_barrier_ptr is not None:
        raise NotImplementedError("f16/bf16 not supported for async atomics")
      if atomic not in ("add", "min", "max"):
        raise NotImplementedError(
            f"f16/bf16 only supports add, min, max atomics, got {atomic}"
        )
      if (is_smem or multimem) and atomic != "add":
        raise NotImplementedError(
            f"f16/bf16 SMEM/multimem atomics only support add, got {atomic}"
        )
      ptx_type = f"{element_type}x2"
      noftz = "" if multimem else ".noftz"
    else:
      raise NotImplementedError(
          f"Unsupported element type for atomic stores: {element_type}"
      )
    [vec_len] = vreg.type.shape
    if element_bitwidth == 16:
      if vec_len % 2 != 0:
        raise NotImplementedError(
            f"f16/bf16 atomic stores require even vector length,"
            f" got {vec_len}"
        )
    i32_vec_len = vec_len * element_bitwidth // 32
    vreg = utils.bitcast(vreg, ir.VectorType.get(
        (i32_vec_len,), i32,
    ))
    regs = [
        llvm.extractelement(vreg, arith.constant(i32, i))
        for i in range(i32_vec_len)
    ]
    width = 1

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a supported dtype: cast the array to f32 or i32 before the atomic store
2. For 64-bit counters, split into two 32-bit halves or use a non-atomic store/reduction strategy
3. Drop atomic (plain store) if atomicity was not actually required

Example fix

# before
fa64.store_tiled_async(ref, atomic='add')
# after
fa32 = fa64.convert(ir.F32Type.get())
fa32.store_tiled_async(ref32, atomic='add')

from jax._src.lib import ir
SUPPORTED = (ir.F32Type, ir.F16Type, ir.BF16Type)
if not (isinstance(fa.mlir_dtype, SUPPORTED) or (isinstance(fa.mlir_dtype, ir.IntegerType) and fa.mlir_dtype.width == 32)):
    fa = fa.convert(ir.F32Type.get())

from jax._src.lib import ir

def atomic_dtype_supported(dtype) -> bool:
    if isinstance(dtype, (ir.F32Type, ir.F16Type, ir.BF16Type)):
        return True
    return isinstance(dtype, ir.IntegerType) and dtype.width == 32

try:
    fa.store_tiled_async(ref, atomic=atomic)
except NotImplementedError:
    fa.convert(ir.F32Type.get()).store_tiled_async(ref32, atomic='add')

When it happens

Trigger: Calling store_tiled_async with atomic set on an array whose mlir_dtype is not F32/F16/BF16/32-bit IntegerType — e.g. f64 accumulation or i64 counters.

Common situations: Double-precision accumulation buffers; 64-bit index/counter atomics; i8 quantized buffers with atomics.

Related errors

• f32 not supported for async atomics
• f32 only supports add atomics, got {atomic}
• f16/bf16 not supported for async atomics
• f16/bf16 only supports add, min, max atomics, got {atomic}
• f16/bf16 SMEM/multimem atomics only support add, got {atomic

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