jax-ml/jax · error · ValueError

Only bitcast between types of the same bitwidth supported

Error message

Only bitcast between types of the same bitwidth supported

What it means

FragmentedArray.bitcast only reinterprets bits between types of identical bitwidth (e.g. f32<->i32, f16<->i16); mismatched widths raise ValueError('Only bitcast between types of the same bitwidth supported'). Use to() for width-changing conversions.

Source

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

          args_slice = [utils.vector_slice(a, slice(offset, slice_end)) for a in args]
          slices.append(fast_instr(*args_slice))
          offset = slice_end
        return utils.vector_concat(slices)
    return fast_instr

  def bitcast(
      self, elt: ir.Type, *, output_is_signed: bool | None = None
  ) -> FragmentedArray:
    if (output_is_signed is not None) != isinstance(elt, ir.IntegerType):
      raise TypeError(
          "output_is_signed must be non-None if and only if the MLIR type is an"
          f" integer type, got {output_is_signed=} for {elt}"
      )

    if elt == self.mlir_dtype:
      return self
    if utils.bitwidth(elt) != utils.bitwidth(self.mlir_dtype):
      raise ValueError("Only bitcast between types of the same bitwidth supported")
    reg_type = self.registers.flat[0].type
    if isinstance(reg_type, ir.VectorType):
      reg_shape = ir.VectorType(reg_type).shape
      ty = ir.VectorType.get(reg_shape, elt)
    else:
      ty = elt

    return self._pointwise(
        lambda x: arith.bitcast(ty, x), output_is_signed=output_is_signed, restrict_bitwidth=False
    )

  def __getitem__(self, idx) -> FragmentedArray:
    base_idx, slice_shape, is_squeezed = utils.parse_indices(idx, self.shape)
    if isinstance(self.layout, WGSplatFragLayout):
      shape = tuple(d for d, s in zip(slice_shape, is_squeezed) if not s)
      return self.splat(self.registers.item(), shape, is_signed=self.is_signed)
    if not isinstance(self.layout, TiledLayout):
      raise NotImplementedError("Only arrays with tiled layouts can be sliced")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use fa.to(target_type) for width-changing, value-preserving casts
2. Pick a target type with matching bitwidth for bitcast (f32<->i32, f16<->i16, bf16<->i16)
3. Check utils.bitwidth(src)==utils.bitwidth(dst) before bitcasting

Example fix

# before
q = f32_frag.bitcast(ir.IntegerType.get_signless(8))
# after
q = f32_frag.to(ir.IntegerType.get_signless(8))  # numeric cast
# or same-width reinterpret:
r = f32_frag.bitcast(ir.IntegerType.get_signless(32), output_is_signed=True)

assert utils.bitwidth(src) == utils.bitwidth(dst), 'bitcast width mismatch'

def same_bitwidth(a: ir.Type, b: ir.Type) -> bool:
    import jax.experimental.mosaic.gpu.utils as u
    return u.bitwidth(a) == u.bitwidth(b)

When it happens

Trigger: fa.bitcast(ir.IntegerType.get_signless(32)) on an f16 fragment (16 vs 32 bits), or bitcasting f32 to i8.

Common situations: Quantization/dequantization code that tries to reinterpret a float as a narrower int, mistaking bitcast for a numeric cast.

Related errors

• output_is_signed must be non-None if and only if the MLIR ty
• Not implemented: bitcast 1D
• Not implemented: the 2nd minor dim can not be perfectly pack
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}

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