{"record":{"id":"4334cc83352832ab","repo":"jax-ml/jax","slug":"the-bitwidth-of-x-must-match-the-bitwidth-of-un","errorCode":null,"errorMessage":"The bitwidth of `x` must match the bitwidth of `unpacked_dtype` for unpack_elementwise, but got {x.dtype} and {unpacked_dtype}","messagePattern":"The bitwidth of `x` must match the bitwidth of `unpacked_dtype` for unpack_elementwise, but got (.+?) and (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1040,"sourceCode":"    x: The packed array.\n    index: The index of the element to unpack.\n    packed_dtype: Elements\n    unpacked_dtype: The dtype of the unpacked array.\n\n  Returns:\n    The unpacked array in `unpacked_dtype`.\n  \"\"\"\n  return unpack_elementwise_p.bind(\n      x, index=index, packed_dtype=packed_dtype, unpacked_dtype=unpacked_dtype\n  )\n\n\n@unpack_elementwise_p.def_abstract_eval\ndef _unpack_elementwise_abstract_eval(\n    x, *, index, packed_dtype, unpacked_dtype\n):\n  if dtypes.itemsize_bits(x.dtype) != dtypes.itemsize_bits(unpacked_dtype):\n    raise ValueError(\n        \"The bitwidth of `x` must match the bitwidth of `unpacked_dtype` for \"\n        f\"unpack_elementwise, but got {x.dtype} and {unpacked_dtype}\"\n    )\n  packing_factor = _get_elementwise_packing_factor(unpacked_dtype, packed_dtype)\n  if index < 0 or index >= packing_factor:\n    raise ValueError(\n        f\"Index {index} is out of bounds for packing factor {packing_factor}\")\n  return jax_core.ShapedArray(x.shape, unpacked_dtype)\n\n\ndef with_memory_space_constraint(\n    x: jax.Array, memory_space: Any\n) -> jax.Array:\n  \"\"\"Constrains the memory space of an array.\n\n  This primitive does not change the value of ``x``, but it constrains the\n  memory space where it should be allocated. This is useful to force\n  Pallas to allocate an array in a specific memory space.","sourceCodeStart":1022,"sourceCodeEnd":1058,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1022-L1058","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Make unpacked_dtype match the bitwidth of the packed source dtype used when packing (e.g. uint32 <-> f32/int32, uint8 <-> int8/bf16-element)","If you want int8 outputs, pack int8 sources so the packed block is uint8"],"exampleFix":"# before\nunpacked = unpack_elementwise(block_u32, index=0, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.int8)\n# after (f32 sources packed into uint32)\nunpacked = unpack_elementwise(block_u32, index=0, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.float32)","handlingStrategy":"type-guard","validationCode":"from jax._src import dtypes\nassert dtypes.itemsize_bits(x.dtype) == dtypes.itemsize_bits(unpacked_dtype)","typeGuard":"def bitwidths_match(x, unpacked_dtype) -> bool:\n    return dtypes.itemsize_bits(x.dtype) == dtypes.itemsize_bits(unpacked_dtype)","tryCatchPattern":null,"preventionTips":["Derive unpacked_dtype from the original packed-source dtype so bitwidths always align","Centralize pack/unpack dtype pairs in helpers"],"tags":["jax","pallas","tpu","dtype","bitwidth"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}