jax-ml/jax · error · ValueError

Ellipsis can't be used in a compound factor

Error message

Ellipsis can't be used in a compound factor

What it means

Ellipsis ('...'/batching notation) cannot appear inside a CompoundFactor, because compound factors map to concrete factor groups on one dimension. Using a batching factor there raises ValueError.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:70

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')
  """
  def __init__(self, *dim_mappings):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Place ellipsis/batching at the start of the ArrayMapping: ArrayMapping('...', CompoundFactor('b','c'))
2. Remove ellipsis from CompoundFactor elements
3. Use the BATCHING constant at position 0 of the ArrayMapping

Example fix

# before
ArrayMapping(CompoundFactor('...', 'b'), 'i')

# after
ArrayMapping(BATCHING, CompoundFactor('b', 'c'), 'i')

def no_ellipsis(factors):
    return all(not f.startswith('...') for f in factors)
assert no_ellipsis(factors)

When it happens

Trigger: CompoundFactor('...', 'b') or CompoundFactor('...0', 'c') while writing an SdyShardingRule.

Common situations: Trying to express 'leading dims plus factors' inside one compound token instead of at the mapping's leading position.

Related errors

• Ellipsis can only be used at the beginning of a dimension
• Unknown keyword arguments: {sharding_rule_dict}
• sharding_rule callable must produce either an SdyShardingRul
• Factor names have to start with a letter, but got '{factor[0
• Unknown character '{char}'

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