jax-ml/jax · error · ValueError

Unsupported bitwidth: {bitwidth}

Error message

Unsupported bitwidth: {bitwidth}

What it means

FragmentedArray.to_layout converts registers between fragment layouts (e.g. for WGMMA). When reassembling sub-32-bit values it bitcasts halves through i32 vectors; only specific bitwidths have a code path. Hitting the else branch means the element bitwidth is not one of the handled widths (e.g. 4/8/16/32) in this conversion.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:1251

          reg_shfl = utils.shfl_bfly(reg, 4)
          new_reg = utils.prmt(reg, reg_shfl, perm)
        elif bitwidth == 32:
          i32_vec = ir.VectorType.get((1,), i32)
          regs = [
              utils.bitcast(utils.vector_slice(reg, slice(i, i + 1)), i32)
              for i in range(2)
          ]
          reg_to_shfl = arith.select(is_even_row, regs[1], regs[0])
          reg_shfl = utils.shfl_bfly(reg_to_shfl, 4)
          new_reg_low = arith.select(is_even_row, regs[0], reg_shfl)
          new_reg_high = arith.select(is_even_row, reg_shfl, regs[1])
          new_reg_i32 = utils.vector_concat([
              utils.bitcast(new_reg_low, i32_vec),
              utils.bitcast(new_reg_high, i32_vec),
          ])
          new_reg = utils.bitcast(new_reg_i32, reg_ty)
        else:
          raise ValueError(f"Unsupported bitwidth: {bitwidth}")
        tmp_new_regs.append(utils.bitcast(new_reg, reg_ty))
      new_regs = np.asarray(
          tmp_new_regs, dtype=object
      ).reshape(new_layout.registers_shape(shape))
      return FragmentedArray(
          _registers=new_regs, _layout=new_layout, _is_signed=self.is_signed
      )
    if (
        isinstance(self.layout, TiledLayout)
        and isinstance(new_layout, TiledLayout)
        and self.layout == tmem_native_layout(self.layout.vector_length)
        and new_layout == tmem_native_layout(new_layout.vector_length)
    ):
      new_registers = np.empty(new_layout.registers_shape(shape), dtype=object)
      if self.layout.vector_length > new_layout.vector_length:
        ratio = self.layout.vector_length // new_layout.vector_length
        new_length = new_layout.vector_length
        for idx, reg in np.ndenumerate(self.registers):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the kernel dtype to a supported 16- or 32-bit type (f16/bf16/f32/i32) before the layout cast
2. Avoid the layout conversion by keeping the array in its original layout (skip to_layout / layout_cast for that value)
3. If you need sub-8-bit types, upcast first and downcast after the relayout
4. File a feature request / patch fragmented_array.py to handle your bitwidth

Example fix

# before
acc = x.to_layout(WGMMA_LAYOUT)  # x is f64
# after
acc = x.cast_f32().to_layout(WGMMA_LAYOUT)

width = FragmentedArray/bitwidth of dtype  # e.g. via utils.bitwidth
assert width in (1, 4, 8, 16, 32), f"bitwidth {width} unsupported by to_layout"

def layout_castable(dtype) -> bool:
    from jax._src.lib import _mlir_dialects as d
    return isinstance(dtype, (d.ir.F16Type, d.ir.BF16Type, d.ir.F32Type, d.ir.IntegerType)) and utils_bitwidth(dtype) <= 32

try:
    y = x.to_layout(L)
except ValueError as e:
    if 'Unsupported bitwidth' in str(e):
        y = x.cast_f32().to_layout(L)
    else:
        raise

When it happens

Trigger: Calling to_layout(new_layout) on a FragmentedArray whose mlir_dtype has an unusual bitwidth (e.g. f64, 64-bit integers, or exotic widths like 2-bit) during layout_cast, wgmma lowering, or building arrays from IR values.

Common situations: Running WGMMA/relayout paths with fp64 or 64-bit integer dtypes that the Mosaic GPU layout-cast code never implemented; version upgrades that route a previously-supported dtype into this relayout path.

Related errors

• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Accumulator shape {inner.shape} does not match value shape {
• Accumulator dtype {inner.dtype} does not match value dtype {
• Cannot convert from {self.layout} to {new_layout}
• Concatenating arrays with strided layout is only supported a

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