jax-ml/jax · error · ValueError

Factor {factor} is only used in compound factors; must speci

Error message

Factor {factor} is only used in compound factors; must specify its size

What it means

The factor appears in the rule only inside compound factors (e.g. '(a,b)'), so its size cannot be inferred from a dimension extent and must be given explicitly in factor_sizes. __init__ raises when such a factor is missing from factor_sizes.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:152

        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):
        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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Add the missing factor(s) to factor_sizes with explicit integer sizes, e.g. factor_sizes={'a':2,'b':4}
If the factor should span a whole dimension, move it out of the compound factor so it maps to a dimension directly

Example fix

# before
rule = 'i,(a,b)->i'  # factor_sizes omitted
# after
rule = 'i,(a,b)->i'
factor_sizes = {'a': 2, 'b': 4}

Defensive patterns

Strategy: validation

Validate before calling

import re
whole = re.findall(r'(?:^|[, (])([A-Za-z_][A-Za-z0-9_]*)(?=$|[,)])', rule.replace('->', ','))
compound_only = set(re.findall(r'[A-Za-z_][A-Za-z0-9_]*', rule)) - set(whole)
missing = compound_only - set(factor_sizes)
assert not missing, f'must specify sizes for {missing}'

Type guard

def compound_factors_sized(rule: str, factor_sizes: dict[str, int]) -> bool:
    return not missing if not (missing := _compound_only(rule) - set(factor_sizes)) else False

Prevention

Any factor appearing only inside parentheses needs an explicit size
Keep factor_sizes and rule in one config structure so they change together

When it happens

Trigger: Rule like 'i,(a,b)->i' where 'a' and 'b' only occur inside parentheses, and factor_sizes omits 'a' or 'b'.

Common situations: Assuming all factors are inferable from dimension sizes; adding a compound factor to a rule without updating factor_sizes.

Understand the failure class

Background: Missing required parameter errors: what 'X is required' and 'the required X param is missing' mean, and how to fix them — this error's family across 27 libraries.

Related errors

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