jax-ml/jax · error · ValueError

Factor {factor} represents a whole dimension; do not specify

Error message

Factor {factor} represents a whole dimension; do not specify its size

What it means

The factor maps to a whole dimension on its own (inferable), so its size is derived from the array shape and must NOT be supplied in factor_sizes. Supplying it is contradictory and rejected.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:156

            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(
            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.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Delete that factor's entry from factor_sizes
2. Restructure the rule so the factor appears only inside compound factors if you truly need to fix its size

Example fix

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

# ensure no factor that spans a whole dimension is in factor_sizes
whole_dims = {f for op in rule.replace('->', ',').split(',') for f in ([op.strip()] if op.strip().isidentifier() else [])}
bad = whole_dims & set(factor_sizes)
assert not bad, f'remove whole-dim factors from factor_sizes: {bad}'

def no_whole_dim_sizes(rule: str, factor_sizes: dict[str, int]) -> bool:
    dims = [d.strip() for d in rule.replace('->', ',').split(',')]
    return not ({d for d in dims if d.isidentifier()} & set(factor_sizes))

When it happens

Trigger: rule='(i,j)->(i,j)' with factor_sizes={'i':8} where 'i' occupies an entire dimension by itself.

Common situations: Over-specifying sizes 'just to be safe'; migrating a rule where a factor used to be compound but now covers a whole dim.

Related errors

• Factor {factor} is not used in the rule, but size is provide
• Factor {factor} is only used in compound factors; must speci
• factor_sizes must be a dict of str to int, but got {factor_s
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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