jax-ml/jax · error · ValueError

Attempting to convert array of shape {operand.shape} from {o

Error message

Attempting to convert array of shape {operand.shape} from {old_dtype} of size {old_nbits} bits to {new_dtype} of size {new_nbits}, bits but {dim_size} * {old_nbits} != {new_nbits}

What it means

bitcast_convert_type's output shape rule: when upsizing (old_nbits < new_nbits, e.g. uint8 -> uint16), the last dimension of the operand must exactly supply the bits of one new element: dim_size * old_nbits == new_nbits. Otherwise the bit-level reinterpretation is impossible and ValueError is raised.

Source

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

    lambda ct, x, dtype: [to_edtype_p.bind(ct, edtype=x.dtype)]
batching.defvectorized(from_edtype_p)
mlir.register_lowering(from_edtype_p, lambda _, x, **__: [x])


def _bitcast_convert_type_shape_rule(operand, *, new_dtype):
  old_dtype = operand.dtype

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

  if old_nbits == new_nbits:
    return operand.shape
  elif old_nbits > new_nbits:
    return (*operand.shape, old_nbits // new_nbits)
  else:
    dim_size = operand.shape[-1] if operand.shape else 1
    if dim_size * old_nbits != new_nbits:
      raise ValueError(
        f"Attempting to convert array of shape {operand.shape} "
        f"from {old_dtype} of size {old_nbits} bits "
        f"to {new_dtype} of size {new_nbits}, bits "
        f"but {dim_size} * {old_nbits} != {new_nbits}")
    return operand.shape[:-1]

def _bitcast_convert_type_sharding_rule(operand, *, new_dtype):
  old_dtype = operand.dtype

  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])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape so the last dimension groups old elements into whole new elements (e.g. (10,) uint8 -> (2, 5) or (10//4, 4) for uint32)
Pick a new_dtype of equal bit width (shape preserved) or narrower (shape gains a trailing axis)
Pad the last dim to a multiple of new_nbits/old_nbits if semantics allow
Handle the scalar case by adding a trailing axis of the required group size

Example fix

// before
y = lax.bitcast_convert_type(x_uint8.reshape(10), jnp.uint32)  # 10*8 != 32

// after
y = lax.bitcast_convert_type(x_uint8.reshape(10, 1), jnp.uint8)  # same width
# or group: x_uint8.reshape(2, 5) is invalid; use (…, 4)->uint32
y = lax.bitcast_convert_type(x_uint8[:8].reshape(2, 4), jnp.uint32)

Defensive patterns

Strategy: validation

Validate before calling

old_b, new_b = np.dtype(x.dtype).itemsize*8, np.dtype(new_dtype).itemsize*8
if old_b < new_b:
    group = new_b // old_b
    assert x.shape[-1] % group == 0 if x.ndim else False or new_b == old_b, 'regroup last axis'

Try / catch

try:
    y = lax.bitcast_convert_type(x, dt)
except ValueError:
    g = new_bits // old_bits
    y = lax.bitcast_convert_type(x.reshape(*x.shape[:-1], -1, g), dt)

Prevention

Compute the bit-group size before widening bitcasts
Reshape to (..., new/old) grouped form as a standard preprocessing step
Prefer equal-width bitcasts when shape stability matters

When it happens

Trigger: jax.lax.bitcast_convert_type(x, jnp.uint32) where x is uint8 with last-dim size not divisible by 4 (e.g. shape (10,)), or bitcasting a scalar to a wider type.

Common situations: Packing/unpacking sub-byte data (uint8 views of float4/float8 payloads); resizing tensors before bitcast; forgetting that bitcast to wider types consumes the last axis.

Related errors

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