Skip to content

Fix einsum dropping a trailing empty subscript - #4299

Open
Adityaj0 wants to merge 1 commit into
ml-explore:mainfrom
Adityaj0:fix-einsum-scalar-operands
Open

Fix einsum dropping a trailing empty subscript#4299
Adityaj0 wants to merge 1 commit into
ml-explore:mainfrom
Adityaj0:fix-einsum-scalar-operands

Conversation

@Adityaj0

Copy link
Copy Markdown
Contributor

Fixes #4298.

Proposed changes

An empty subscript denotes a scalar operand, but parse() tokenized the left hand side with getline:

std::stringstream ss(lhs);
while (getline(ss, token, ',')) {
  input_list.push_back(token);
}

getline stops at the last delimiter, so a trailing empty field is never produced. "i," parsed as ["i"] instead of ["i", ""], and the operand-count check then rejected the call:

mx.einsum("i,->i", mx.zeros((3,)), mx.array(2.0))
# ValueError: [einsum] Number of operands, 2, does not match number of input subscripts, 1

A leading empty subscript worked, because getline does emit an empty first field, which is why ",i->i" was accepted while "i,->i" was not. That asymmetry is what makes it look like a tokenizer bug rather than an intentional restriction.

This splits on commas directly so every field is kept, including trailing and all-empty ones.

equation NumPy before after
",i->i" (3,) (3,) (3,)
"i,->i" (3,) error (3,)
",->" () error ()
"->" () error ()
"ij,->ij" (2,3) error (2,3)
",,->" () error ()

The validation that rejects genuinely bad input is untouched: the operand-count check and the per-operand in.size() != operands[i].ndim() check still run, so the existing negative cases in tests/einsum_tests.cpp continue to throw:

CHECK_THROWS(einsum("", {}));                 // 1 subscript, 0 operands
CHECK_THROWS(einsum("", {array({1, 2})}));    // "" requires a 0-d operand

Tests

Added test_scalar_operands to python/tests/test_einsum.py covering scalar operands in leading, trailing and middle positions, all-scalar equations, and "->", each checked against NumPy. It also asserts the negative cases still raise: mismatched operand count, and an empty subscript paired with a non 0-d operand.

  • I have read the CONTRIBUTING document
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (not needed, no API change)
  • I have run pre-commit run --all-files to format my code and installed pre-commit prior to committing changes

Verified with a CPU-only build (-DMLX_BUILD_METAL=OFF):

  • C++ suite: 247 test cases, 3326 assertions, all pass
  • Python: test_einsum, test_ops, test_autograd, test_vmap, test_linalg, test_blas, test_nn, test_array, test_reduce (464 tests) pass

This is independent of #4125, which touches batch_tensordot in the same file but a different function.

An empty subscript denotes a scalar operand, but the parser tokenized the
left hand side with getline:

    std::stringstream ss(lhs);
    while (getline(ss, token, ',')) {
      input_list.push_back(token);
    }

getline stops at the last delimiter, so a trailing empty field is never
produced. "i,->i" parsed as a single input and einsum rejected the call:

    mx.einsum("i,->i", mx.zeros((3,)), mx.array(2.0))
    ValueError: [einsum] Number of operands, 2, does not match number of
    input subscripts, 1

A leading empty subscript worked, since getline does emit an empty first
field, which is why ",i->i" was fine while "i,->i" was not.

Split on commas directly so every field is kept. Operand count and per
operand dimension checks are unchanged, so einsum("", {}) and
einsum("", {1-d array}) still throw as before.
@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

einsum rejects scalar operands when the empty subscript is not first

2 participants