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).

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

• A must be ( , ) matrix A, got output
• Accumulator and LHS have incompatible shapes. Expected LHS…
• Accumulator and RHS have incompatible shapes. Expected RHS…
• Accumulator and RHS have incompatible shapes. Expected RHS…
• Accumulator shape does not match value shape

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

Appendix: 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)