jax-ml/jax · error · ValueError

Factor {factor} can only be in one of the reduction, need re

Error message

Factor {factor} can only be in one of the reduction, need replication, or permutation factor sets.

What it means

A factor may belong to only one of the three special sets (reduction, need-replication, permutation). special_factors accumulates seen factors across all three checks and rejects a repeat.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:173

      if factor in factor_sizes and inferable:
        raise ValueError(
          f"Factor {factor} represents a whole dimension; do not specify its"
          " size")

    special_factors = set()
    def check_special_factors(kind, factors):
      if not isinstance(factors, tuple):
        raise ValueError(f"{kind} must be a tuple of factors")

      if len(factors) != len(set(factors)):
        raise ValueError(f"{kind} contains duplicated factors")

      for factor in factors:
        if factor not in factors_inferrable:
          raise ValueError(
            f"Factor {factor} in {kind} is not used in the rule")
        if factor in special_factors:
          raise ValueError(f"Factor {factor} can only be in one of the "
              f"reduction, need replication, or permutation factor sets.")
        special_factors.add(factor)

    check_special_factors("reduction_factors", reduction_factors)
    check_special_factors("need_replication_factors", need_replication_factors)
    check_special_factors("permutation_factors", permutation_factors)

    self.operand_mappings = operand_mappings
    self.result_mappings = result_mappings
    self.factor_sizes = factor_sizes
    self.reduction_factors = reduction_factors
    self.need_replication_factors = need_replication_factors
    self.permutation_factors = permutation_factors


  def __str__(self):
    def to_str(kind, factors):
      if len(factors) > 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Keep the factor in exactly one of the three tuples and delete it from the others
If different behavior is needed for the same name, rename one occurrence in the rule to a distinct factor

Example fix

# before
reduction_factors=('m',), need_replication_factors=('m',)
# after
reduction_factors=('m',), need_replication_factors=('r',)

Defensive patterns

Strategy: validation

Validate before calling

sets = [reduction_factors, need_replication_factors, permutation_factors]
seen = set()
for s in sets:
    assert not (seen & set(s)), f'factor in multiple sets: {seen & set(s)}'
    seen |= set(s)

Type guard

def disjoint_factor_sets(*sets) -> bool:
    seen = set()
    return not any(seen & set(s) or (seen := seen | set(s)) for s in sets)

Prevention

When moving a factor between categories, delete it from the old tuple in the same commit

When it happens

Trigger: reduction_factors=('m',) and need_replication_factors=('m',) in the same call.

Common situations: Iterating on sharding rules and moving a factor between categories without removing it from the old one.

Related errors

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