jax-ml/jax · error · ValueError

Sharding rule has {len(rule.operand_mappings)} operands, but

Error message

Sharding rule has {len(rule.operand_mappings)} operands, but the operation has {len(operand_types)} operands

What it means

Raised by sdy_sharding_rule_to_mlir when a custom sharding rule (Einsum-like notation) declares a different number of operands than the operation being lowered actually has. JAX validates user-supplied sharding rules against the op's operand/result types before converting them to the SDY OpShardingRuleAttr in MLIR.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:374

  return SdyShardingRule(operand_mappings, result_mappings,
                         reduction_factors=reduction_factors,
                         need_replication_factors=need_replication_factors,
                         permutation_factors=permutation_factors,
                         **factor_sizes)


def sdy_sharding_rule_to_mlir(
  rule: SdyShardingRule,
  operand_types: list[ir.Type],
  result_types: list[ir.Type],) -> ir.Attribute:
  """Builds the MLIR representation for the sharding rule.

  This is done by verifying that the rule is consistent with the types of
  the operation and converting the Einsum notation like string to
  OpShardingRuleAttr.
  """
  if len(rule.operand_mappings) != len(operand_types):
    raise ValueError(
      f"Sharding rule has {len(rule.operand_mappings)} operands, but the operation"
      f" has {len(operand_types)} operands")
  if len(rule.result_mappings) != len(result_types):
    raise ValueError(
      f"Sharding rule has {len(rule.result_mappings)} results, but the operation"
      f" has {len(result_types)} results")
  if not all(isinstance(t, ir.Type) for t in operand_types + result_types):
    raise TypeError(
        f"operand_types and result_types must be a list of ir.Type, but got"
        f" {operand_types} and {result_types}")

  factors_to_indices_sizes: OrderedDict[str, list[int]] = OrderedDict()
  types = operand_types + result_types
  UNKNOWN = -1  # Representation for unknown factor size or factor index.

  def get_message_for_value(i):
    if i >= len(operand_types):
      return f"{i - len(operand_types)}th result"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Count operands in your sharding rule string (left of '->', comma-separated) and make it equal the number of array arguments of the decorated function
2. Remember scalars/python constants are not operands — exclude them from the rule
3. Re-run after any signature change to the custom-partitioned function

Example fix

// before
@custom_partitioning
def dot3(a, b, c):
  ...
# rule 'ij,jk->ik' has 2 operands but op has 3

// after
rule='ij,jk,kl->il'  # 3 operands matching (a, b, c)

# before registering the rule
n_rule_ops = rule.split('->')[0].count(',') + 1
assert n_rule_ops == len(array_args), (
    f'rule has {n_rule_ops} operands, fn takes {len(array_args)} array args')

def rule_matches_operands(rule: str, args: tuple) -> bool:
    lhs = rule.split('->')[0]
    return lhs.count(',') + 1 == sum(not (arg is None or isinstance(arg, (int, float))) for arg in args)

try:
    fn_lowered = fn.lower(x, y)
except ValueError as e:
    if 'operands' in str(e):
        raise SystemExit(f'sharding rule arity mismatch: {e}') from e
    raise

When it happens

Trigger: Calling jax.custom_partitioning with a sharding_rule string like 'ij,jk->ik' on a function whose actual operand count differs (e.g. rule lists 2 operands but the function takes 3, or a python number/scalar constant changes the operand count).

Common situations: Adding/removing an argument to a custom-partitioned function without updating the sharding-rule string; passing scalars that JAX does not count as array operands; copied rule strings from another op with different arity.

Related errors

• Custom Partitioning rules must return Sharding.
• Mismatch in result shapes. %s vs %s
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• to_dlpack can only pack a dlpack tensor from an array on a s

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