jax-ml/jax · error · ValueError

Each element of CompoundFactor must be a str, but got {type(

Error message

Each element of CompoundFactor must be a str, but got {type(factor)}

What it means

Every element of a CompoundFactor must be a plain Python string. Passing any non-str (int, list, another CompoundFactor, etc.) raises ValueError with the offending type.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:68

    return False
  return len(factor) == 1 or factor[1:].isdigit()

def _get_batching_group(factor: str) -> str:
  """Extracts the batching group from a factor for leading batching dimensions."""
  return factor[1:] if len(factor) > 1 else "0"

class CompoundFactor(tuple):
  """Describes the factors for a compound factor.

  A compound factor should contain at least two factors, e.g.
  * CompoundFactor('b', 'c').
  """
  def __init__(self, *factors):
    if len(factors) < 2:
      raise ValueError("A compound factor should contain at least two factors")
    for factor in factors:
      if not isinstance(factor, str):
        raise ValueError(f"Each element of CompoundFactor must be a str, but got {type(factor)}")
      if _is_batching(factor):
        raise ValueError("Ellipsis can't be used in a compound factor")
      else:
        _check_factor(factor)

  def __new__(cls, *factors):
    return tuple.__new__(CompoundFactor, factors)


class ArrayMapping(tuple):
  """Describes the factors for an operand or result.

  Each element is either a factor or a CompoundFactor. A leading element can
  also be BATCHING, which represents batching dimensions. examples:
  * ArrayMapping('a')
  * ArrayMapping('b', 'c')
  * ArrayMapping(CompoundFactor('b', 'c'), 'd')
  * ArrayMapping(BATCHING, CompoundFactor('b', 'c'), 'd')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert elements to strings before constructing: CompoundFactor(str(f1), str(f2))
  2. Flatten nested CompoundFactors into their string components
  3. Add a type check in rule-generation code

Example fix

# before
CompoundFactor(mesh_axis_id, 'c')  # int

# after
CompoundFactor(str(mesh_axis_id), 'c')
Defensive patterns

Strategy: type-guard

Validate before calling

assert all(isinstance(f, str) for f in factors), [type(f) for f in factors]

Type guard

def all_str(factors):
    return all(isinstance(f, str) for f in factors)

Prevention

When it happens

Trigger: CompoundFactor(2, 'c'), CompoundFactor(['b'], 'c'), or nested CompoundFactor usage while building sharding rules.

Common situations: Building rules programmatically where factor lists may contain ints or nested tuples.

Related errors


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