jax-ml/jax · error · ValueError

Cannot concatenate vectors of different element types

Error message

Cannot concatenate vectors of different element types

What it means

vector_concat requires a uniform element type across all operands because the result type is derived from vectors[0]. Any operand whose element type differs raises this ValueError before concatenation.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2135

      [slice_length],
      [1],
  )


def vector_concat(
    vectors: Sequence[ir.Value[ir.VectorType]],
) -> ir.Value[ir.VectorType]:
  if not vectors:
    raise ValueError("Cannot concatenate an empty list of vectors")
  vty = vectors[0].type
  if not isinstance(vty, ir.VectorType):
    raise ValueError("Cannot concatenate non-vector values")
  vty = ir.VectorType(vty)
  if vty.rank != 1:
    raise NotImplementedError("Only 1D vectors are supported")
  for v in vectors:
    if v.type.element_type != vty.element_type:
      raise ValueError("Cannot concatenate vectors of different element types")
    if v.type.rank != 1:
      raise ValueError("Can only concatenate 1D vectors")
  return _vector_concat_rec(vectors)


def _vector_concat_rec(
    vectors: Sequence[ir.Value[ir.VectorType]],
) -> ir.Value[ir.VectorType]:
  match vectors:
    case [v]:
      return v
    case [v, w]:
      [v_len] = ir.VectorType(v.type).shape
      [w_len] = ir.VectorType(w.type).shape
      mask = ir.DenseI64ArrayAttr.get(list(range(v_len + w_len)))
      return vector.shuffle(*vectors, mask=mask)
    case _:
      assert vectors

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize element types first: bitcast if same bitwidth, else arith.extf/truncf/uitofp conversions
2. Broadcast/convert all operands to the result element type before the concat
3. Add per-operand asserts on .type.element_type during development

Example fix

# before
v = vector_concat([acc_f32, w_bf16])
# after
w_f32 = arith.extf(f32, w_bf16)
v = vector_concat([acc_f32, w_f32])

et = ir.VectorType(vectors[0].type).element_type
assert all(ir.VectorType(v.type).element_type == et for v in vectors), 'mixed element types'

When it happens

Trigger: Calling vector_concat with e.g. [vector<4xf32>, vector<4xbf16>] — mixed f32/bf16 or i32/f32 operands.

Common situations: Mixing values from an accumulator (f32) and reloaded weights (bf16) in an epilogue; one pipeline stage inserting an implicit upcast while another does not.

Related errors

• The stored value has dtype {src.dtype}, but the target refer
• Accumulator dtype {inner.dtype} does not match value dtype {
• All arrays must have the same dtype, got {arr.mlir_dtype} at
• Expected same element type, got {element_type} and {dst_ref_
• Reduction op {reduction_op} not supported by the TMA impleme

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