Replace rank inference with shape inference for Einsum op #6010

peishenyan · 2024-03-11T05:52:53Z

In the development of ONNX Runtime, we need know the output shape of each Op node for statical graph compilation. However, we found that we could use onnx shape inference to achieve almost all output shapes except the output shape of Einsum. In onnx/defs/math/defs.cc, we found that there was only Rank Inference function for Einsum instead of Shape Inference.
Given the results of equation parsing and the input shapes of the Einsum operation, we can easily infer the output shape. As such, we have expanded the rank inference to include shape inference for Einsum.

onnx/defs/math/defs.cc

justinchuby · 2024-03-11T22:07:26Z

Thanks for your contribution! Could you check the CI errors?

justinchuby · 2024-03-11T22:09:50Z

It would be helpful to improve the tests for Einsum starting from this line:

onnx/onnx/test/shape_inference_test.py

Line 7023 in a600e2f

def test_einsum_transpose(self) -> None:

Signed-off-by: peishenyan <peishen.yan@intel.com>

codecov · 2024-03-12T06:52:13Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 56.94%. Comparing base (0bb2775) to head (44024b1).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6010      +/-   ##
==========================================
+ Coverage   56.82%   56.94%   +0.11%     
==========================================
  Files         506      506              
  Lines       30377    30461      +84     
  Branches     4592     4592              
==========================================
+ Hits        17263    17347      +84     
  Misses      12285    12285              
  Partials      829      829

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Signed-off-by: peishenyan <peishen.yan@intel.com>

peishenyan · 2024-03-12T07:46:57Z

Hi @justinchuby , thanks for your comments. I have improved the tests for einsum and fix some bugs in my previous code with the help of tests! The code is ready for review now. PTAL, thanks!

onnx/defs/math/defs.cc

justinchuby · 2024-03-13T00:30:49Z

LGTM. I wonder if there are more test cases that should be added? Consider referring to the pytorch tests here: https://github.com/pytorch/pytorch/blob/02bb2180f497d6d10dd11da51fffefaa5af80aca/torch/testing/_internal/common_methods_invocations.py#L6083-L6117

justinchuby · 2024-03-13T00:33:03Z

As well as some complicated ones: 'ijk,ilm,njm,nlk,abc->' (https://numpy.org/doc/stable/reference/generated/numpy.einsum.html)

justinchuby · 2024-03-13T00:34:21Z

Or this table in your reference:

https://fdwr.github.io/MachineLearningOperators/OperatorFormulas.html

Call signature	NumPy equivalent	Description
('i', A1)	A1	returns a view of A1
('i->', A1)	sum(A1)	sums the values of A1
('i,i->i', A1, B1)	A1 * B1	element-wise multiplication of A1 and B1
('i,i->', A1, B1)	inner(A1, B1) or dot(A1, B1)	inner product of A1 and B1
('i,i', A1, B1)	inner(A1, B1) or dot(A1, B1)	inner product of A1 and B1
('i,j->ij', A1, B1)	outer(A1, B1)	outer product of A1 and B1
('ij->ij', A2)	A2	returns a view of A2
('ij', A2)	A2	returns a view of A2
('ji', A2)	A2.T	view transpose of A2
('ji->ij', A2)	A2.T	view transpose of A2
('ii->i', A2)	diag(A2)	view main diagonal of A2
('ii->', A2)	trace(A2)	sums main diagonal of A2
('ij->', A2)	sum(A2)	sums the values of A2
('ij->j', A2)	sum(A2, axis=0)	sum down the columns of A2 (across rows)
('ij->i', A2)	sum(A2, axis=1)	sum horizontally along the rows of A2
('ij,ij->ij', A2, B2)	A2 * B2	element-wise multiplication of A2 and B2
('ij,ji->ij', A2, B2)	A2 * B2.transpose()	element-wise multiplication of A2 and B2.T
('ij,jk', A2, B2)	matmul(A2, B2) or dot(A2, B2)	matrix multiplication of A2 and B2
('ij,jk->ik', A2, B2)	matmul(A2, B2) or dot(A2, B2)	matrix multiplication of A2 and B2
('bij,bjk->bik', A2, B2)	matmul(A3, B3)	matrix multiplication of A3 and B3 (a stack of 2D matrices)
('bij,bkj->bik', A2, B2)	matmul(A3, transpose(B3))	matrix multiplication of A3 and B3 (a stack of 2D matrices)
('ij,kj->ik', A2, B2)	inner(A2, B2)	inner product of A2 and B2
('ij,kj->ikj', A2, B2)	A2[:, None] * B2	each row of A2 multiplied by B2
('ij,kl->ijkl', A2, B2)	A2[:, :, None, None] * B2	each value of A2 multiplied by B2
(',ij', 3, B2)		Scalar times array: array([[ 0, 3, 6], [ 9, 12, 15]])
("ij,j", A2, B1)	matvec(A2, B1)	Matrix and vector.
("ii,ii->i", A2, B2)	A2.diag() * B2.diag()	diagonals multiplied by each other
("ii,ii->", A2, B2)	dot(A2.diag(), B2.diag())	dot product of diagonals

You may leverage parameterized to parameterize the tests

…tps://fdwr.github.io/MachineLearningOperators/OperatorFormulas.html Signed-off-by: peishenyan <peishen.yan@intel.com>

Signed-off-by: peishenyan <peishen.yan@intel.com>

justinchuby

Thank you!

onnx/test/shape_inference_test.py

Signed-off-by: peishenyan <peishen.yan@intel.com>

onnx/defs/math/defs.cc

justinchuby · 2024-03-22T00:29:30Z

@gramalingam @xadupre could you help approve? Thanks!

fdwr · 2024-03-22T01:15:52Z

That reminds me, I should review this one too (I only reviewed the ORT WebNN EP one)...

fdwr

Two minor comments.

onnx/defs/math/defs.cc

Signed-off-by: peishenyan <peishen.yan@intel.com>

fdwr

👍

onnx/defs/math/defs.cc

+        continue;
+      }
+
+      const auto inserted = label_maps.insert({term[index], num_labels}).second;


