{"record":{"id":"c1592f69ea66faf9","repo":"jax-ml/jax","slug":"cannot-stack-arrays-with-different-numbers-of-dime","errorCode":null,"errorMessage":"Cannot stack arrays with different numbers of dimensions: got {}.","messagePattern":"Cannot stack arrays with different numbers of dimensions: got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7350,"sourceCode":"                  for i in range(0, len(current_xs), k)]\n  return current_xs[0]\n\ndef _concatenate_lower(ctx, *xs, dimension):\n  aval_out, = ctx.avals_out\n  out = _concatenate_tree(xs, dimension)\n  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]\n\nmlir.register_lowering(concatenate_p, _concatenate_lower)\n\n# --- stack and unstack primitives ---\n\ndef _stack_shape_rule(*operands, axis):\n  if not operands:\n    msg = \"stack expects at least one operand, got 0.\"\n    raise ValueError(msg)\n  if len({op.ndim for op in operands}) != 1:\n    msg = \"Cannot stack arrays with different numbers of dimensions: got {}.\"\n    raise ValueError(msg.format(\", \".join(str(o.shape) for o in operands)))\n  if len({op.shape for op in operands}) != 1:\n    msg = \"All input arrays must have the same shape. Got {}.\"\n    raise ValueError(msg.format(\", \".join(str(o.shape) for o in operands)))\n\n  shape = list(operands[0].shape)\n  shape.insert(axis, len(operands))\n  return tuple(shape)\n\ndef _stack_dtype_rule(*operands, axis):\n  check_same_dtypes('stack', *operands)\n  return operands[0].dtype\n\ndef _stack_sharding_rule(*operands, axis):\n  non_empty_s = [o.sharding for o in operands if not o.sharding.mesh.empty]\n  if not non_empty_s:\n    return core.get_cur_mesh_sharding()\n  if not all(s == non_empty_s[0] for s in non_empty_s):\n    ss = \", \".join(str(o.sharding) for o in operands)","sourceCodeStart":7332,"sourceCodeEnd":7368,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7332-L7368","documentation":"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.","triggerScenarios":"jnp.stack([jnp.zeros(3), jnp.zeros((2,3))]); stacking a scalar with vectors; mixing outputs of layers with different ranks.","commonSituations":"Stacking model outputs where one branch was squeezed; list of per-step states where some are scalars; mixing raw scalars with traced arrays.","solutions":["Normalize ranks: wrap scalars/vectors with jnp.atleast_ndim(x, n) or x[None] before stacking","Fix upstream squeeze/reshape calls that dropped an axis inconsistently","Use jnp.stack on a uniformly-shaped list produced by vmap instead of manual loops"],"exampleFix":"# before\nout = jnp.stack([xs, total], axis=0)  # xs:(n,), total scalar\n# after\nout = jnp.stack([xs, jnp.broadcast_to(total, xs.shape)], axis=0)","handlingStrategy":"validation","validationCode":"nd = max(a.ndim for a in arrays)\narrays = [jnp.atleast_ndim(a, nd) for a in arrays]\nout = jnp.stack(arrays, axis=0)","typeGuard":"def uniform_rank(xs) -> bool:\n    return len({x.ndim for x in xs}) == 1","tryCatchPattern":null,"preventionTips":["Normalize ranks with atleast_ndim before stacking","Avoid mixing scalars and arrays in stacked collections"],"tags":["jax","stack","rank-mismatch","shape-validation"],"backgroundTag":"ndim-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}