jax-ml/jax · error · ValueError

sharding_rule callable must produce either an SdyShardingRul

Error message

sharding_rule callable must produce either an SdyShardingRule object or an Einsum-like notation string.

What it means

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.

Source

Thrown at jax/_src/custom_partitioning.py:651

      has_side_effect=ir.BoolAttr.get(False),
      api_version=mlir.i32_attr(2),
      called_computations=ir.ArrayAttr.get([]),
      backend_config=ir.StringAttr.get(key),
      operand_layouts=None,
      result_layouts=None)
  if sharding_rule is not None:
    value_types, _ = mlir.ir_tree_registry.flatten(
        [mlir.aval_to_ir_types(ctx.module_context, a) for a in call.in_avals])
    if callable(sharding_rule):
      sharding_rule = sharding_rule(*static_args, mesh, value_types, result_types)
      if isinstance(sharding_rule, (list, tuple)) and len(sharding_rule) == 2:
        sharding_rule, sharding_rule_dict = sharding_rule
      else:
        sharding_rule_dict = {}
      if isinstance(sharding_rule, str):
        sharding_rule = str_to_sdy_sharding_rule(sharding_rule, **sharding_rule_dict)
      elif not isinstance(sharding_rule, SdyShardingRule):
          raise ValueError("sharding_rule callable must produce either an "
                           "SdyShardingRule object or an Einsum-like notation "
                           "string.")
    out.attributes['sdy.sharding_rule'] = sdy_sharding_rule_to_mlir(
      sharding_rule, value_types, result_types)
  return out.results

mlir.register_lowering(custom_partitioning_p,
                       _custom_partitioning_lowering_rule)

xc.register_custom_call_partitioner(
    _CUSTOM_PARTITIONING_CALL_NAME,
    _custom_partitioning_propagate_user_sharding,
    _custom_partitioning_partition,
    _custom_partitioning_infer_sharding_from_operands,
    can_side_effecting_have_replicated_sharding=True,
)
xb.register_plugin_callbacks(
    partial(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return a string like 'i j, j k -> i k' or an SdyShardingRule instance from the callback
2. Optionally return a (rule_string, kwargs_dict) pair for factor options
3. Never return None; cover all shape cases with a valid rule

Example fix

# before
def rule(mesh, shapes):
  if unsupported: return None

# after
def rule(mesh, shapes):
  return 'i j -> i j'  # always a valid rule string

r = rule_cb(mesh, shapes)
assert isinstance(r, (str, SdyShardingRule)) or (isinstance(r, tuple) and isinstance(r[0], str) and isinstance(r[1], dict)), type(r)

def valid_rule(r):
    return isinstance(r, (str, SdyShardingRule)) or (isinstance(r, tuple) and len(r) == 2 and isinstance(r[0], str) and isinstance(r[1], dict))

When it happens

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

Common situations: Callbacks that conditionally return None for unsupported shapes, or that return SDY protos/MLIR attributes instead of SdyShardingRule.

Related errors

• Unknown keyword arguments: {sharding_rule_dict}
• Shardy is used, but sharding propagation callbacks instead o
• Custom-partitioned function {function!r} does not support GS
• Factor names have to start with a letter, but got '{factor[0
• Unknown character '{char}'

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