Fix einsum operator for empty inputs

onnxruntime/core/providers/cpu/math/einsum_utils/einsum_typed_compute_processor.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <algorithm>
+
 #include "einsum_typed_compute_processor.h"
 #include "core/common/narrow.h"
 #include "core/common/span_utils.h"
@@ -357,6 +359,33 @@ Status EinsumTypedComputeProcessor<T>::Run() {
 
   auto num_inputs = context_->InputCount();
 
+  {
+    bool has_empty_input = std::any_of(raw_inputs.begin(), raw_inputs.end(), [](const auto& input) {
+      return input->Shape().Size() == 0;
+    });
+    // Skip all the work, fill with zeros if needed
+    if (has_empty_input) {
+      AllocatorPtr cpu_allocator;
+      ORT_RETURN_IF_ERROR(context_->GetTempSpaceAllocator(&cpu_allocator));
+
+      const auto output_dims = einsum_compute_preprocessor_.GetOutputDims();
+      Tensor& output = *context_->Output(0, output_dims);
+      Tensor candidate_output(raw_inputs[0]->DataType(), output_dims, cpu_allocator);
+
+      if constexpr (std::is_integral<T>::value) {
+        std::fill_n(reinterpret_cast<T*>(candidate_output.MutableDataRaw()), candidate_output.Shape().Size(), T(0));
+      } else {
+        std::fill_n(reinterpret_cast<T*>(candidate_output.MutableDataRaw()), candidate_output.Shape().Size(), T(0.f));
+      }
+
+      auto status = device_data_copy_func_(candidate_output, output, einsum_ep_assets_);
+      ORT_ENFORCE(status.IsOK(), "Einsum op: Could not copy the intermediate output's buffer into the op's output buffer. Error: ",
+                  status.ErrorMessage());
+
+      return status;
+    }
+  }
+
   // Pre-process the first input so as to reduce any dims that only it has
   std::unique_ptr<const Tensor> result;
 

onnxruntime/test/providers/cpu/math/einsum_test.cc

@@ -755,6 +755,52 @@ TEST(Einsum, ExplicitEinsumAsTensorContraction_Half) {
   test.Run();
 }
 
+// Theme: Empty inputs
+TEST(Einsum, EinsumEmptyInputOuterProduct) {
+  OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
+  test.AddAttribute<std::string>("equation", "i,j->ij");
+  test.AddInput<float>("x", {0}, {});
+  test.AddInput<float>("y", {0}, {});
+  test.AddOutput<float>("o", {0, 0}, {});
+  test.Run();
+}
+
+TEST(Einsum, EinsumEmptyInputTranspose) {
+  OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
+  test.AddAttribute<std::string>("equation", "ab,ba->ab");
+  test.AddInput<float>("x", {0, 1}, {});
+  test.AddInput<float>("y", {1, 0}, {});
+  test.AddOutput<float>("o", {0, 1}, {});
+  test.Run();
+}
+
+TEST(Einsum, EinsumEmptyInputVanish) {
+  OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
+  test.AddAttribute<std::string>("equation", "ab,ba->a");
+  test.AddInput<float>("x", {1, 0}, {});
+  test.AddInput<float>("y", {0, 1}, {});
+  test.AddOutput<float>("o", {1}, {0.f});
+  test.Run();
+}
+
+TEST(Einsum, EinsumEmptyInputVanish3d) {
+  OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
+  test.AddAttribute<std::string>("equation", "abc,bad->ad");
+  test.AddInput<float>("x", {10, 0, 10}, {});
+  test.AddInput<float>("y", {0, 10, 1}, {});
+  test.AddOutput<float>("o", {10, 1}, std::vector<float>(10, 0.f));
+  test.Run();
+}
+
+TEST(Einsum, EinsumEmptyInputVanish3d2empty) {
+  OpTester test("Einsum", 12, onnxruntime::kOnnxDomain);
+  test.AddAttribute<std::string>("equation", "abc,bcd->ad");
+  test.AddInput<float>("x", {0, 0, 0}, {});
+  test.AddInput<float>("y", {0, 0, 1}, {});
+  test.AddOutput<float>("o", {0, 1}, {});
+  test.Run();
+}
+
 // Theme: Tests involving MatMul(s) interleaved with Transpose(s)
 // for two and three inputs (most common use-case of Einsum operator)