jax-ml/jax · error · ValueError

Ellipsis can only be used at the beginning of a dimension

Error message

Ellipsis can only be used at the beginning of a dimension

What it means

In ArrayMapping, ellipsis/batching notation is only allowed as the first element (mapping the leading batch dimensions). Using it at any later position raises ValueError.

Source

Thrown at jax/_src/custom_partitioning_sharding_rule.py:97

  """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
  the ArrayMappings. If a factor is only used in CompoundFactors, its size must
  be specified.

  By default, a factor is a passthrough factor. Keyword arguments can be used to
  specify other factor kinds including reduction_factors, need_replication_factors,
  and permutation_factors.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Move the ellipsis/BATCHING token to the first position of the ArrayMapping
  2. For trailing dims, enumerate explicit factor names instead of ellipsis
  3. Re-express the rule so variable-rank dims are leading (transpose the problem)

Example fix

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

# after
ArrayMapping('...', 'i')
Defensive patterns

Strategy: validation

Validate before calling

def valid_ellipsis_pos(ms):
    return all(not str(m).startswith('...') for m in ms[1:])
assert valid_ellipsis_pos(dim_mappings)

Prevention

When it happens

Trigger: ArrayMapping('i', '...', 'j') or ArrayMapping('i', '...0'), i.e. ellipsis anywhere but index 0.

Common situations: Attempting to represent trailing variable-rank dims with '...' at the end of a mapping; misunderstanding that batching dims are leading only.

Related errors


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