jax-ml/jax · error · ValueError

Only 32-bit types supported

Error message

Only 32-bit types supported

What it means

redux lowers to PTX redux.sync, which on NVIDIA hardware only operates on 32-bit values. When x is a vector, redux recursively extracts elements and checks each is 32-bit; a vector of f16/bf16/i8/i64 elements fails this check.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2021

  if (x_bitwidth := bitwidth(result_type)) < 32:
    bits_ty = ir.IntegerType.get_signless(x_bitwidth)
    y_vec = bitcast(y, ir.VectorType.get((32 // x_bitwidth,), bits_ty))
    y = vector.extract(
        y_vec,
        dynamic_position=[],
        static_position=ir.DenseI64ArrayAttr.get([0]),
    )
  return bitcast(y, result_type)


ReductionKind = nvvm.ReductionKind


def redux(x: ir.Value, mask: ir.Value, kind: ReductionKind):
  i32 = ir.IntegerType.get_signless(32)
  if isinstance(vec_ty := x.type, ir.VectorType):
    if bitwidth(vec_ty.element_type) != 32:
      raise ValueError("Only 32-bit types supported")
    [vec_len] = vec_ty.shape
    result = llvm.mlir_undef(x.type)
    for i in range(vec_len):
      xi = llvm.extractelement(x, arith.constant(i32, i))
      yi = redux(xi, mask, kind)
      result = llvm.insertelement(result, yi, arith.constant(i32, i))
    return result
  if bitwidth(x.type) != 32:
    raise ValueError("Only 32-bit scalar types supported")
  if isinstance(x.type, ir.IntegerType):
    pass
  elif isinstance(x.type, ir.F32Type):
    if get_arch().major != 10:
      raise ValueError("F32 redux only supported on Blackwell GPUs")
  else:
    raise NotImplementedError(x.type)
  assert mask.type == i32
  extra_kwargs: dict[str, Any] = {}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upcast the vector to a 32-bit element type before redux (e.g. via arith.extf to f32 or fp_ext) and downcast after
2. Use a different reduction strategy for sub-32-bit types (shuffle-based warp_reduce or LLVM vector reduce ops)
3. Check dtype at the kernel boundary and insert explicit conversion nodes

Example fix

# before
result = redux(x_bf16, mask, ReductionKind.ADD)
# after
x_f32 = arith.extf(f32, x_bf16)
result = redux(x_f32, mask, ReductionKind.ADD)

def is_vec32(v):
    t = ir.VectorType(v.type)
    return bitwidth(t.element_type) == 32
assert is_vec32(x), 'upcast vector to 32-bit elements before redux'

def is_redux_capable_vector(v) -> bool:
    return isinstance(v.type, ir.VectorType) and bitwidth(ir.VectorType(v.type).element_type) == 32

When it happens

Trigger: Calling redux(x, mask, kind) where x is e.g. vector<4xbf16> or vector<8xi8> — any vector whose element bitwidth != 32.

Common situations: Feeding low-precision accumulators (bf16/f16) from attention or GEMM epilogues directly into redux without upcasting; assuming redux works like a generic reduce utility across dtypes.

Related errors

• Only 32-bit scalar types supported
• Unsupported dtype: {ref.dtype}
• Expected metadata dtype to be uint2, got: {meta.dtype}
• Sparse metadata format not implemented for {operand_dtype=}
• Data type bitcast is only supported from i8 to other types.

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