{"record":{"id":"de6fb94ea2ec2f13","repo":"jax-ml/jax","slug":"sharding-rule-callable-must-produce-either-an-sdys","errorCode":null,"errorMessage":"sharding_rule callable must produce either an SdyShardingRule object or an Einsum-like notation string.","messagePattern":"sharding_rule callable must produce either an SdyShardingRule object or an Einsum-like notation string\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_partitioning.py","lineNumber":651,"sourceCode":"      has_side_effect=ir.BoolAttr.get(False),\n      api_version=mlir.i32_attr(2),\n      called_computations=ir.ArrayAttr.get([]),\n      backend_config=ir.StringAttr.get(key),\n      operand_layouts=None,\n      result_layouts=None)\n  if sharding_rule is not None:\n    value_types, _ = mlir.ir_tree_registry.flatten(\n        [mlir.aval_to_ir_types(ctx.module_context, a) for a in call.in_avals])\n    if callable(sharding_rule):\n      sharding_rule = sharding_rule(*static_args, mesh, value_types, result_types)\n      if isinstance(sharding_rule, (list, tuple)) and len(sharding_rule) == 2:\n        sharding_rule, sharding_rule_dict = sharding_rule\n      else:\n        sharding_rule_dict = {}\n      if isinstance(sharding_rule, str):\n        sharding_rule = str_to_sdy_sharding_rule(sharding_rule, **sharding_rule_dict)\n      elif not isinstance(sharding_rule, SdyShardingRule):\n          raise ValueError(\"sharding_rule callable must produce either an \"\n                           \"SdyShardingRule object or an Einsum-like notation \"\n                           \"string.\")\n    out.attributes['sdy.sharding_rule'] = sdy_sharding_rule_to_mlir(\n      sharding_rule, value_types, result_types)\n  return out.results\n\nmlir.register_lowering(custom_partitioning_p,\n                       _custom_partitioning_lowering_rule)\n\nxc.register_custom_call_partitioner(\n    _CUSTOM_PARTITIONING_CALL_NAME,\n    _custom_partitioning_propagate_user_sharding,\n    _custom_partitioning_partition,\n    _custom_partitioning_infer_sharding_from_operands,\n    can_side_effecting_have_replicated_sharding=True,\n)\nxb.register_plugin_callbacks(\n    partial(","sourceCodeStart":633,"sourceCodeEnd":669,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_partitioning.py#L633-L669","documentation":"A custom_partitioning sharding_rule callback may return either an Einsum-like notation string or an SdyShardingRule object. Returning any other type (dict handled separately; other objects invalid) raises ValueError at lowering time.","triggerScenarios":"A sharding_rule callable returning None, a tuple that is not (str, dict), a Sharding object, or other custom class when the custom-partitioned function is compiled.","commonSituations":"Callbacks that conditionally return None for unsupported shapes, or that return SDY protos/MLIR attributes instead of SdyShardingRule.","solutions":["Return a string like 'i j, j k -> i k' or an SdyShardingRule instance from the callback","Optionally return a (rule_string, kwargs_dict) pair for factor options","Never return None; cover all shape cases with a valid rule"],"exampleFix":"# before\ndef rule(mesh, shapes):\n  if unsupported: return None\n\n# after\ndef rule(mesh, shapes):\n  return 'i j -> i j'  # always a valid rule string","handlingStrategy":"type-guard","validationCode":"r = rule_cb(mesh, shapes)\nassert isinstance(r, (str, SdyShardingRule)) or (isinstance(r, tuple) and isinstance(r[0], str) and isinstance(r[1], dict)), type(r)","typeGuard":"def valid_rule(r):\n    return isinstance(r, (str, SdyShardingRule)) or (isinstance(r, tuple) and len(r) == 2 and isinstance(r[0], str) and isinstance(r[1], dict))","tryCatchPattern":null,"preventionTips":["Never return None from sharding_rule callbacks; cover every shape case"],"tags":["jax","shardy","sharding-rule","custom-partitioning"],"backgroundTag":"invalid-sharding-rule-return-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}