jax-ml/jax · error · ValueError

Run out of characters for batch dimension in einsum.

Error message

Run out of characters for batch dimension in einsum.

What it means

Raised by Einsum.batch when vmapping an einsum whose subscripts already consume all 52 ascii letters (upper+lower case). Batching requires injecting one fresh batch letter into the subscripts; if none remain unused, the primitive cannot construct a batched rule.

Source

Thrown at jax/_src/numpy/hijax.py:400

        list(operands_out), contractions, precision=self.precision,
        preferred_element_type=self.preferred_element_type)

  def batch(
      self,
      axis_data: Any,
      args: tuple[Array, ...],
      dims: tuple[int | None, ...]
  ) -> tuple[Array, int | None]:
    del axis_data  # unused
    if all(d is None for d in dims):
      return self(*args), None
    input_subs, output_sub = self.subscripts.split('->')
    input_subs_list = input_subs.split(',')

    used_chars = set(input_subs) | set(output_sub) | {',', '-', '>'}
    available_chars = [c for c in string.ascii_letters if c not in used_chars]
    if not available_chars:
      raise ValueError("Run out of characters for batch dimension in einsum.")
    batch_char = available_chars[0]

    new_input_subs_list = []
    for i, dim in enumerate(dims):
      sub = input_subs_list[i]
      if dim is not None:
        sub_list = list(sub)
        sub_list.insert(dim, batch_char)
        new_sub = "".join(sub_list)
      else:
        new_sub = sub
      new_input_subs_list.append(new_sub)

    new_input_subs = ",".join(new_input_subs_list)
    new_output_sub = batch_char + output_sub
    new_subscripts = f"{new_input_subs}->{new_output_sub}"

    batched_prim = Einsum(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reduce the number of distinct subscript letters by splitting the contraction into multiple einsum calls with intermediates
Batch manually instead of via vmap: move the batch axis into an existing subscript letter that already appears in every operand
Use ellipsis (...) for broadcasting dims to free explicit letters

Example fix

# before
out = jax.vmap(lambda *args: jnp.einsum(all_letters_subscripts, *args))(operands)
# after
# use ellipsis to free letters: '...abc,...bcd->...acd'
out = jax.vmap(lambda *args: jnp.einsum('...abc,...bcd->...acd', *args))(operands)

Defensive patterns

Strategy: fallback

Validate before calling

import string
def einsum_has_free_letter(subscripts: str) -> bool:
    used = set(subscripts.replace('->', '').replace(',', ''))
    return bool(set(string.ascii_letters) - used)
assert einsum_has_free_letter(subscripts), 'no letter left for vmap batching'

Type guard

def vmappable_subscripts(s: str) -> bool:
    used = set(s) - {',', '-', '>'}
    return len(set(string.ascii_letters) - used) > 0

Try / catch

try:
    out = jax.vmap(f)(xs)
except ValueError as e:
    if 'Run out of characters' in str(e):
        # manual batching: fold batch axis into an existing subscript letter
        out = jnp.stack([f(x) for x in xs])
    else:
        raise

Prevention

Use ellipsis (...) for batch dims in generated einsums
Cap distinct letters in generated subscripts; split giant contractions into stages
Add a generator test asserting at least one free ascii letter remains

When it happens

Trigger: Applying jax.vmap (directly or nested) to an einsum with an extremely wide contraction using all a-zA-Z letters in its subscripts, e.g. tens of operands contracting dozens of distinct axes.

Common situations: Programmatically generated einsums over high-dimensional tensor networks (physics, attention batching) that accumulate letters until the alphabet is exhausted.

Related errors

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