{"record":{"id":"1bcf1732fd8bab35","repo":"jax-ml/jax","slug":"the-number-of-sources-must-match-the-packing-facto","errorCode":null,"errorMessage":"The number of sources must match the packing factor ({packing_factor}), got {len(xs)}","messagePattern":"The number of sources must match the packing factor \\((.+?)\\), got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":982,"sourceCode":"def _pack_elementwise_abstract_eval(*xs, packed_dtype):\n  if not xs:\n    raise ValueError(\"At least one source is required\")\n  first = xs[0]\n  if not all(x.shape == first.shape for x in xs):\n    raise ValueError(\"All sources must have the same shape\")\n  if not all(x.dtype == first.dtype for x in xs):\n    raise ValueError(\"All sources must have the same dtype\")\n  if not (first.dtype == jnp.float32 and packed_dtype == jnp.bfloat16) and not (\n      jnp.issubdtype(first.dtype, jnp.integer)\n      and jnp.issubdtype(packed_dtype, jnp.integer)\n  ):\n    raise ValueError(\n        \"Only f32 -> bf16 and int -> int are supported. Got\"\n        f\" {first.dtype} and {packed_dtype}\"\n    )\n  packing_factor = _get_elementwise_packing_factor(first.dtype, packed_dtype)\n  if len(xs) != packing_factor:\n    raise ValueError(\n        \"The number of sources must match the packing factor \"\n        f\"({packing_factor}), got {len(xs)}\"\n    )\n  out_dtype = jnp.dtype(f\"uint{dtypes.itemsize_bits(first.dtype)}\")\n  return jax_core.ShapedArray(first.shape, out_dtype)\n\n\nunpack_elementwise_p = jax_core.Primitive(\"unpack_elementwise\")\n\n\ndef unpack_elementwise(x, *, index, packed_dtype, unpacked_dtype):\n  \"\"\"Unpacks an elementwise packed array.\n\n  The function follows the *interleaved format* during unpacking, and it's the\n  reverse of `pack_elementwise`.\n\n  For example, if `packed_dtype` is `int4`, `unpacked_dtype` is `int8`,\n  and `x` is packed `int8` with x'y'z'w'm'n'i'j' in a word, where each","sourceCodeStart":964,"sourceCodeEnd":1000,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L964-L1000","documentation":"Raised by the abstract eval of pack_elementwise in JAX's TPU Mosaic Pallas primitives. The number of source arrays passed to pack_elementwise must exactly equal the packing factor implied by the source and packed dtypes (e.g. 4 sources for f32->bf16 since 32/8=4). Any other count is rejected at tracing time.","triggerScenarios":"Calling pack_elementwise(*xs, packed_dtype=...) where len(xs) != _get_elementwise_packing_factor(xs[0].dtype, packed_dtype), e.g. passing 2 or 8 f32 arrays with packed_dtype=bfloat16 (factor is 4).","commonSituations":"Assuming the packing factor is 2 instead of 4 for f32->bf16; packing int32 to int8 (factor 4) with the wrong number of operands; changing packed_dtype without adjusting the number of sources.","solutions":["Check the factor: packing_factor = itemsize_bits(src_dtype) // itemsize_bits(packed_dtype) and pass exactly that many arrays","For f32 -> bf16 pass 4 source arrays","For int32 -> int8 pass 4 source arrays; int8->int8 passes 1"],"exampleFix":"# before\npacked = pack_elementwise(x0, x1, packed_dtype=jnp.bfloat16)  # factor is 4, not 2\n# after\npacked = pack_elementwise(x0, x1, x2, x3, packed_dtype=jnp.bfloat16)","handlingStrategy":"validation","validationCode":"from jax._src import dtypes\nfactor = dtypes.itemsize_bits(src_dtype) // dtypes.itemsize_bits(packed_dtype)\nassert len(xs) == factor, f'need {factor} sources for {src_dtype}->{packed_dtype}'","typeGuard":"def is_valid_pack_sources(xs, src_dtype, packed_dtype) -> bool:\n    factor = dtypes.itemsize_bits(src_dtype) // dtypes.itemsize_bits(packed_dtype)\n    return len(xs) == factor and all(getattr(x, 'dtype', None) == src_dtype for x in xs)","tryCatchPattern":null,"preventionTips":["Compute the packing factor from dtype bitwidths before assembling sources","Keep pack/unpack dtype pairs in one config constant"],"tags":["jax","pallas","tpu","dtype","shape-mismatch"],"backgroundTag":"argument-count-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}