onnx/defs/math/defs.cc

      }
    }
  } else { // Infer the dimension for right-hand side
-    // If there's an ellipsis, add it's corresponding dimensions
+    // If there's an ellipsis, add its corresponding dimensions


onnx/defs/math/defs.cc

Signed-off-by: peishenyan <peishen.yan@intel.com>

gramalingam · 2024-03-22T06:02:41Z

@fdwr : do you know if the onnxruntime implementation supports broadcasting across the ellipsis dimensions? The ONNX spec says "broadcasting" ... which would mean the shape-inference should do the broadcasting logic for the ellipsis dimensions (or, at least, leave them as unknown).

fdwr · 2024-03-22T06:53:06Z

@fdwr : do you know if the onnxruntime implementation supports broadcasting across the ellipsis dimensions? The ONNX spec says "broadcasting" ... which would mean the shape-inference should do the broadcasting logic for the ellipsis dimensions (or, at least, leave them as unknown).

~~🤔 I'm not sure about ellipsis broadcasting without spelunking. What I know is:~~
(a) the ORT DML EP does not (we fallback to the CPU EP if an ellipsis is found inside the equation)
(b) the ORT CPU EP supports at least some degree of ellipsis support, because I had test cases for ...ii->...i, which passed.

(update) Judging from this comment in the CPU EP, it probably does support ellipsis broadcasting:

  // A function to process broadcasted dims (ellipsis) of inputs that they occur in
  Status PostProcessBroadcastedDims();

Signed-off-by: peishenyan <peishen.yan@intel.com>

fdwr

👍

gramalingam

Thanks for the extension ... greatly appreciated.

In the development of ONNX Runtime, we need know the output shape of each Op node for statical graph compilation. However, we found that we could use onnx shape inference to achieve almost all output shapes except the output shape of Einsum. In `onnx/defs/math/defs.cc`, we found that there was only Rank Inference function for Einsum instead of Shape Inference. Given the results of equation parsing and the input shapes of the Einsum operation, we can easily infer the output shape. As such, we have expanded the rank inference to include shape inference for Einsum. --------- Signed-off-by: peishenyan <peishen.yan@intel.com> Co-authored-by: G. Ramalingam <grama@microsoft.com> Signed-off-by: Linsho Kaku <linsho@preferred.jp>

…22376) ### Description  Pick up onnx/onnx#6010 to support EinSum shape inference ### Motivation and Context  This change allows EinSum operator's output shape to be inferenced so that it can run on accelerators.

peishenyan requested a review from a team as a code owner March 11, 2024 05:52

peishenyan mentioned this pull request Mar 11, 2024

[Feature request] Shape Inference for Einsum instead of Rank Inference #6011

Closed

justinchuby added the module: shape inference Issues related to shape inference label Mar 11, 2024

justinchuby added this to the 1.17 milestone Mar 11, 2024

github-advanced-security bot found potential problems Mar 11, 2024

View reviewed changes

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

peishenyan added 2 commits March 12, 2024 14:37

replace rank inference with shape inference for einsum op

f6027b1

Signed-off-by: peishenyan <peishen.yan@intel.com>

fix bugs and add improve the tests for einsum

6ce4699

Signed-off-by: peishenyan <peishen.yan@intel.com>

peishenyan requested a review from a team as a code owner March 12, 2024 06:45

add comments

15785eb

Signed-off-by: peishenyan <peishen.yan@intel.com>

peishenyan force-pushed the einsum_shape_inference branch from bf9f7b1 to 6ce4699 Compare March 12, 2024 07:09

remove unused code

bc5d6fa

Signed-off-by: peishenyan <peishen.yan@intel.com>

peishenyan mentioned this pull request Mar 12, 2024

[WebNN EP] Support Einsum op microsoft/onnxruntime#19558

Merged

justinchuby reviewed Mar 12, 2024

View reviewed changes

