jax-ml/jax · error · TypeError

{}: arrays must have same number of dimensions, got {}.

Error message

{}: arrays must have same number of dimensions, got {}.

What it means

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.

Source

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

    raise RuntimeError(
      "First argument of broadcasting_sharding_rule should be a name."
      f" Got {name}")
  mesh = None
  for a in avals:
    if a.sharding is not None and not a.sharding.mesh.empty:
      if mesh is not None and mesh != a.sharding.mesh:
        raise core.ShardingTypeError(
            f'Mesh for all inputs should be equal. Got one mesh: {mesh} and'
            f' another mesh: {a.sharding.mesh}')
      mesh = a.sharding.mesh
  mesh = get_abstract_mesh() if mesh is None else mesh

  shapes = [aval.shape for aval in avals if aval.shape]
  if not shapes:
    return NamedSharding(mesh, P())
  if len({len(shape) for shape in shapes}) != 1:
    msg = '{}: arrays must have same number of dimensions, got {}.'
    raise TypeError(msg.format(name, ', '.join(map(str, map(tuple, shapes)))))

  specs = [a.sharding.spec.partitions for a in avals if a.shape]

  result_specs = [None] * len(shapes[0])
  for i, (ss, ds) in enumerate(zip(zip(*specs), zip(*shapes))):
    if all(ss[0] == s for s in ss[1:]):
      # if all dimension shardings are same, the resulting dimension sharding is
      # the same.
      result_specs[i] = ss[0]
    else:
      non_trivial_s = [s for s, d in zip(ss, ds)
                       if not (core.definitely_equal(d, 1) and s is None)]
      if not non_trivial_s:
        result_specs[i] = None
      elif all(non_trivial_s[0] == s for s in non_trivial_s[1:]):
        result_specs[i] = non_trivial_s[0]
      else:
        for s in ss:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Explicitly expand dims to match ranks before the op: bias[None, :] so ranks align with shardings
2. Apply with_sharding_constraint to align operand shardings/ranks
3. Move the broadcast outside the sharded computation (pre-broadcast on the host)

Example fix

// before
@jit(in_shardings=(P('x', None), P(None)))
def f(act, bias): return act + bias  # 2-d + 1-d
// after
@jit(in_shardings=(P('x', None), P(None,)))
def f(act, bias): return act + bias[None, :]

if x.ndim != y.ndim:
    y = jnp.broadcast_to(y, (1, *y.shape)).reshape((-1, *y.shape)[0:]) if False else y[None, ...]
out = x + y

def same_rank(*arrays) -> bool:
    rs = {np.ndim(a) for a in arrays}
    return len(rs) == 1

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Sharding spec {spec} implies that array axis {dim} is partit
• Expected unreduced_kind to be of type `jax.sharding.Unreduce
• {name} argument of ManualAxisType should of type `frozenset`
• shardings specs rank should be 3, but got lhs: {len(lhs.spec
• {kind} must be a tuple of factors

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