jax-ml/jax · error · ValueError

Unknown keyword arguments: {sharding_rule_dict}

Error message

Unknown keyword arguments: {sharding_rule_dict}

What it means

In custom_partitioning's def_partition, when constructing an SdyShardingRule from user-supplied pieces, any leftover unsupported keyword entries in sharding_rule_dict raise ValueError. This happens when a rule definition contains options the installed SDY path does not accept.

Source

Thrown at jax/_src/custom_partitioning.py:497

                    propagate_user_sharding=None, decode_shardings=True,
                    sharding_rule=None, *, reduction_factors=(),
                    need_replication_factors=(), permutation_factors=(),
                    **factor_sizes):
    self.partition = partition
    self.propagate_user_sharding = propagate_user_sharding
    self.infer_sharding_from_operands = infer_sharding_from_operands
    self.decode_shardings = decode_shardings
    if (sharding_rule is None or isinstance(sharding_rule, Callable) or
        isinstance(sharding_rule, SdyShardingRule)):
      sharding_rule_dict = factor_sizes
      if len(reduction_factors) > 0:
        sharding_rule_dict["reduction_factors"] = reduction_factors
      if len(need_replication_factors) > 0:
        sharding_rule_dict["need_replication_factors"] = need_replication_factors
      if len(permutation_factors) > 0:
        sharding_rule_dict["permutation_factors"] = permutation_factors
      if sharding_rule_dict:
        raise ValueError(f"Unknown keyword arguments: {sharding_rule_dict}")
      self.sharding_rule = sharding_rule
    else:
      self.sharding_rule = str_to_sdy_sharding_rule(
          sharding_rule,
          reduction_factors=reduction_factors,
          need_replication_factors=need_replication_factors,
          permutation_factors=permutation_factors,
          **factor_sizes)
    return partition

  def __call__(self, *args, **kwargs):
    args = _resolve_kwargs(self.fun, args, kwargs)
    debug = api_util.debug_info("custom_partitioning", self.fun,
                                args, {},
                                static_argnums=self.static_argnums)
    if self.static_argnums:
      static_argnums = set(self.static_argnums)
      dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove reduction_factors / need_replication_factors / permutation_factors and express them in the sharding_rule itself
2. Pass those arguments only via str_to_sdy_sharding_rule when using a string rule, not alongside a prebuilt SdyShardingRule
3. Upgrade JAX so the rule options you need are supported

Example fix

# before
rule = def_partition(infer_sharding_from_operands=..., reduction_factors=[...])

# after
rule = def_partition(sharding_rule='i j, j k -> i k')  # factors encoded in the rule string

allowed = {'reduction_factors','need_replication_factors','permutation_factors'}
# only pass these when using a rule string, never with a prebuilt SdyShardingRule

When it happens

Trigger: Building a sharding rule via def_partition with kwargs such as reduction_factors, need_replication_factors, or permutation_factors that end up unconsumed (e.g. combined with an explicit SdyShardingRule where they are not applicable).

Common situations: Mixing GSPMD-style factor arguments with the new sharding_rule API; version skew between JAX and the Shardy compiler.

Related errors

• sharding_rule callable must produce either an SdyShardingRul
• 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/4a635c4465e4bc95. Report an issue: GitHub.