jax-ml/jax · error · ValueError

Cannot bitcast from {x.dtype} ({old_bitwidth} bits) to {dtyp

Error message

Cannot bitcast from {x.dtype} ({old_bitwidth} bits) to {dtype} ({new_bitwidth} bits), because {x.shape[-1]=} * {old_bitwidth} is not divisible by {new_bitwidth}

What it means

For non-scalar arrays, bitcasting to a different bitwidth requires the last dimension's total bits (shape[-1] * old_bitwidth) to be divisible by the new bitwidth. If not, the bits cannot be repacked into whole elements of the new dtype.

Source

Thrown at jax/_src/pallas/mosaic/sc_primitives.py:448


bitcast_p = jax_core.Primitive("bitcast")


@bitcast_p.def_abstract_eval
def _bitcast_abstract_eval(x, dtype):
  old_bitwidth = dtypes.itemsize_bits(x.dtype)
  new_bitwidth = dtypes.itemsize_bits(dtype)
  if old_bitwidth == new_bitwidth:
    return jax_core.ShapedArray(x.shape, dtype)
  if x.ndim == 0:
    raise ValueError(
        "Cannot bitcast a ()-shaped array to a dtype with a different bitwidth:"
        f" {old_bitwidth=} vs {new_bitwidth=}"
    )
  new_last_dim, rem = divmod(x.shape[-1] * old_bitwidth, new_bitwidth)
  if rem:
    raise ValueError(
        f"Cannot bitcast from {x.dtype} ({old_bitwidth} bits) to"
        f" {dtype} ({new_bitwidth} bits), because {x.shape[-1]=} *"
        f" {old_bitwidth} is not divisible by {new_bitwidth}"
    )
  return jax_core.ShapedArray((*x.shape[:-1], new_last_dim), dtype)


@sc_lowering.register_lowering_rule(bitcast_p)
def _bitcast_lowering_rule(ctx: sc_lowering.LoweringRuleContext, x, *, dtype):
  del dtype  # Unused.
  [out_aval] = ctx.avals_out
  return vector.bitcast(ctx.aval_to_ir_type(out_aval), x)


def bitcast(x: jax.Array, dtype: jax.typing.DTypeLike) -> jax.Array:
  """Bitcasts an array to a different dtype.

  Unlike ``lax.bitcast_convert_type``, this function returns an array of the

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pad or slice the last dimension so shape[-1]*old_bits is divisible by new_bits (e.g. make last dim even for 8->16 bit casts)
  2. Choose a target dtype whose bitwidth divides the row's total bits
  3. Verify the last-dim size matches the ratio new_bitwidth//gcd(old,new)

Example fix

// before
y = bitcast(x_u8, jnp.uint16)  # x.shape[-1]=3 -> 24 bits not divisible by 16

// after
x = jnp.pad(x_u8, [(0,0),(0,1)])  # last dim 4
y = bitcast(x, jnp.uint16)
Defensive patterns

Strategy: validation

Validate before calling

ob, nb = dtypes.itemsize_bits(x.dtype), dtypes.itemsize_bits(dtype)
if ob != nb and (x.shape[-1] * ob) % nb:
    pad = (-x.shape[-1]) % (nb // math.gcd(ob, nb))
    x = jnp.pad(x, [(0, 0)] * (x.ndim - 1) + [(0, pad)])

Type guard

def bitcast_ok(x, dtype) -> bool:
    ob, nb = dtypes.itemsize_bits(x.dtype), dtypes.itemsize_bits(dtype)
    return ob == nb or (x.ndim > 0 and (x.shape[-1] * ob) % nb == 0)

Prevention

When it happens

Trigger: e.g. bitcast of a float32 array with last dim 3 to int32 (3*32/32 ok) vs last dim 3 to bfloat16 from int8 where 3*8=24 not divisible by 16; any last-dim size where the bit count doesn't divide evenly.

Common situations: Viewing buffers as a different-width dtype (u8->f16, f32->2xbf16) with a trailing dimension that isn't a multiple of the width ratio.

Related errors


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