jax-ml/jax · error · ValueError

Unsupported dtype for reduction: {val_aval.dtype}. Only floa

Error message

Unsupported dtype for reduction: {val_aval.dtype}. Only float32,  int32 and uint32 are supported.

What it means

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.

Source

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

    result, reduced = results
    if layout is not None:
      result = gpu_core.layout_cast(result, layout)
    return result, reduced


@async_load_tmem_p.def_effectful_abstract_eval
def _async_load_tmem_abstract_eval(
    src, *avals_flat, tree, reduce: tcgen05.LoadReduceOp | None = None
):
  if src.memory_space != gpu_core.MemorySpace.TMEM:
    raise ValueError("Async load only supports TMEM refs")
  val_aval, effects = state_primitives._get_abstract_eval(
      src, *avals_flat, tree=tree
  )
  if reduce is None:
    return (val_aval,), effects
  if val_aval.dtype not in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32)):
    raise ValueError(
        f"Unsupported dtype for reduction: {val_aval.dtype}. Only float32, "
        " int32 and uint32 are supported."
    )
  reduced_aval = jax_core.ShapedArray(val_aval.shape[:-1], val_aval.dtype)
  return (val_aval, reduced_aval), effects


@lowering.register_lowering_rule(async_load_tmem_p, mgpu.LoweringSemantics.Lane)
def _async_load_tmem_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    x_ref,
    *leaves,
    tree,
    reduce: Literal["max", "min", "absmax", "absmin"] | None = None,
):
  assert isinstance(x_ref, tcgen05.TMEMRef)
  x_aval = ctx.avals_in[0]
  assert isinstance(x_aval, state_types.AbstractRef)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Load without reduce and perform the reduction in registers afterwards at a supported dtype
2. Cast the TMEM contents to float32 (or int32/uint32) before the reduced load
3. Reconfigure the accumulator dtype of the preceding MMA to float32

Example fix

// before
v, r = async_load_tmem(tmem_ref, reduce='max')  # tmem is bf16
// after
v = async_load_tmem(tmem_ref)
r = jnp.max(v.astype(jnp.float32), axis=-1)

assert val_aval.dtype in map(jnp.dtype, (jnp.float32, jnp.int32, jnp.uint32))

When it happens

Trigger: Calling async_load_tmem(..., reduce='min'/'max'/...) where the TMEM ref holds bfloat16, float16, float64, int8, etc.

Common situations: Applying f32-only reduction idioms to bf16 accumulators on Blackwell; assuming all dtypes support tcgen05 load reduction.

Related errors

• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsupported load reduce operation: {orig_reduce}
• Unsupported reduction for i32. Only min and max are supporte
• Index map function {debug_info.func_src_info} for {origin} m

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