{"record":{"id":"cb00f67c9ba13abf","repo":"jax-ml/jax","slug":"run-out-of-characters-for-batch-dimension-in-einsu","errorCode":null,"errorMessage":"Run out of characters for batch dimension in einsum.","messagePattern":"Run out of characters for batch dimension in einsum\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/hijax.py","lineNumber":400,"sourceCode":"        list(operands_out), contractions, precision=self.precision,\n        preferred_element_type=self.preferred_element_type)\n\n  def batch(\n      self,\n      axis_data: Any,\n      args: tuple[Array, ...],\n      dims: tuple[int | None, ...]\n  ) -> tuple[Array, int | None]:\n    del axis_data  # unused\n    if all(d is None for d in dims):\n      return self(*args), None\n    input_subs, output_sub = self.subscripts.split('->')\n    input_subs_list = input_subs.split(',')\n\n    used_chars = set(input_subs) | set(output_sub) | {',', '-', '>'}\n    available_chars = [c for c in string.ascii_letters if c not in used_chars]\n    if not available_chars:\n      raise ValueError(\"Run out of characters for batch dimension in einsum.\")\n    batch_char = available_chars[0]\n\n    new_input_subs_list = []\n    for i, dim in enumerate(dims):\n      sub = input_subs_list[i]\n      if dim is not None:\n        sub_list = list(sub)\n        sub_list.insert(dim, batch_char)\n        new_sub = \"\".join(sub_list)\n      else:\n        new_sub = sub\n      new_input_subs_list.append(new_sub)\n\n    new_input_subs = \",\".join(new_input_subs_list)\n    new_output_sub = batch_char + output_sub\n    new_subscripts = f\"{new_input_subs}->{new_output_sub}\"\n\n    batched_prim = Einsum(","sourceCodeStart":382,"sourceCodeEnd":418,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/hijax.py#L382-L418","documentation":"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.","triggerScenarios":"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.","commonSituations":"Programmatically generated einsums over high-dimensional tensor networks (physics, attention batching) that accumulate letters until the alphabet is exhausted.","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"],"exampleFix":"# before\nout = jax.vmap(lambda *args: jnp.einsum(all_letters_subscripts, *args))(operands)\n# after\n# use ellipsis to free letters: '...abc,...bcd->...acd'\nout = jax.vmap(lambda *args: jnp.einsum('...abc,...bcd->...acd', *args))(operands)","handlingStrategy":"fallback","validationCode":"import string\ndef einsum_has_free_letter(subscripts: str) -> bool:\n    used = set(subscripts.replace('->', '').replace(',', ''))\n    return bool(set(string.ascii_letters) - used)\nassert einsum_has_free_letter(subscripts), 'no letter left for vmap batching'","typeGuard":"def vmappable_subscripts(s: str) -> bool:\n    used = set(s) - {',', '-', '>'}\n    return len(set(string.ascii_letters) - used) > 0","tryCatchPattern":"try:\n    out = jax.vmap(f)(xs)\nexcept ValueError as e:\n    if 'Run out of characters' in str(e):\n        # manual batching: fold batch axis into an existing subscript letter\n        out = jnp.stack([f(x) for x in xs])\n    else:\n        raise","preventionTips":["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"],"tags":["jax","einsum","vmap","resource-limit"],"backgroundTag":"einsum-subscript-overflow","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}