{"record":{"id":"c8b4e9c688b37243","repo":"jax-ml/jax","slug":"index-index-is-out-of-bounds-for-packing-factor","errorCode":null,"errorMessage":"Index {index} is out of bounds for packing factor {packing_factor}","messagePattern":"Index (.+?) is out of bounds for packing factor (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1046,"sourceCode":"    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.\n\n  As of now, this only operates on the inputs pallas_calls, as in you can\n  apply this to the arguments of a pallas_call and it will constrain them, but\n  other operations will not respect this constraint.\n\n  Args:","sourceCodeStart":1028,"sourceCodeEnd":1064,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1028-L1064","documentation":"Raised by the abstract eval of unpack_elementwise in JAX Mosaic Pallas. The index selects which lane of the packed word to extract, so it must satisfy 0 <= index < packing_factor (e.g. 0..3 for f32->bf16 packing).","triggerScenarios":"Calling unpack_elementwise with index outside [0, packing_factor), e.g. index=4 or index=-1 when the packing factor is 4.","commonSituations":"Hardcoding index values after changing packed_dtype (factor changes from 4 to 2 or 1); looping over range(packing_factor) computed for a different dtype pair.","solutions":["Compute the packing factor and clamp/iterate: for i in range(_packing_factor(unpacked, packed))","Check index >= 0 and index < packing_factor before calling"],"exampleFix":"# before\nlane = unpack_elementwise(block, index=4, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.float32)\n# after\nlane = unpack_elementwise(block, index=i % 4, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.float32)","handlingStrategy":"validation","validationCode":"factor = itemsize_bits(unpacked_dtype) // itemsize_bits(packed_dtype)\nassert 0 <= index < factor","typeGuard":"def valid_lane_index(index, factor) -> bool:\n    return isinstance(index, int) and 0 <= index < factor","tryCatchPattern":null,"preventionTips":["Iterate lanes with range(packing_factor) instead of hardcoded indices","Use index % packing_factor when indices come from loops"],"tags":["jax","pallas","tpu","index-out-of-bounds"],"backgroundTag":"index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}