jax-ml/jax · error · TypeError

Cannot concatenate arrays with shapes that differ in dimensi

Error message

Cannot concatenate arrays with shapes that differ in dimensions other than the one being concatenated: concatenating along dimension {} for shapes {}.

What it means

For concatenation along `dimension`, every non-concatenated axis must match exactly across operands. JAX computes each shape with the concat axis removed and raises this TypeError if they differ — lax.concatenate never broadcasts.

Source

Thrown at jax/_src/lax/lax.py:7249

    raise TypeError(msg)
  if not all(isinstance(operand, ShapedArray) for operand in operands):
    msg = "All objects to concatenate must be arrays, got {}."
    op = next(op for op in operands if not isinstance(op, ShapedArray))
    raise TypeError(msg.format(type(op)))
  if len({operand.ndim for operand in operands}) != 1:
    msg = "Cannot concatenate arrays with different numbers of dimensions: got {}."
    raise TypeError(msg.format(", ".join(str(o.shape) for o in operands)))
  if not 0 <= dimension < operands[0].ndim:
    msg = "concatenate dimension out of bounds: dimension {} for shapes {}."
    raise TypeError(msg.format(dimension, ", ".join([str(o.shape) for o in operands])))
  shapes = [operand.shape[:dimension] + operand.shape[dimension+1:]
            for operand in operands]
  if shapes[:-1] != shapes[1:]:
    msg = ("Cannot concatenate arrays with shapes that differ in dimensions "
           "other than the one being concatenated: concatenating along "
           "dimension {} for shapes {}.")
    shapes = [operand.shape for operand in operands]
    raise TypeError(msg.format(dimension, ", ".join(map(str, shapes))))

  concat_size = sum(o.shape[dimension] for o in operands)
  ex_shape = operands[0].shape
  return ex_shape[:dimension] + (concat_size,) + ex_shape[dimension+1:]

def _concatenate_sharding_rule(*operands, **kwargs):
  non_empty_s = [o.sharding for o in operands if not o.sharding.mesh.empty]
  if not non_empty_s:
    return core.get_cur_mesh_sharding()
  if not all(s == non_empty_s[0] for s in non_empty_s):
    ss = ", ".join(str(o.sharding) for o in operands)
    raise core.ShardingTypeError(
        f"All operands should have the same sharding. Got shardings {ss}")
  return non_empty_s[0]

def _concatenate_reduced_rule(*operands, **kwargs):
  reduced_specs = {r for o in operands if (r := getr(o))}
  if len(reduced_specs) > 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print/assert all shapes except the concat axis match before concatenating
2. Fix the upstream producer so the non-concat dimensions agree (correct feature size, transpose, or reshape)
3. For ragged data, pad to a common size or store as a list instead of concatenating

Example fix

# before
out = jnp.concatenate([a, b], axis=0)  # a:(2,3), b:(2,4)
# after
assert a.shape[1:] == b.shape[1:], (a.shape, b.shape)
out = jnp.concatenate([a, b], axis=0)

ref = arrays[0].shape[:axis] + arrays[0].shape[axis+1:]
assert all(a.shape[:axis] + a.shape[axis+1:] == ref for a in arrays), \
    [a.shape for a in arrays]
out = jnp.concatenate(arrays, axis=axis)

def concat_compatible(xs, axis) -> bool:
    return len({x.shape[:axis] + x.shape[axis+1:] for x in xs}) == 1

When it happens

Trigger: jnp.concatenate([jnp.zeros((2,3)), jnp.zeros((2,4))], axis=0) — second axis 3 vs 4; feature dimensions differing across batches.

Common situations: Concatenating sequences of embeddings/timesteps where a feature dim was mistyped; ragged per-sample data naively converted to arrays; unit drift after reshapes or transposes upstream.

Related errors

• concatenate expects at least one operand, got 0.
• Cannot concatenate arrays with different numbers of dimensio
• concatenate dimension out of bounds: dimension {} for shapes
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out

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