jax-ml/jax · error · ValueError

Factor {factor} in {kind} is not used in the rule

Error message

Factor {factor} in {kind} is not used in the rule

What it means

Every factor listed in reduction_factors / need_replication_factors / permutation_factors must actually appear in the rule string. This check fires when a special factor is nowhere in operands or results.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:170

        raise ValueError(
          f"Factor {factor} is only used in compound factors; must specify"
          " its size")
      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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the factor from the special-factor tuple or fix its name to match the rule
2. Cross-check each special factor against tokens in the rule string

Example fix

# before
rule='(i,j)->(i,j)', permutation_factors=('m',)
# after
rule='(i,j,m)->(i,j)', permutation_factors=('m',)

import re
used = set(re.findall(r'[A-Za-z_][A-Za-z0-9_]*', rule))
assert set(reduction_factors) <= used and set(need_replication_factors) <= used and set(permutation_factors) <= used

def special_factors_in_rule(rule: str, *sets) -> bool:
    used = set(re.findall(r'[A-Za-z_][A-Za-z0-9_]*', rule))
    return all(set(s) <= used for s in sets)

When it happens

Trigger: rule='(i,j)->(i,j)' with permutation_factors=('m',) — 'm' never occurs in the rule.

Common situations: Renaming rule factors without updating the special-factor tuples; typos in factor names.

Related errors

• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin
• `unreduced_kind` should be `None` when `unreduced` is an emp
• Factor {factor} is not used in the rule, but size is provide

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