Fix einsum dropping a trailing empty subscript - #4299
Open
Adityaj0 wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #4298.
Proposed changes
An empty subscript denotes a scalar operand, but
parse()tokenized the left hand side withgetline:getlinestops 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:A leading empty subscript worked, because
getlinedoes 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.
",i->i"(3,)(3,)(3,)"i,->i"(3,)(3,)",->"()()"->"()()"ij,->ij"(2,3)(2,3)",,->"()()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 intests/einsum_tests.cppcontinue to throw:Tests
Added
test_scalar_operandstopython/tests/test_einsum.pycovering 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.pre-commit run --all-filesto format my code and installed pre-commit prior to committing changesVerified with a CPU-only build (
-DMLX_BUILD_METAL=OFF):test_einsum,test_ops,test_autograd,test_vmap,test_linalg,test_blas,test_nn,test_array,test_reduce(464 tests) passThis is independent of #4125, which touches
batch_tensordotin the same file but a different function.