jax-ml/jax · error · TypeError

output_is_signed must be non-None if and only if the MLIR ty

Error message

output_is_signed must be non-None if and only if the MLIR type is an integer type, got {output_is_signed=} for {elt}

What it means

FragmentedArray.bitcast validates that output_is_signed must be non-None exactly when the target MLIR type elt is an integer type. Passing output_is_signed for a float target, or omitting it for an integer target, raises this TypeError.

Source

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

        return utils.bitcast(result_int, original_arg_ty)
      else:
        assert vec_bitwidth > 32
        slice_len = 32 // utils.bitwidth(arg_ty.element_type)
        offset = 0
        slices = []
        while offset < vec_len:
          slice_end = min(offset + slice_len, vec_len)
          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
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

For integer targets always pass output_is_signed=True/False
For float/other targets omit output_is_signed entirely
Use ir.IntegerType.get_signless(...) only together with explicit output_is_signed

Example fix

# before
f = i32_frag.bitcast(ir.F32Type.get(), output_is_signed=False)
# after
f = i32_frag.bitcast(ir.F32Type.get())
# and for int targets:
i = f32_frag.bitcast(ir.IntegerType.get_signless(32), output_is_signed=True)

Defensive patterns

Strategy: validation

Validate before calling

if (output_is_signed is not None) != isinstance(elt, ir.IntegerType):
    raise TypeError('bitcast signedness/type mismatch')

Type guard

def bitcast_args_ok(elt, output_is_signed) -> bool:
    return (output_is_signed is not None) == isinstance(elt, ir.IntegerType)

Prevention

Always pair integer targets with explicit output_is_signed
Omit output_is_signed for float targets

When it happens

Trigger: fa.bitcast(ir.F32Type.get(), output_is_signed=False) (signedness given for non-int), or fa.bitcast(ir.IntegerType.get_signless(32)) with output_is_signed=None (int target without signedness).

Common situations: Reinterpreting register bits between int and float for fast exponent tricks or quantization, and forgetting that Mosaic tracks signedness for integer types explicitly.

Related errors

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