{"record":{"id":"d4bb3fac72517f45","repo":"jax-ml/jax","slug":"sharding-rule-has-len-rule-operand-mappings-ope","errorCode":null,"errorMessage":"Sharding rule has {len(rule.operand_mappings)} operands, but the operation has {len(operand_types)} operands","messagePattern":"Sharding rule has (.+?) operands, but the operation has (.+?) operands","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_partitioning_sharding_rule.py","lineNumber":374,"sourceCode":"  return SdyShardingRule(operand_mappings, result_mappings,\n                         reduction_factors=reduction_factors,\n                         need_replication_factors=need_replication_factors,\n                         permutation_factors=permutation_factors,\n                         **factor_sizes)\n\n\ndef sdy_sharding_rule_to_mlir(\n  rule: SdyShardingRule,\n  operand_types: list[ir.Type],\n  result_types: list[ir.Type],) -> ir.Attribute:\n  \"\"\"Builds the MLIR representation for the sharding rule.\n\n  This is done by verifying that the rule is consistent with the types of\n  the operation and converting the Einsum notation like string to\n  OpShardingRuleAttr.\n  \"\"\"\n  if len(rule.operand_mappings) != len(operand_types):\n    raise ValueError(\n      f\"Sharding rule has {len(rule.operand_mappings)} operands, but the operation\"\n      f\" has {len(operand_types)} operands\")\n  if len(rule.result_mappings) != len(result_types):\n    raise ValueError(\n      f\"Sharding rule has {len(rule.result_mappings)} results, but the operation\"\n      f\" has {len(result_types)} results\")\n  if not all(isinstance(t, ir.Type) for t in operand_types + result_types):\n    raise TypeError(\n        f\"operand_types and result_types must be a list of ir.Type, but got\"\n        f\" {operand_types} and {result_types}\")\n\n  factors_to_indices_sizes: OrderedDict[str, list[int]] = OrderedDict()\n  types = operand_types + result_types\n  UNKNOWN = -1  # Representation for unknown factor size or factor index.\n\n  def get_message_for_value(i):\n    if i >= len(operand_types):\n      return f\"{i - len(operand_types)}th result\"","sourceCodeStart":356,"sourceCodeEnd":392,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_partitioning_sharding_rule.py#L356-L392","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Count operands in your sharding rule string (left of '->', comma-separated) and make it equal the number of array arguments of the decorated function","Remember scalars/python constants are not operands — exclude them from the rule","Re-run after any signature change to the custom-partitioned function"],"exampleFix":"// before\n@custom_partitioning\ndef dot3(a, b, c):\n  ...\n# rule 'ij,jk->ik' has 2 operands but op has 3\n\n// after\nrule='ij,jk,kl->il'  # 3 operands matching (a, b, c)","handlingStrategy":"validation","validationCode":"# before registering the rule\nn_rule_ops = rule.split('->')[0].count(',') + 1\nassert n_rule_ops == len(array_args), (\n    f'rule has {n_rule_ops} operands, fn takes {len(array_args)} array args')","typeGuard":"def rule_matches_operands(rule: str, args: tuple) -> bool:\n    lhs = rule.split('->')[0]\n    return lhs.count(',') + 1 == sum(not (arg is None or isinstance(arg, (int, float))) for arg in args)","tryCatchPattern":"try:\n    fn_lowered = fn.lower(x, y)\nexcept ValueError as e:\n    if 'operands' in str(e):\n        raise SystemExit(f'sharding rule arity mismatch: {e}') from e\n    raise","preventionTips":["Keep sharding-rule strings next to the function signature and update both together","Write a unit test asserting rule arity equals the function's array-argument count"],"tags":["jax","sharding","custom-partitioning","mlir","sdy"],"backgroundTag":"sharding-rule-operand-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}