jax-ml/jax · error · TypeError

lax.bitcast_convert_type does not support bool or complex va

Error message

lax.bitcast_convert_type does not support bool or complex values unless the operand and destination types match. Got operand dtype={old_dtype}, {new_dtype=}. Consider using the arr.view() method instead.

What it means

bitcast_convert_type cannot reinterpret bool or complex values unless the operand and destination types are identical, because bool/complex bit layouts make cross-type reinterpretation ambiguous in XLA. Any bool/complex involved with a different target dtype raises TypeError.

Source

Thrown at jax/_src/lax/lax.py:5634

  old_nbits = dtypes.itemsize_bits(old_dtype)
  new_nbits = dtypes.itemsize_bits(new_dtype)

  if old_nbits == new_nbits:
    return operand.sharding
  elif old_nbits > new_nbits:
    return operand.sharding.update(spec=(*operand.sharding.spec, None))
  else:
    return operand.sharding.update(spec=operand.sharding.spec[:-1])

def _bitcast_convert_type_dtype_rule(operand, *, new_dtype):
  old_dtype = operand.dtype
  if (dtypes.issubdtype(old_dtype, np.bool_) or
      dtypes.issubdtype(old_dtype, np.complexfloating) or
      dtypes.issubdtype(new_dtype, np.bool_) or
      dtypes.issubdtype(new_dtype, np.complexfloating)):
    if old_dtype != new_dtype:
      raise TypeError("lax.bitcast_convert_type does not support bool or complex values "
                      "unless the operand and destination types match. "
                      f"Got operand dtype={old_dtype}, {new_dtype=}. "
                      "Consider using the arr.view() method instead.")
  return new_dtype

bitcast_convert_type_p = standard_primitive(
    _bitcast_convert_type_shape_rule, _bitcast_convert_type_dtype_rule,
    'bitcast_convert_type', weak_type_rule=_strip_weak_type,
    sharding_rule=_bitcast_convert_type_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'bitcast_convert_type'))
ad.defjvp_zero(bitcast_convert_type_p)
batching.defvectorized(bitcast_convert_type_p)

def _bitcast_convert_type_lower(ctx, operand, *, new_dtype):
  aval_out, = ctx.avals_out
  out_type = mlir.aval_to_ir_type(ctx.module_context, aval_out)
  out = hlo.bitcast_convert(out_type, operand)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use identical operand and destination dtypes where bool/complex are involved
2. Bitcast via an intermediate non-bool, non-complex dtype is not allowed — instead split complex into real/imag with jnp.real/jnp.imag then reinterpret each
3. Use arr.view() (jnp arrays) which supports the needed reinterpretation as the message suggests
4. Convert bool to uint8 with astype (value cast) if bit-exactness is not required

Example fix

// before
bits = lax.bitcast_convert_type(z_c64, jnp.float32)

// after
bits_re = lax.bitcast_convert_type(jnp.real(z_c64), jnp.float32)  # value-preserving cast not bitcast; better:
flat = z_c64.view(jnp.float32)  # arr.view method

def bitcast_ok(old, new):
    bad = lambda d: np.issubdtype(d, np.bool_) or np.issubdtype(d, np.complexfloating)
    return old == new or not (bad(old) or bad(new))

def is_bitcastable_pair(old_dt, new_dt) -> bool:
    bad = lambda d: np.issubdtype(d, np.bool_) or np.issubdtype(d, np.complexfloating)
    return old_dt == new_dt or not (bad(old_dt) or bad(new_dt))

try:
    bits = lax.bitcast_convert_type(x, dt)
except TypeError:
    bits = x.view(dt)  # arr.view supports the reinterpretation

When it happens

Trigger: lax.bitcast_convert_type(jnp.complex64_arr, jnp.float32); bitcasting bool to uint8; bitcasting float32 to complex64 — all with differing types.

Common situations: Trying to inspect raw bits of complex numbers; extracting bool bit patterns; porting numpy view()-based code to JAX.

Related errors

• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• bool() not supported for instances of type '{0}' (did you me
• Default value must be of type bool, got {default} of type {g

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