{"record":{"id":"abc92a98c62a899a","repo":"jax-ml/jax","slug":"unsupported-dtype-for-reduction-val-aval-dtype","errorCode":null,"errorMessage":"Unsupported dtype for reduction: {val_aval.dtype}. Only float32,  int32 and uint32 are supported.","messagePattern":"Unsupported dtype for reduction: (.+?)\\. Only float32,  int32 and uint32 are supported\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":4080,"sourceCode":"    result, reduced = results\n    if layout is not None:\n      result = gpu_core.layout_cast(result, layout)\n    return result, reduced\n\n\n@async_load_tmem_p.def_effectful_abstract_eval\ndef _async_load_tmem_abstract_eval(\n    src, *avals_flat, tree, reduce: tcgen05.LoadReduceOp | None = None\n):\n  if src.memory_space != gpu_core.MemorySpace.TMEM:\n    raise ValueError(\"Async load only supports TMEM refs\")\n  val_aval, effects = state_primitives._get_abstract_eval(\n      src, *avals_flat, tree=tree\n  )\n  if reduce is None:\n    return (val_aval,), effects\n  if val_aval.dtype not in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32)):\n    raise ValueError(\n        f\"Unsupported dtype for reduction: {val_aval.dtype}. Only float32, \"\n        \" int32 and uint32 are supported.\"\n    )\n  reduced_aval = jax_core.ShapedArray(val_aval.shape[:-1], val_aval.dtype)\n  return (val_aval, reduced_aval), effects\n\n\n@lowering.register_lowering_rule(async_load_tmem_p, mgpu.LoweringSemantics.Lane)\ndef _async_load_tmem_lowering_rule(\n    ctx: lowering.LoweringRuleContext,\n    x_ref,\n    *leaves,\n    tree,\n    reduce: Literal[\"max\", \"min\", \"absmax\", \"absmin\"] | None = None,\n):\n  assert isinstance(x_ref, tcgen05.TMEMRef)\n  x_aval = ctx.avals_in[0]\n  assert isinstance(x_aval, state_types.AbstractRef)","sourceCodeStart":4062,"sourceCodeEnd":4098,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L4062-L4098","documentation":"TMEM load reductions (the reduce= argument) are only hardware-supported for float32, int32, and uint32. Any other dtype on the loaded value raises ValueError listing the offending dtype.","triggerScenarios":"Calling async_load_tmem(..., reduce='min'/'max'/...) where the TMEM ref holds bfloat16, float16, float64, int8, etc.","commonSituations":"Applying f32-only reduction idioms to bf16 accumulators on Blackwell; assuming all dtypes support tcgen05 load reduction.","solutions":["Load without reduce and perform the reduction in registers afterwards at a supported dtype","Cast the TMEM contents to float32 (or int32/uint32) before the reduced load","Reconfigure the accumulator dtype of the preceding MMA to float32"],"exampleFix":"// before\nv, r = async_load_tmem(tmem_ref, reduce='max')  # tmem is bf16\n// after\nv = async_load_tmem(tmem_ref)\nr = jnp.max(v.astype(jnp.float32), axis=-1)","handlingStrategy":"validation","validationCode":"assert val_aval.dtype in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep TMEM reduced loads to f32/i32/u32; cast accumulators to float32"],"tags":["jax","pallas","tmem","dtype","reduction"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}