jax-ml/jax · error · ValueError

Factor {factor} is not used in the rule, but size is provide

Error message

Factor {factor} is not used in the rule, but size is provided

What it means

Thrown by SdyShardingRule.__init__ when a factor name present in factor_sizes does not appear anywhere in the sharding rule string (operands or results). The rule parser builds the set of factors actually referenced; any extra key in factor_sizes is rejected because its size can never be used.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:144

               *, reduction_factors: tuple[str, ...] = (),
               need_replication_factors: tuple[str, ...] = (),
               permutation_factors: tuple[str, ...] = (),
               **factor_sizes: int):
    # Find all factors and mark whether their size can be inferred.
    factors_inferrable = {}
    for value in operand_mappings + result_mappings:
      for dim in value:
        if isinstance(dim, str):
          factors_inferrable[dim] = True
        else:
          for factor in dim:
            if factor not in factors_inferrable.keys():
              factors_inferrable[factor] = False

    # Check that factors in factor_sizes are used in the rule.
    for factor in factor_sizes:
      if factor not in factors_inferrable:
        raise ValueError(
          f"Factor {factor} is not used in the rule, but size is provided")

    # Check that factors that are used for a whole dimension aren't in
    # factor_sizes and factors that are never used for a whole dimension are
    # in factor_sizes.
    for factor, inferable in factors_inferrable.items():
      if factor not in factor_sizes and not inferable:
        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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Remove the unused factor key from factor_sizes (or fix its spelling so it matches the rule)
  2. Check every key of factor_sizes against the factor tokens (including those inside compound factors like '(a,b)') in the rule string
  3. Regenerate factor_sizes from the rule programmatically to avoid drift

Example fix

# before
def_partition(fn, rule='(i,j)->(i,j)', factor_sizes={'i':8,'k':4})
# after
def_partition(fn, rule='(i,j)->(i,j)', factor_sizes={'i':8})
Defensive patterns

Strategy: validation

Validate before calling

rule_factors = set(re.findall(r'[A-Za-z_][A-Za-z0-9_]*', rule.split('->')[0] + rule.split('->')[1]))
assert set(factor_sizes) <= rule_factors, f'unused factor_sizes keys: {set(factor_sizes) - rule_factors}'

Type guard

def has_only_used_factors(rule: str, factor_sizes: dict[str, int]) -> bool:
    import re
    used = set(re.findall(r'[A-Za-z_][A-Za-z0-9_]*', rule))
    return set(factor_sizes) <= used

Prevention

When it happens

Trigger: Calling jax.experimental.custom_partitioning.def_partition (or str_to_sdy_sharding_rule) with factor_sizes={'a':2,'b':4} while the rule string only mentions 'a', e.g. rule='(a,a)->(a)' with a stray 'b' key.

Common situations: Renaming factors in the rule but forgetting to update factor_sizes; copy-pasting a rule from another op and leaving stale factor_sizes keys; typos in factor names inside the rule string.

Related errors

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