jax-ml/jax · error · ValueError

Inconsistent size for {name}: {dim_sizes[name]} vs {shape_va

Error message

Inconsistent size for {name}: {dim_sizes[name]} vs {shape_val}

What it means

einshape binds each named dimension (letter) to a size from the input shape. If the same name appears in multiple positions whose sizes differ (e.g. 'aa' with shape (4, 8)), the binding is contradictory and _get_einshape_dims raises ValueError showing the conflicting sizes. Repeated letters act like einsum labels: they must all have the same extent.

Source

Thrown at jax/_src/pallas/einshape.py:179

  lhs_str, rhs_str = equation.split("->")
  return _parse_side(lhs_str), _parse_side(rhs_str)


def _get_einshape_dims(
    parsed_side: list[list[str]],
    shape: tuple[int, ...],
    sizes: dict[str, int],
) -> dict[str, int]:
  """Parses an einshape equation into a dictionary of dimension sizes."""
  dim_sizes: dict[str, int] = {}

  # Populate known sizes from input
  for i, group in enumerate(parsed_side):
    shape_val = shape[i]
    if len(group) == 1:
      name = group[0]
      if name in dim_sizes and dim_sizes[name] != shape_val:
        raise ValueError(
            f"Inconsistent size for {name}: {dim_sizes[name]} vs {shape_val}"
        )
      dim_sizes[name] = shape_val
    else:
      # We have a merged dimension on LHS, need to split
      known_product = 1
      unknown_dims = []
      for name in group:
        if name in sizes:
          dim_sizes[name] = sizes[name]
          known_product *= sizes[name]
        elif name in dim_sizes:
          known_product *= dim_sizes[name]
        else:
          unknown_dims.append(name)

      if not unknown_dims:
        if known_product != shape_val:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Fix the equation so each repeated letter corresponds to axes of equal size, or use distinct letters for unrelated axes
Verify the input shape matches the equation's LHS structure before calling einshape
If merging/splitting dims, use parenthesized groups with correct products instead of repeated letters

Example fix

# before
t = get_einshape_transforms('aa -> a', x.shape)  # x.shape == (4, 8)

# after
t = get_einshape_transforms('ab -> ab', x.shape)  # distinct names for distinct sizes

Defensive patterns

Strategy: validation

Validate before calling

def dims_consistent(parsed_side, shape):
    sizes = {}
    for group, size in zip(parsed_side, shape):
        if len(group) == 1 and group[0] in sizes:
            if sizes[group[0]] != size:
                return False
        elif len(group) == 1:
            sizes[group[0]] = size
    return True

Type guard

def letter_sizes_consistent(lhs_groups: list[list[str]], shape: tuple) -> bool:
    seen = {}
    for g, s in zip(lhs_groups, shape):
        if len(g) == 1:
            if g[0] in seen and seen[g[0]] != s:
                return False
            seen[g[0]] = s
    return True

Try / catch

try:
    t = get_einshape_transforms(eq, shape)
except ValueError as e:
    if 'Inconsistent size' in str(e):
        raise ValueError(f'equation {eq!r} reuses a letter for axes of different size in shape {shape}') from None
    raise

Prevention

Treat repeated letters as equality constraints; use unique letters for independent axes
Assert input shape matches the equation's implied structure before calling
Regenerate equations when tensor layouts change instead of reusing constants

When it happens

Trigger: Passing an equation whose LHS repeats a dimension letter with inconsistent extents, e.g. get_einshape_transforms('aa -> a', shape=(4, 8)); or a merged group like '(ab)a' where 'a' also appears standalone with a different size.

Common situations: Using repeated letters merely as placeholders instead of einsum-style equality constraints; input arrays whose axis sizes changed after refactoring while the equation stayed fixed; transposing code from reshape where letters carried no meaning.

Related errors

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