{"record":{"id":"534de432fc6c334d","repo":"jax-ml/jax","slug":"only-32-bit-scalar-types-supported","errorCode":null,"errorMessage":"Only 32-bit scalar types supported","messagePattern":"Only 32-bit scalar types supported","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":2030,"sourceCode":"\n\nReductionKind = nvvm.ReductionKind\n\n\ndef redux(x: ir.Value, mask: ir.Value, kind: ReductionKind):\n  i32 = ir.IntegerType.get_signless(32)\n  if isinstance(vec_ty := x.type, ir.VectorType):\n    if bitwidth(vec_ty.element_type) != 32:\n      raise ValueError(\"Only 32-bit types supported\")\n    [vec_len] = vec_ty.shape\n    result = llvm.mlir_undef(x.type)\n    for i in range(vec_len):\n      xi = llvm.extractelement(x, arith.constant(i32, i))\n      yi = redux(xi, mask, kind)\n      result = llvm.insertelement(result, yi, arith.constant(i32, i))\n    return result\n  if bitwidth(x.type) != 32:\n    raise ValueError(\"Only 32-bit scalar types supported\")\n  if isinstance(x.type, ir.IntegerType):\n    pass\n  elif isinstance(x.type, ir.F32Type):\n    if get_arch().major != 10:\n      raise ValueError(\"F32 redux only supported on Blackwell GPUs\")\n  else:\n    raise NotImplementedError(x.type)\n  assert mask.type == i32\n  extra_kwargs: dict[str, Any] = {}\n  if kind == ReductionKind.FMAX or kind == ReductionKind.FMIN:\n    extra_kwargs = dict(nan=True)\n  return nvvm.redux_sync(x, kind, mask, **extra_kwargs)\n\n\ndef prmt(high: ir.Value, low: ir.Value, permutation: ir.Value):\n  i32 = ir.IntegerType.get_signless(32)\n  if (result_type := high.type) != low.type:\n    raise ValueError(f\"Types must match, got {high.type} and {low.type}\")","sourceCodeStart":2012,"sourceCodeEnd":2048,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L2012-L2048","documentation":"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.","triggerScenarios":"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.","commonSituations":"Accumulating per-thread partial sums in bf16 for speed and then calling redux; using i64 loop counters fed into a redux-based reduction.","solutions":["Convert the scalar to i32 or f32 before calling redux","Keep accumulators in 32-bit dtypes throughout the reduction, converting at load/store boundaries","For i64/f64 reductions, use a manual shuffle-based reduction instead of redux"],"exampleFix":"# before\nout = redux(acc_i64, mask, kind)\n# after\nout = redux(arith.trunci(i32, acc_i64), mask, kind)","handlingStrategy":"type-guard","validationCode":"assert bitwidth(x.type) == 32, f'redux needs 32-bit scalar, got {x.type}'","typeGuard":"def is_redux_scalar(v) -> bool:\n    return bitwidth(v.type) == 32 and isinstance(v.type, (ir.IntegerType, ir.F32Type))","tryCatchPattern":null,"preventionTips":["Convert scalars to i32/f32 before redux","Use warp_reduce for wider dtypes"],"tags":["mosaic-gpu","redux","dtype","hardware-limit"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}