jax-ml/jax · error · ValueError

The number of sources must match the packing factor ({packin

Error message

The number of sources must match the packing factor ({packing_factor}), got {len(xs)}

What it means

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.

Source

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

def _pack_elementwise_abstract_eval(*xs, packed_dtype):
  if not xs:
    raise ValueError("At least one source is required")
  first = xs[0]
  if not all(x.shape == first.shape for x in xs):
    raise ValueError("All sources must have the same shape")
  if not all(x.dtype == first.dtype for x in xs):
    raise ValueError("All sources must have the same dtype")
  if not (first.dtype == jnp.float32 and packed_dtype == jnp.bfloat16) and not (
      jnp.issubdtype(first.dtype, jnp.integer)
      and jnp.issubdtype(packed_dtype, jnp.integer)
  ):
    raise ValueError(
        "Only f32 -> bf16 and int -> int are supported. Got"
        f" {first.dtype} and {packed_dtype}"
    )
  packing_factor = _get_elementwise_packing_factor(first.dtype, packed_dtype)
  if len(xs) != packing_factor:
    raise ValueError(
        "The number of sources must match the packing factor "
        f"({packing_factor}), got {len(xs)}"
    )
  out_dtype = jnp.dtype(f"uint{dtypes.itemsize_bits(first.dtype)}")
  return jax_core.ShapedArray(first.shape, out_dtype)


unpack_elementwise_p = jax_core.Primitive("unpack_elementwise")


def unpack_elementwise(x, *, index, packed_dtype, unpacked_dtype):
  """Unpacks an elementwise packed array.

  The function follows the *interleaved format* during unpacking, and it's the
  reverse of `pack_elementwise`.

  For example, if `packed_dtype` is `int4`, `unpacked_dtype` is `int8`,
  and `x` is packed `int8` with x'y'z'w'm'n'i'j' in a word, where each

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the factor: packing_factor = itemsize_bits(src_dtype) // itemsize_bits(packed_dtype) and pass exactly that many arrays
2. For f32 -> bf16 pass 4 source arrays
3. For int32 -> int8 pass 4 source arrays; int8->int8 passes 1

Example fix

# before
packed = pack_elementwise(x0, x1, packed_dtype=jnp.bfloat16)  # factor is 4, not 2
# after
packed = pack_elementwise(x0, x1, x2, x3, packed_dtype=jnp.bfloat16)

from jax._src import dtypes
factor = dtypes.itemsize_bits(src_dtype) // dtypes.itemsize_bits(packed_dtype)
assert len(xs) == factor, f'need {factor} sources for {src_dtype}->{packed_dtype}'

def is_valid_pack_sources(xs, src_dtype, packed_dtype) -> bool:
    factor = dtypes.itemsize_bits(src_dtype) // dtypes.itemsize_bits(packed_dtype)
    return len(xs) == factor and all(getattr(x, 'dtype', None) == src_dtype for x in xs)

When it happens

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

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

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• Can only load scalars from SMEM
• masked swap with non-32-bit data
• Expected value and mask to have the same shape, but got valu
• Can only store scalars to SMEM

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