jax-ml/jax · error · NotImplementedError

bitcast_convert_type with different bitwidths not supported

Error message

bitcast_convert_type with different bitwidths not supported yet: {old_dtype=}, {new_dtype=}

What it means

The fuser only supports bitcast_convert_type when the source and destination dtypes have identical bit widths (itemsize). Bitcasting across widths (e.g. float32 -> float16, or int32 -> int8) changes the array extent per block, which block specs can't yet express, hence NotImplementedError.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:2205

  return [block_transform]


@register_eval_rule(lax.bitcast_convert_type_p)
def _bitcast_convert_type_eval_rule(eval_ctx: KernelEvalContext, x, new_dtype):
  del eval_ctx
  return jax.lax.bitcast_convert_type(x, new_dtype)


@register_pull_block_spec_rule(lax.bitcast_convert_type_p)
def _bitcast_convert_type_pull_rule(
    ctx: PullRuleContext,
    block_transform: BlockIndexTransform,
    *,
    new_dtype: jnp.dtype,
):
  old_dtype = ctx.avals_in[0].dtype
  if old_dtype.itemsize != new_dtype.itemsize:
    raise NotImplementedError(
        'bitcast_convert_type with different bitwidths not supported yet:'
        f' {old_dtype=}, {new_dtype=}'
    )
  return [block_transform]


@register_eval_rule(prng.random_bits_p)
def _random_bits_eval_rule(eval_ctx: KernelEvalContext, key, bit_width, shape):
  del shape
  block_spec = eval_ctx.out_block_specs[0]
  indices = eval_ctx.get_out_block_indices()[0]
  block_shape = block_spec.block_shape
  # This is the important part here: we fold in block indices into the key so
  # each block gets different random numbers.
  for idx in indices:
    key = jax.random.fold_in(key, idx)
  return prng.random_bits(key, bit_width=bit_width, shape=block_shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use equal-width bitcasts (f32<->u32, f16<->u16, etc.) and do any width change with a real cast (lax.convert_element_type) instead
Perform the width-changing bitcast outside the fused kernel and pass the result as input
If narrowing is required, reshape/atomically split via supported ops before bitcasting equal-width pieces

Example fix

// before
y = lax.bitcast_convert_type(x, jnp.dtype('float16'))  # x is float32: widths differ
// after
y = lax.bitcast_convert_type(x, jnp.dtype('uint32'))  # equal width, then handle narrowing separately

Defensive patterns

Strategy: validation

Validate before calling

assert x.dtype.itemsize == jnp.dtype(new_dtype).itemsize, 'bitcast requires equal itemsize'

Type guard

def bitcast_widths_match(old: jnp.dtype, new: jnp.dtype) -> bool:
    return old.itemsize == new.itemsize

Try / catch

try:
    y = fused_bitcast(x, new_dtype)
except NotImplementedError as e:
    if 'different bitwidths' in str(e):
        y = lax.bitcast_convert_type(x, new_dtype)  # outside fusion
    else:
        raise

Prevention

Bitcast only equal-width types (f32<->u32, f16<->u16)
Use convert_element_type for width changes
Centralize dtype-pair checks in a helper

When it happens

Trigger: lax.bitcast_convert_type inside a fused Pallas region where old_dtype.itemsize != new_dtype.itemsize, e.g. bitcasting a f32 array to uint16, or float16 to int8.

Common situations: Packing/unpacking sub-word types (fp8/bf16) for TPU/GPU kernels; converting a wide accumulator to a narrow storage dtype via bitcast instead of a value-preserving cast; assuming bitcast behaves like astype.

Related errors

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