jax-ml/jax · error · TypeError

Cannot concatenate arrays with different numbers of dimensio

Error message

Cannot concatenate arrays with different numbers of dimensions: got {}.

What it means

Concatenation only joins arrays along one axis, so all operands must have the same number of dimensions. JAX checks that the set of ndims has size 1 and otherwise raises this TypeError listing every shape, since NumPy-style broadcasting is not applied by lax.concatenate.

Source

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

          select(bitwise_and(gt(operand, min), lt(operand, max)),
                 g, _zeros(operand)),
          lambda g, min, operand, max:
          select(lt(max, operand), g, _zeros(operand)))
batching.primitive_batchers[clamp_p] = _clamp_batch_rule
mlir.register_lowering(clamp_p, partial(_nary_lower_hlo, hlo.clamp))

def _concatenate_shape_rule(*operands, **kwargs):
  dimension = kwargs.pop('dimension')
  if not operands:
    msg = "concatenate expects at least one operand, got 0."
    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]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add or remove axes so all arrays match rank: use x[None, :] / jnp.expand_dims or jnp.atleast_2d
2. Use jnp.stack instead if you want to add a new axis
3. Verify intermediate shapes with prints or jax.debug.print before the concat

Example fix

# before
out = jnp.concatenate([batch, row], axis=0)  # row has shape (n,)
# after
out = jnp.concatenate([batch, row[None, :]], axis=0)

if len({a.ndim for a in arrays}) != 1:
    arrays = [jnp.atleast_2d(a) for a in arrays]
out = jnp.concatenate(arrays, axis=0)

def same_rank(xs) -> bool:
    nd = xs[0].ndim if xs else None
    return all(x.ndim == nd for x in xs)

When it happens

Trigger: jnp.concatenate([jnp.zeros((3,)), jnp.zeros((2,3))], axis=0) — mixing rank-1 and rank-2 arrays; concatenating scalars with vectors.

Common situations: Appending a scalar or 1-D row to a 2-D batch without reshaping; mixed data pipelines where some tensors went through squeeze/reshape.

Related errors

• convolution requires lhs and rhs ndim to be equal, got {} an
• can only convert to extended dtype from an array of its repr
• concatenate expects at least one operand, got 0.
• concatenate dimension out of bounds: dimension {} for shapes
• Cannot concatenate arrays with shapes that differ in dimensi

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