jax-ml/jax · error · ValueError

All input arrays must have the same shape. Got {}.

Error message

All input arrays must have the same shape. Got {}.

What it means

All inputs to jnp.stack must have exactly the same shape, since stack builds shape + new axis of length num_operands. JAX compares the set of operand shapes and raises this ValueError showing each shape if they differ (no broadcasting, unlike NumPy in some cases).

Source

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

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)
    raise core.ShardingTypeError(
        f"All operands should have the same sharding. Got shardings {ss}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or truncate inputs to a common shape before stacking
2. Verify model config so all stacked tensors come from identically-shaped sources
3. If lengths legitimately differ, use a list/pytree or jax.tree_util instead of stacking

Example fix

# before
out = jnp.stack(seqs, axis=0)  # seqs have varying lengths
# after
maxlen = max(s.shape[0] for s in seqs)
seqs = [jnp.pad(s, (0, maxlen - s.shape[0])) for s in seqs]
out = jnp.stack(seqs, axis=0)

ref = arrays[0].shape
assert all(a.shape == ref for a in arrays), [a.shape for a in arrays]

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

When it happens

Trigger: jnp.stack([jnp.zeros((2,3)), jnp.zeros((2,4))], axis=0); stacking variable-length sequences; stacking tensors from layers with mismatched output dims.

Common situations: Stacking per-sequence states of different lengths; mistyped hidden sizes across model config; results of different tokenizers/paddings.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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