jax-ml/jax · error · ValueError

Cannot stack arrays with different numbers of dimensions: go

Error message

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

What it means

jnp.stack inserts a new axis, so every input must have identical rank. The shape rule checks that the set of operand ndims has exactly one element and otherwise raises this ValueError listing all shapes. Unlike concatenate, stack requires fully matching shapes, not just matching rank.

Source

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

                  for i in range(0, len(current_xs), k)]
  return current_xs[0]

def _concatenate_lower(ctx, *xs, dimension):
  aval_out, = ctx.avals_out
  out = _concatenate_tree(xs, dimension)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(concatenate_p, _concatenate_lower)

# --- stack and unstack primitives ---

def _stack_shape_rule(*operands, axis):
  if not operands:
    msg = "stack expects at least one operand, got 0."
    raise ValueError(msg)
  if len({op.ndim for op in operands}) != 1:
    msg = "Cannot stack arrays with different numbers of dimensions: got {}."
    raise ValueError(msg.format(", ".join(str(o.shape) for o in operands)))
  if len({op.shape for op in operands}) != 1:
    msg = "All input arrays must have the same shape. Got {}."
    raise ValueError(msg.format(", ".join(str(o.shape) for o in operands)))

  shape = list(operands[0].shape)
  shape.insert(axis, len(operands))
  return tuple(shape)

def _stack_dtype_rule(*operands, axis):
  check_same_dtypes('stack', *operands)
  return operands[0].dtype

def _stack_sharding_rule(*operands, axis):
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize ranks: wrap scalars/vectors with jnp.atleast_ndim(x, n) or x[None] before stacking
2. Fix upstream squeeze/reshape calls that dropped an axis inconsistently
3. Use jnp.stack on a uniformly-shaped list produced by vmap instead of manual loops

Example fix

# before
out = jnp.stack([xs, total], axis=0)  # xs:(n,), total scalar
# after
out = jnp.stack([xs, jnp.broadcast_to(total, xs.shape)], axis=0)

nd = max(a.ndim for a in arrays)
arrays = [jnp.atleast_ndim(a, nd) for a in arrays]
out = jnp.stack(arrays, axis=0)

def uniform_rank(xs) -> bool:
    return len({x.ndim for x in xs}) == 1

When it happens

Trigger: jnp.stack([jnp.zeros(3), jnp.zeros((2,3))]); stacking a scalar with vectors; mixing outputs of layers with different ranks.

Common situations: Stacking model outputs where one branch was squeezed; list of per-step states where some are scalars; mixing raw scalars with traced arrays.

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
• Cannot concatenate arrays with different numbers of dimensio
• stack expects at least one operand, got 0.
• convolve2d() only supports 2-dimensional inputs.

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