{"record":{"id":"1611c38bc735f029","repo":"jax-ml/jax","slug":"unsupported-bitwidth-bitwidth","errorCode":null,"errorMessage":"Unsupported bitwidth: {bitwidth}","messagePattern":"Unsupported bitwidth: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":1251,"sourceCode":"          reg_shfl = utils.shfl_bfly(reg, 4)\n          new_reg = utils.prmt(reg, reg_shfl, perm)\n        elif bitwidth == 32:\n          i32_vec = ir.VectorType.get((1,), i32)\n          regs = [\n              utils.bitcast(utils.vector_slice(reg, slice(i, i + 1)), i32)\n              for i in range(2)\n          ]\n          reg_to_shfl = arith.select(is_even_row, regs[1], regs[0])\n          reg_shfl = utils.shfl_bfly(reg_to_shfl, 4)\n          new_reg_low = arith.select(is_even_row, regs[0], reg_shfl)\n          new_reg_high = arith.select(is_even_row, reg_shfl, regs[1])\n          new_reg_i32 = utils.vector_concat([\n              utils.bitcast(new_reg_low, i32_vec),\n              utils.bitcast(new_reg_high, i32_vec),\n          ])\n          new_reg = utils.bitcast(new_reg_i32, reg_ty)\n        else:\n          raise ValueError(f\"Unsupported bitwidth: {bitwidth}\")\n        tmp_new_regs.append(utils.bitcast(new_reg, reg_ty))\n      new_regs = np.asarray(\n          tmp_new_regs, dtype=object\n      ).reshape(new_layout.registers_shape(shape))\n      return FragmentedArray(\n          _registers=new_regs, _layout=new_layout, _is_signed=self.is_signed\n      )\n    if (\n        isinstance(self.layout, TiledLayout)\n        and isinstance(new_layout, TiledLayout)\n        and self.layout == tmem_native_layout(self.layout.vector_length)\n        and new_layout == tmem_native_layout(new_layout.vector_length)\n    ):\n      new_registers = np.empty(new_layout.registers_shape(shape), dtype=object)\n      if self.layout.vector_length > new_layout.vector_length:\n        ratio = self.layout.vector_length // new_layout.vector_length\n        new_length = new_layout.vector_length\n        for idx, reg in np.ndenumerate(self.registers):","sourceCodeStart":1233,"sourceCodeEnd":1269,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L1233-L1269","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Change the kernel dtype to a supported 16- or 32-bit type (f16/bf16/f32/i32) before the layout cast","Avoid the layout conversion by keeping the array in its original layout (skip to_layout / layout_cast for that value)","If you need sub-8-bit types, upcast first and downcast after the relayout","File a feature request / patch fragmented_array.py to handle your bitwidth"],"exampleFix":"# before\nacc = x.to_layout(WGMMA_LAYOUT)  # x is f64\n# after\nacc = x.cast_f32().to_layout(WGMMA_LAYOUT)","handlingStrategy":"type-guard","validationCode":"width = FragmentedArray/bitwidth of dtype  # e.g. via utils.bitwidth\nassert width in (1, 4, 8, 16, 32), f\"bitwidth {width} unsupported by to_layout\"","typeGuard":"def layout_castable(dtype) -> bool:\n    from jax._src.lib import _mlir_dialects as d\n    return isinstance(dtype, (d.ir.F16Type, d.ir.BF16Type, d.ir.F32Type, d.ir.IntegerType)) and utils_bitwidth(dtype) <= 32","tryCatchPattern":"try:\n    y = x.to_layout(L)\nexcept ValueError as e:\n    if 'Unsupported bitwidth' in str(e):\n        y = x.cast_f32().to_layout(L)\n    else:\n        raise","preventionTips":["Keep Mosaic GPU kernels on f16/bf16/f32/i32 dtypes","Grep release notes for mosaic layout-cast dtype support before adding new dtypes"],"tags":["jax","mosaic-gpu","layout-cast","bitwidth","wgmma","unsupported-dtype"],"backgroundTag":"unsupported-dtype-conversion","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}