jax-ml/jax · error · ValueError

Expected input and output shapes are the same after multiply

Error message

Expected input and output shapes are the same after multiplying the second-minor dimension by the bitwidths.

What it means

jax._src.state.utils.bitcast raises ValueError when the second-minor dimension's total bits (shape[-2] * x_bitwidth) aren't evenly divisible by the target bitwidth, so no valid output shape exists for the width-changing bitcast. The comment notes this packing scheme is TPU-specific.

Source

Thrown at jax/_src/state/utils.py:94

  hoisted_jaxpr, _ = pe.trace_to_jaxpr(
      _hoist, ft.flatten_args(*in_avals),
      jaxpr.debug_info.with_unknown_names())
  assert not hoisted_jaxpr.consts, "All consts should have been converted to refs"
  return hoisted_jaxpr


def bitcast(x, dtype: DTypeLike):
  x_bitwidth = dtypes.itemsize_bits(x.dtype)
  y_bitwidth = dtypes.itemsize_bits(dtype)
  shape = list(x.shape)
  if x_bitwidth != y_bitwidth:
    if len(shape) < 2:
      raise NotImplementedError(
          "Bitcast 1D ref with bitwidth change is not supported."
      )
    # Note: this is only valid on TPU.
    if shape[-2] * x_bitwidth % y_bitwidth != 0:
      raise ValueError(
          "Expected input and output shapes are the same after multiplying"
          " the second-minor dimension by the bitwidths."
      )
  shape[-2] = shape[-2] * x_bitwidth // y_bitwidth
  if x_bitwidth < y_bitwidth:
    ratio = y_bitwidth // x_bitwidth
    x = x.reshape(*x.shape[:-2], x.shape[-2] // ratio, ratio, -1).swapaxes(
        -1, -2
    )
  y = lax.bitcast_convert_type(x, dtype)
  if x_bitwidth > y_bitwidth:
    y = y.swapaxes(-1, -2).reshape(shape)
  return y


def eval_bitcast_shape(x, dtype: DTypeLike):
  f = partial(bitcast, dtype=dtype)
  return api.eval_shape(f, api.ShapeDtypeStruct.like(x)).shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or trim shape[-2] so shape[-2]*old_bits is divisible by new_bits
2. Reshape so the second-minor dim is a multiple of (new_bits//old_bits) or vice versa
3. Use same-bitwidth dtypes to avoid the constraint

Example fix

# before
y = bitcast(x_f32_shape_5x4, jnp.f16)
# after
x = x.reshape(10, 2)  # make second-minor dim compatible
y = bitcast(x, jnp.f16)

from jax._src import dtypes
xb, yb = dtypes.itemsize_bits(x.dtype), dtypes.itemsize_bits(dtype)
assert xb == yb or (len(x.shape) >= 2 and x.shape[-2] * xb % yb == 0), "incompatible bitcast shape"

try:
    y = bitcast(x, dtype)
except ValueError as e:
    if "bitwidths" in str(e):
        pad = (-x.shape[-2]) % (dtypes.itemsize_bits(dtype) // dtypes.itemsize_bits(x.dtype))
        x = x.reshape(x.shape[-2] + pad, -1)
        y = bitcast(x, dtype)
    else:
        raise

When it happens

Trigger: bitcast(x, dtype) with differing bitwidths where shape[-2] * old_bits % new_bits != 0, e.g. a (5, 4) f32 ref bitcast to f16: 5*32=160 not divisible by 16... (concretely, dims whose bit total doesn't factor the new width).

Common situations: Bitcasting to wider/narrower dtypes on TPU without padding dimensions to a multiple of the width ratio.

Related errors

• Cannot bitcast from {x.dtype} ({old_bitwidth} bits) to {dtyp
• Attempting to convert array of shape {operand.shape} from {o
• lax.bitcast_convert_type does not support bool or complex va
• Block size must be a multiple of the input size. Got block {
• Every block dimension must be either a multiple or factor of

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