{"record":{"id":"836be4c002e77734","repo":"jax-ml/jax","slug":"all-input-arrays-must-have-the-same-shape-got","errorCode":null,"errorMessage":"All input arrays must have the same shape. Got {}.","messagePattern":"All input arrays must have the same shape\\. Got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7353,"sourceCode":"def _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)\n    raise core.ShardingTypeError(\n        f\"All operands should have the same sharding. Got shardings {ss}\")\n","sourceCodeStart":7335,"sourceCodeEnd":7371,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7335-L7371","documentation":"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).","triggerScenarios":"jnp.stack([jnp.zeros((2,3)), jnp.zeros((2,4))], axis=0); stacking variable-length sequences; stacking tensors from layers with mismatched output dims.","commonSituations":"Stacking per-sequence states of different lengths; mistyped hidden sizes across model config; results of different tokenizers/paddings.","solutions":["Pad or truncate inputs to a common shape before stacking","Verify model config so all stacked tensors come from identically-shaped sources","If lengths legitimately differ, use a list/pytree or jax.tree_util instead of stacking"],"exampleFix":"# before\nout = jnp.stack(seqs, axis=0)  # seqs have varying lengths\n# after\nmaxlen = max(s.shape[0] for s in seqs)\nseqs = [jnp.pad(s, (0, maxlen - s.shape[0])) for s in seqs]\nout = jnp.stack(seqs, axis=0)","handlingStrategy":"validation","validationCode":"ref = arrays[0].shape\nassert all(a.shape == ref for a in arrays), [a.shape for a in arrays]","typeGuard":"def uniform_shape(xs) -> bool:\n    return len({x.shape for x in xs}) == 1","tryCatchPattern":null,"preventionTips":["Pad variable-length inputs to a common shape before stacking","Assert shape uniformity in sequence-pipeline unit tests"],"tags":["jax","stack","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}