jax-ml/jax · error · ValueError

The bitwidth of `x` must match the bitwidth of `unpacked_dty

Error message

The bitwidth of `x` must match the bitwidth of `unpacked_dtype` for unpack_elementwise, but got {x.dtype} and {unpacked_dtype}

What it means

Raised by the abstract eval of unpack_elementwise in JAX Mosaic Pallas. The input array x must have the same total bitwidth as the declared unpacked_dtype (e.g. a uint32 input unpacks to f32; a uint8 to int8), because unpacking reinterprets the packed bits lane by lane.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:1040

    x: The packed array.
    index: The index of the element to unpack.
    packed_dtype: Elements
    unpacked_dtype: The dtype of the unpacked array.

  Returns:
    The unpacked array in `unpacked_dtype`.
  """
  return unpack_elementwise_p.bind(
      x, index=index, packed_dtype=packed_dtype, unpacked_dtype=unpacked_dtype
  )


@unpack_elementwise_p.def_abstract_eval
def _unpack_elementwise_abstract_eval(
    x, *, index, packed_dtype, unpacked_dtype
):
  if dtypes.itemsize_bits(x.dtype) != dtypes.itemsize_bits(unpacked_dtype):
    raise ValueError(
        "The bitwidth of `x` must match the bitwidth of `unpacked_dtype` for "
        f"unpack_elementwise, but got {x.dtype} and {unpacked_dtype}"
    )
  packing_factor = _get_elementwise_packing_factor(unpacked_dtype, packed_dtype)
  if index < 0 or index >= packing_factor:
    raise ValueError(
        f"Index {index} is out of bounds for packing factor {packing_factor}")
  return jax_core.ShapedArray(x.shape, unpacked_dtype)


def with_memory_space_constraint(
    x: jax.Array, memory_space: Any
) -> jax.Array:
  """Constrains the memory space of an array.

  This primitive does not change the value of ``x``, but it constrains the
  memory space where it should be allocated. This is useful to force
  Pallas to allocate an array in a specific memory space.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make unpacked_dtype match the bitwidth of the packed source dtype used when packing (e.g. uint32 <-> f32/int32, uint8 <-> int8/bf16-element)
2. If you want int8 outputs, pack int8 sources so the packed block is uint8

Example fix

# before
unpacked = unpack_elementwise(block_u32, index=0, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.int8)
# after (f32 sources packed into uint32)
unpacked = unpack_elementwise(block_u32, index=0, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.float32)

from jax._src import dtypes
assert dtypes.itemsize_bits(x.dtype) == dtypes.itemsize_bits(unpacked_dtype)

def bitwidths_match(x, unpacked_dtype) -> bool:
    return dtypes.itemsize_bits(x.dtype) == dtypes.itemsize_bits(unpacked_dtype)

When it happens

Trigger: Calling unpack_elementwise(x, index=..., packed_dtype=..., unpacked_dtype=...) where itemsize_bits(x.dtype) != itemsize_bits(unpacked_dtype), e.g. passing a uint32 block with unpacked_dtype=jnp.int8 (32 vs 8 bits).

Common situations: Feeding a block whose dtype was chosen by the compiler (typically uint{bits}) while declaring a narrower/wider unpacked_dtype; mismatch between the dtype used in the preceding pack_elementwise and the later unpack.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot

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