{"record":{"id":"ef0f116deb46fc35","repo":"jax-ml/jax","slug":"unsupported-element-type-for-atomic-stores-eleme","errorCode":null,"errorMessage":"Unsupported element type for atomic stores: {element_type}","messagePattern":"Unsupported element type for atomic stores: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3916,"sourceCode":"      if atomic in (\"and\", \"or\", \"xor\"):\n        ptx_type = \"b32\"\n      else:\n        ptx_type = \"s32\" if self.is_signed else \"u32\"\n    elif isinstance(element_type, (ir.F16Type, ir.BF16Type)):\n      if cluster_barrier_ptr is not None:\n        raise NotImplementedError(\"f16/bf16 not supported for async atomics\")\n      if atomic not in (\"add\", \"min\", \"max\"):\n        raise NotImplementedError(\n            f\"f16/bf16 only supports add, min, max atomics, got {atomic}\"\n        )\n      if (is_smem or multimem) and atomic != \"add\":\n        raise NotImplementedError(\n            f\"f16/bf16 SMEM/multimem atomics only support add, got {atomic}\"\n        )\n      ptx_type = f\"{element_type}x2\"\n      noftz = \"\" if multimem else \".noftz\"\n    else:\n      raise NotImplementedError(\n          f\"Unsupported element type for atomic stores: {element_type}\"\n      )\n    [vec_len] = vreg.type.shape\n    if element_bitwidth == 16:\n      if vec_len % 2 != 0:\n        raise NotImplementedError(\n            f\"f16/bf16 atomic stores require even vector length,\"\n            f\" got {vec_len}\"\n        )\n    i32_vec_len = vec_len * element_bitwidth // 32\n    vreg = utils.bitcast(vreg, ir.VectorType.get(\n        (i32_vec_len,), i32,\n    ))\n    regs = [\n        llvm.extractelement(vreg, arith.constant(i32, i))\n        for i in range(i32_vec_len)\n    ]\n    width = 1","sourceCodeStart":3898,"sourceCodeEnd":3934,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3898-L3934","documentation":"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.","triggerScenarios":"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.","commonSituations":"Double-precision accumulation buffers; 64-bit index/counter atomics; i8 quantized buffers with atomics.","solutions":["Use a supported dtype: cast the array to f32 or i32 before the atomic store","For 64-bit counters, split into two 32-bit halves or use a non-atomic store/reduction strategy","Drop atomic (plain store) if atomicity was not actually required"],"exampleFix":"# before\nfa64.store_tiled_async(ref, atomic='add')\n# after\nfa32 = fa64.convert(ir.F32Type.get())\nfa32.store_tiled_async(ref32, atomic='add')","handlingStrategy":"type-guard","validationCode":"from jax._src.lib import ir\nSUPPORTED = (ir.F32Type, ir.F16Type, ir.BF16Type)\nif not (isinstance(fa.mlir_dtype, SUPPORTED) or (isinstance(fa.mlir_dtype, ir.IntegerType) and fa.mlir_dtype.width == 32)):\n    fa = fa.convert(ir.F32Type.get())","typeGuard":"from jax._src.lib import ir\n\ndef atomic_dtype_supported(dtype) -> bool:\n    if isinstance(dtype, (ir.F32Type, ir.F16Type, ir.BF16Type)):\n        return True\n    return isinstance(dtype, ir.IntegerType) and dtype.width == 32","tryCatchPattern":"try:\n    fa.store_tiled_async(ref, atomic=atomic)\nexcept NotImplementedError:\n    fa.convert(ir.F32Type.get()).store_tiled_async(ref32, atomic='add')","preventionTips":["Standardize atomic accumulation buffers on f32 or i32","Guard kernel entry points with an atomic_dtype_supported check"],"tags":["jax","mosaic","gpu","atomics","dtype","not-implemented"],"backgroundTag":"unsupported-datatype-width","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}