jax-ml/jax · error · ValueError

Equation must contain exactly one '->'

Error message

Equation must contain exactly one '->'

What it means

einshape requires a single explicit '->' mapping an input-side layout to an output-side layout (unlike general einsum, which allows implicit mode). _parse_equation counts occurrences of '->' and raises ValueError if there isn't exactly one — both zero occurrences and two or more fail.

Source

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

      j = s.find(")", i)
      if j == -1:
        raise ValueError(f"Unmatched parenthesis in {s!r}")
      group = list(s[i + 1 : j])
      groups.append(group)
      i = j + 1
    elif s[i] == ")":
      raise ValueError(f"Unmatched parenthesis in {s!r}")
    else:
      # distinct dimension
      groups.append([s[i]])
      i += 1
  return groups


def _parse_equation(equation: str) -> tuple[list[list[str]], list[list[str]]]:
  """Parses an einshape equation."""
  if equation.count("->") != 1:
    raise ValueError("Equation must contain exactly one '->'")
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add exactly one '->' specifying the output layout, e.g. 'abc -> a(bc)'
2. If you intended implicit mode, compute the output side explicitly instead — einshape has no implicit form
3. Sanity-check equation.count('->') == 1 before calling einshape helpers

Example fix

# before
t = get_einshape_transforms('abc', x.shape)  # no '->'

# after
t = get_einshape_transforms('abc -> a(bc)', x.shape)

assert equation.count('->') == 1, 'einshape equation needs exactly one "->"'

def is_valid_einshape_equation(eq: str) -> bool:
    return eq.count('->') == 1 and all(eq_side_ok(s) for s in eq.split('->'))

try:
    t = get_einshape_transforms(eq, shape)
except ValueError as e:
    if 'exactly one' in str(e):
        eq = eq if '->' in eq else eq + ' -> ' + ''.join(c for c in eq if c not in '()->')
        t = get_einshape_transforms(eq, shape)
    else:
        raise

When it happens

Trigger: Calling get_einshape_transforms with an equation like 'abc' (no arrow) or 'a->b->c' (two arrows). e.g. jax._src.pallas.einshape.get_einshape_transforms('abc', shape).

Common situations: Assuming einsum-style implicit output ordering works in einshape; string concatenation accidentally duplicating the arrow; typos like '-' or '=>'.

Related errors

• Unmatched parenthesis in {s!r}
• Inconsistent size for {name}: {dim_sizes[name]} vs {shape_va
• Size mismatch for group {group}: expected {shape_val}, got {
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64

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