{"record":{"id":"57e03208f7782bee","repo":"jax-ml/jax","slug":"arrays-must-have-same-number-of-dimensions-go","errorCode":null,"errorMessage":"{}: arrays must have same number of dimensions, got {}.","messagePattern":"(.+?): arrays must have same number of dimensions, got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4355,"sourceCode":"    raise RuntimeError(\n      \"First argument of broadcasting_sharding_rule should be a name.\"\n      f\" Got {name}\")\n  mesh = None\n  for a in avals:\n    if a.sharding is not None and not a.sharding.mesh.empty:\n      if mesh is not None and mesh != a.sharding.mesh:\n        raise core.ShardingTypeError(\n            f'Mesh for all inputs should be equal. Got one mesh: {mesh} and'\n            f' another mesh: {a.sharding.mesh}')\n      mesh = a.sharding.mesh\n  mesh = get_abstract_mesh() if mesh is None else mesh\n\n  shapes = [aval.shape for aval in avals if aval.shape]\n  if not shapes:\n    return NamedSharding(mesh, P())\n  if len({len(shape) for shape in shapes}) != 1:\n    msg = '{}: arrays must have same number of dimensions, got {}.'\n    raise TypeError(msg.format(name, ', '.join(map(str, map(tuple, shapes)))))\n\n  specs = [a.sharding.spec.partitions for a in avals if a.shape]\n\n  result_specs = [None] * len(shapes[0])\n  for i, (ss, ds) in enumerate(zip(zip(*specs), zip(*shapes))):\n    if all(ss[0] == s for s in ss[1:]):\n      # if all dimension shardings are same, the resulting dimension sharding is\n      # the same.\n      result_specs[i] = ss[0]\n    else:\n      non_trivial_s = [s for s, d in zip(ss, ds)\n                       if not (core.definitely_equal(d, 1) and s is None)]\n      if not non_trivial_s:\n        result_specs[i] = None\n      elif all(non_trivial_s[0] == s for s in non_trivial_s[1:]):\n        result_specs[i] = non_trivial_s[0]\n      else:\n        for s in ss:","sourceCodeStart":4337,"sourceCodeEnd":4373,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L4337-L4373","documentation":"In sharding-propagation logic for NamedSharding, all operands of an n-ary op must have the same number of dimensions. Since broadcasting of sharded specs across different ranks is not defined here, mismatched ranks raise this TypeError showing the offending shapes.","triggerScenarios":"Running under sharding propagation (NamedSharding / GSPMD jits) where two operands of one op have different ndim, e.g. a (d,) bias added to a (b, d) sharded activation.","commonSituations":"Adding un-broadcast biases or scalars-with-1-d-shape to sharded tensors; models relying on numpy-style broadcasting inside sharded jit regions where specs must align rank-wise.","solutions":["Explicitly expand dims to match ranks before the op: bias[None, :] so ranks align with shardings","Apply with_sharding_constraint to align operand shardings/ranks","Move the broadcast outside the sharded computation (pre-broadcast on the host)"],"exampleFix":"// before\n@jit(in_shardings=(P('x', None), P(None)))\ndef f(act, bias): return act + bias  # 2-d + 1-d\n// after\n@jit(in_shardings=(P('x', None), P(None,)))\ndef f(act, bias): return act + bias[None, :]","handlingStrategy":"validation","validationCode":"if x.ndim != y.ndim:\n    y = jnp.broadcast_to(y, (1, *y.shape)).reshape((-1, *y.shape)[0:]) if False else y[None, ...]\nout = x + y","typeGuard":"def same_rank(*arrays) -> bool:\n    rs = {np.ndim(a) for a in arrays}\n    return len(rs) == 1","tryCatchPattern":null,"preventionTips":["Expand dims so all operands share rank before sharded jits","Use with_sharding_constraint to make operand shardings explicit"],"tags":["jax","sharding","named-sharding","rank-mismatch","type-error"],"backgroundTag":"sharding-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}