onnx/defs/math/defs.cc Outdated Show resolved Hide resolved

github-advanced-security bot found potential problems Mar 12, 2024

View reviewed changes

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

github-advanced-security bot found potential problems Mar 12, 2024

View reviewed changes

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

onnx/defs/math/defs.cc Fixed Show fixed Hide fixed

justinchuby self-assigned this Mar 13, 2024

peishenyan added 2 commits March 15, 2024 10:12

cover all test case for einsum according to the operator formulas: ht…

086345f

…tps://fdwr.github.io/MachineLearningOperators/OperatorFormulas.html Signed-off-by: peishenyan <peishen.yan@intel.com>

lint and apply the fixes

e8eaf7e

Signed-off-by: peishenyan <peishen.yan@intel.com>

justinchuby approved these changes Mar 15, 2024

View reviewed changes

github-advanced-security bot found potential problems Mar 15, 2024

View reviewed changes

onnx/test/shape_inference_test.py Fixed Show fixed Hide fixed

onnx/test/shape_inference_test.py Fixed Show fixed Hide fixed

address comments

31b900c

Signed-off-by: peishenyan <peishen.yan@intel.com>

justinchuby added the review needed: operators approvers Require reviews from members of operators-approvers label Mar 15, 2024

justinchuby reviewed Mar 15, 2024

View reviewed changes

onnx/defs/math/defs.cc Outdated Show resolved Hide resolved

justinchuby enabled auto-merge March 22, 2024 00:29

fdwr approved these changes Mar 22, 2024

View reviewed changes

onnx/defs/math/defs.cc Outdated Show resolved Hide resolved

onnx/defs/math/defs.cc Outdated Show resolved Hide resolved

address comments

25ae21c

Signed-off-by: peishenyan <peishen.yan@intel.com>

auto-merge was automatically disabled March 22, 2024 04:05
Head branch was pushed to by a user without write access

fdwr approved these changes Mar 22, 2024

View reviewed changes

github-advanced-security bot found potential problems Mar 22, 2024

View reviewed changes

gramalingam reviewed Mar 22, 2024

View reviewed changes

onnx/defs/math/defs.cc Show resolved Hide resolved

gramalingam reviewed Mar 22, 2024

View reviewed changes

onnx/defs/math/defs.cc Show resolved Hide resolved

fix ellipsis broadcast error and add one test for it

fd00114

Signed-off-by: peishenyan <peishen.yan@intel.com>

support for lower and upper case

9bd76da

Signed-off-by: peishenyan <peishen.yan@intel.com>

fdwr approved these changes Mar 22, 2024

View reviewed changes

justinchuby mentioned this pull request Mar 22, 2024

[Shape Inference] Einsum: Ellipsis represents incompatible dimensions. #5739

Closed

Merge branch 'main' into einsum_shape_inference

44024b1

gramalingam approved these changes Mar 22, 2024

View reviewed changes

justinchuby enabled auto-merge March 22, 2024 22:08

justinchuby added this pull request to the merge queue Mar 22, 2024

Merged via the queue into onnx:main with commit 5b191de Mar 22, 2024

fdwr mentioned this pull request Jun 21, 2024

DML EP EinSum make more generic to avoid EP fallback microsoft/onnxruntime#21114

Merged

indygit mentioned this pull request Oct 10, 2024

Pick changes from onnx/onnx#6010 to support EinSum shape inference microsoft/onnxruntime#22376

Merged

yuslepukhin mentioned this pull request May 19, 2025

Segfault in Einsum shape inference #6981

Closed

Copilot AI mentioned this pull request Jun 19, 2025

Fix Einsum shape inference segfault for scalar inputs #7055

Merged

dependabot bot mentioned this pull request Sep 8, 2025

deps: Bump Microsoft.ML.OnnxRuntime from 1.19.0 to 1.22.1 rivoli-ai/andy-guard#17

Open

Replace rank inference with shape inference for Einsum op #6010

Replace rank inference with shape inference for Einsum op #6010

Uh oh!

Conversation

peishenyan commented Mar 11, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Mar 11, 2024

Uh oh!

justinchuby commented Mar 11, 2024

Uh oh!

codecov bot commented Mar 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

peishenyan commented Mar 12, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Mar 13, 2024

Uh oh!

justinchuby commented Mar 13, 2024

Uh oh!

justinchuby commented Mar 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

justinchuby left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Mar 22, 2024

Uh oh!

fdwr commented Mar 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fdwr left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

fdwr left a comment

Choose a reason for hiding this comment

Uh oh!

Check notice

Check notice

Uh oh!

Uh oh!

gramalingam commented Mar 22, 2024

Uh oh!

fdwr commented Mar 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fdwr left a comment

Choose a reason for hiding this comment

Uh oh!

gramalingam left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov bot commented Mar 12, 2024 •

edited

Loading

justinchuby commented Mar 13, 2024 •

edited

Loading

fdwr commented Mar 22, 2024 •

edited

Loading

fdwr commented Mar 22, 2024 •

edited

Loading