jax-ml/jax · error · ValueError

Index {index} is out of bounds for packing factor {packing_f

Error message

Index {index} is out of bounds for packing factor {packing_factor}

What it means

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).

Source

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

    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.

  As of now, this only operates on the inputs pallas_calls, as in you can
  apply this to the arguments of a pallas_call and it will constrain them, but
  other operations will not respect this constraint.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute the packing factor and clamp/iterate: for i in range(_packing_factor(unpacked, packed))
2. Check index >= 0 and index < packing_factor before calling

Example fix

# before
lane = unpack_elementwise(block, index=4, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.float32)
# after
lane = unpack_elementwise(block, index=i % 4, packed_dtype=jnp.bfloat16, unpacked_dtype=jnp.float32)

factor = itemsize_bits(unpacked_dtype) // itemsize_bits(packed_dtype)
assert 0 <= index < factor

def valid_lane_index(index, factor) -> bool:
    return isinstance(index, int) and 0 <= index < factor

When it happens

Trigger: Calling unpack_elementwise with index outside [0, packing_factor), e.g. index=4 or index=-1 when the packing factor is 4.

Common situations: 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.

Related errors

• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Axis {axis} is out of bounds for grid {self.grid}
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}

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