Merge pull request #1420 from gs-olive/einsum_schema_fix

fix: `aten::einsum` schema switch for Torch 1.14 nightly
pytorch · Oct 28, 2022 · e75123c · e75123c
2 parents 3053ecc + f6b5bad
commit e75123c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/core/conversion/converters/impl/einsum.cpp b/core/conversion/converters/impl/einsum.cpp
@@ -12,7 +12,7 @@ namespace impl {
 namespace {
 
 auto einsum_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
-    {"aten::einsum(str equation, Tensor[] tensors) -> (Tensor)",
+    {"aten::einsum(str equation, Tensor[] tensors, *, int[]? path=None) -> (Tensor)",
      [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
        // Extract equation and list of tensors
        auto equation = args[0].unwrapToString();

diff --git a/tests/core/conversion/converters/test_einsum.cpp b/tests/core/conversion/converters/test_einsum.cpp
@@ -9,7 +9,8 @@ TEST(Converters, ATenEinsumConvertsMatMulCorrectly) {
       graph(%x.1 : Tensor, %x.2 : Tensor):
         %0 : str = prim::Constant[value="ij,jk->ik"]()
         %3 : Tensor[] = prim::ListConstruct(%x.1, %x.2)
-        %4 : Tensor = aten::einsum(%0, %3)
+        %none : NoneType = prim::Constant()
+        %4 : Tensor = aten::einsum(%0, %3, %none)
         return (%4))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();
@@ -34,7 +35,8 @@ TEST(Converters, ATenEinsumConvertsElementwiseProdCorrectly) {
       graph(%x.1 : Tensor, %x.2 : Tensor):
         %0 : str = prim::Constant[value="abcd,abcd->abcd"]()
         %3 : Tensor[] = prim::ListConstruct(%x.1, %x.2)
-        %4 : Tensor = aten::einsum(%0, %3)
+        %none : NoneType = prim::Constant()
+        %4 : Tensor = aten::einsum(%0, %3, %none)
         return (%4))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();
@@ -59,7 +61,8 @@ TEST(Converters, ATenEinsumConvertsTransposeCorrectly) {
       graph(%x.1 : Tensor):
         %0 : str = prim::Constant[value="jk->kj"]()
         %3 : Tensor[] = prim::ListConstruct(%x.1)
-        %4 : Tensor = aten::einsum(%0, %3)
+        %none : NoneType = prim::Constant()
+        %4 : Tensor = aten::einsum(%0, %3, %none)
         return (%4))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();
@@ -83,7 +86,8 @@ TEST(Converters, ATenEinsumConvertsVectorsCorrectly) {
       graph(%x.1 : Tensor, %x.2 : Tensor):
         %0 : str = prim::Constant[value="a,b->ab"]()
         %3 : Tensor[] = prim::ListConstruct(%x.1, %x.2)
-        %4 : Tensor = aten::einsum(%0, %3)
+        %none : NoneType = prim::Constant()
+        %4 : Tensor = aten::einsum(%0, %3, %none)
         return (%4))IR";
 
   auto g = std::make_shared<torch::jit::Graph>();