jax-ml/jax · error · ValueError

Can't bitcast {x.type} to {new_type}

Error message

Can't bitcast {x.type} to {new_type}

What it means

In the vector-to-vector bitcast path, both vectors must have identical total bitwidth even though their shapes/element types may differ. A mismatch (e.g. vector<4xf32> to vector<2xf32>) raises this ValueError.

Source

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

    return vector.extract(
        vector.bitcast(ir.VectorType.get((1,), new_type), x),
        dynamic_position=[],
        static_position=ir.DenseI64ArrayAttr.get([0]),
    )
  if isinstance(x.type, ir.IntegerType) and isinstance(new_type, ir.VectorType):
    new_type = ir.VectorType(new_type)
    x_ty = ir.IntegerType(x.type)
    assert x_ty.width == bitwidth(new_type.element_type) * math.prod(
        new_type.shape
    )
    return vector.bitcast(
        new_type, vector.broadcast(ir.VectorType.get((1,), x_ty), x)
    )
  if isinstance(x.type, ir.VectorType) and isinstance(new_type, ir.VectorType):
    x_ty = ir.VectorType(x.type)
    new_ty = ir.VectorType(new_type)
    if bitwidth(x_ty) != bitwidth(new_ty):
      raise ValueError(f"Can't bitcast {x.type} to {new_type}")
    return vector.bitcast(new_type, x)
  if isinstance(x.type, ir.IntegerType) and isinstance(new_type, ir.FloatType):
    return arith.bitcast(new_type, x)
  if isinstance(x.type, ir.FloatType) and isinstance(new_type, ir.IntegerType):
    return arith.bitcast(new_type, x)
  if isinstance(x.type, ir.FloatType) and isinstance(new_type, ir.FloatType):
    return arith.bitcast(new_type, x)
  raise ValueError(f"Can't bitcast {x.type} to {new_type}")


def ceil_div(x: int, y: int):
  return (x + y - 1) // y


def vector_slice(v: ir.Value, s: slice):
  v_ty = ir.VectorType(v.type)
  if len(v_ty.shape) != 1:
    raise NotImplementedError(f"Only 1D vectors are supported {v_ty}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Recompute lengths: target element count = source element count * source bitwidth / target element bitwidth
  2. Use vector.shape_cast to change the number of elements of the same type first, then bitcast element-wise widths
  3. Add an assert comparing bitwidth(x_ty) == bitwidth(new_ty) in debug builds of your kernel

Example fix

# before
bitcast(v_f32_4, ir.VectorType.get((2,), f32))  # 128 vs 64 bits
# after
bitcast(v_f32_4, ir.VectorType.get((8,), f16))  # 128 == 128 bits
Defensive patterns

Strategy: validation

Validate before calling

xt, nt = ir.VectorType(x.type), ir.VectorType(new_type)
assert bitwidth(xt) == bitwidth(nt), f'vector bitcast width mismatch {xt} vs {nt}'

Prevention

When it happens

Trigger: Calling bitcast with two ir.VectorType operands whose element counts times element bitwidths differ, e.g. vector<8xi8> -> vector<4xi8>, or vector<4xf32> -> vector<16xi8> is fine but vector<4xf32> -> vector<8xi8> is not.

Common situations: Reshaping packed low-precision data (bf16/f16 vectors) where the element-count math is off by one; changing a vector width in one place but not its bitcast mirror.

Related errors


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