BUG: Fix einsum optimize logic for singleton dimensions #10352
Conversation
|
@larsoner - thanks for the PR. It took me a long time to see what the heck is going on in the parsing, but now I understand that, basically, the Anyway, I think this PR is correct. @dgasmith - could you have a look? |
|
I need to think on it a bit, that parsing tech is fairly complex as folks want quite the variety of input possibilities. I would consider adding more tests that cover the "double broadcasting" cases where 1-dims are mixed with ellipses like in the original issue. |
|
This might hit issues in the BLAS calls. Please try the following two cases: a = np.random.rand(1, 5)
b = np.random.rand(5, 5)
np.einsum("...ij,...jk->...ik", a, a)
np.einsum("...ij,...jk->...ik", a, b)I suspect we may have to go in and change the BLAS logic to bounce broadcasting cases as well. |
|
Rebased due to conflict and test added. I'm not on my dev machine now so I can't easily run the new tests so I'll let the CIs do the work :) |
|
Yup, ill need to go in and really take a look at whats going on. Ill get to it in the next few days. |
|
Okay @dgasmith I deduplicated the code a bit, feel free to look to see if you're happy. CIs are happy, ready for review/merge from my end. |
|
Can you loop me in on what the |
| left_dims = {dim: size for dim, size in | ||
| zip(input_left, tmp_operands[0].shape)} | ||
| right_dims = {dim: size for dim, size in | ||
| zip(input_right, tmp_operands[1].shape)} |
There was a problem hiding this comment.
Why not use idx_rm here? The other entries in this dict are being ignored
There was a problem hiding this comment.
Here input_left is e.g. 'ik', so this seems like a straightforward / quick way to get the indices to line up. Are you thinking something like this should be done instead?
if any(tmp_operands[0].shape[input_left.index(ind)] != tmp_operands[1].shape[input_right.index(ind)] for ind in idx_rm)
There was a problem hiding this comment.
You're right, what you had was better and probably faster too
It's a way for me (contributor) to signal to you (reviewers) that the PR is ready for review / merge from my perspective. In other projects I work on, maintainers historically have then changed the title as they approve to e.g. |
numpy/core/einsumfunc.py
Outdated
| if 1 in tmp_operands[0] or 1 in tmp_operands[1]: | ||
| if any(tmp_operands[0].shape[input_left.index(ind)] != | ||
| tmp_operands[1].shape[input_right.index(ind)] | ||
| for ind in idx_rm): |
There was a problem hiding this comment.
To reiterate my now-hidden comment, I think what you had before was better - sorry!
There was a problem hiding this comment.
That's fine, easy enough to remove the commit
|
Merged, thanks @larsoner . |
Closes #10343.
I don't know if it's the correct fix because these functions have a lot of logic going on that I have never looked at before (and haven't looked at too closely). Feel free to close if the right solution is a bit more involved.