jax-ml/jax · error · ValueError

Each element of ArrayMapping must be a str or CompoundFactor

Error message

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

What it means

Each element of ArrayMapping (a per-dimension factor spec in sharding rules) must be a str or CompoundFactor. Any other type raises ValueError with the offending type name.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:91

  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):
    for i, d in enumerate(dim_mappings):
      if not isinstance(d, str) and not isinstance(d, CompoundFactor):
        raise ValueError(
            "Each element of ArrayMapping must be a str or CompoundFactor, but"
            f" got {type(d)}")
      if isinstance(d, str):
        if _is_batching(d):
          if i != 0:
            raise ValueError("Ellipsis can only be used at the beginning of a dimension")
        else:
          _check_factor(d)

  def __new__(cls, *dim_mappings):
    return tuple.__new__(ArrayMapping, dim_mappings)


class SdyShardingRule:
  """Represents a Shardy sharding rule.

  An SdyShardingRule contains the ArrayMappings for operands and results,
  optional special factors and optional factor sizes. A factor is a name used in

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use plain factor strings and CompoundFactor objects only
2. Convert tuples ('a','b') to CompoundFactor('a','b')
3. Substitute a valid factor (or omit the mapping element) instead of None

Example fix

# before
ArrayMapping(None, 'i')

# after
ArrayMapping(BATCHING, 'i')  # or ArrayMapping('a', 'i')

from jax._src.custom_partitioning_sharding_rule import CompoundFactor
def valid_mapping(ms):
    return all(isinstance(m, (str, CompoundFactor)) for m in ms)

def valid_mapping_element(m):
    return isinstance(m, (str, CompoundFactor))

When it happens

Trigger: ArrayMapping(None, 'i'), ArrayMapping(('a','b'), 'i'), or passing an SdyShardingRule/axis object where a factor string belongs.

Common situations: Programmatically building mappings where optional factors default to None, or passing tuples instead of CompoundFactor.

Related errors

• Each element of CompoundFactor must be a str, but got {type(
• 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/17af119bf7dff4d4. Report an issue: GitHub.