jax-ml/jax · error · ValueError

Only 32-bit scalar types supported

Error message

Only 32-bit scalar types supported

What it means

The scalar branch of redux requires x to be exactly 32 bits because redux.sync only supports 32-bit operands. Scalars of f16, bf16, i8, i64, etc. fail this check before the integer/float dispatch.

Source

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

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] = {}
  if kind == ReductionKind.FMAX or kind == ReductionKind.FMIN:
    extra_kwargs = dict(nan=True)
  return nvvm.redux_sync(x, kind, mask, **extra_kwargs)


def prmt(high: ir.Value, low: ir.Value, permutation: ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if (result_type := high.type) != low.type:
    raise ValueError(f"Types must match, got {high.type} and {low.type}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the scalar to i32 or f32 before calling redux
2. Keep accumulators in 32-bit dtypes throughout the reduction, converting at load/store boundaries
3. For i64/f64 reductions, use a manual shuffle-based reduction instead of redux

Example fix

# before
out = redux(acc_i64, mask, kind)
# after
out = redux(arith.trunci(i32, acc_i64), mask, kind)

assert bitwidth(x.type) == 32, f'redux needs 32-bit scalar, got {x.type}'

def is_redux_scalar(v) -> bool:
    return bitwidth(v.type) == 32 and isinstance(v.type, (ir.IntegerType, ir.F32Type))

When it happens

Trigger: Calling redux with a scalar i8/i64/f16/bf16 value, e.g. after vector elements are extracted or when accumulating per-thread scalars in a non-32-bit dtype.

Common situations: Accumulating per-thread partial sums in bf16 for speed and then calling redux; using i64 loop counters fed into a redux-based reduction.

Related errors

• Only 32-bit 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/534de432fc6c334d. Report an issue: GitHub.