From 4bb41d610bf2f1d822bb438631655931f83afc6b Mon Sep 17 00:00:00 2001
From: RandySheriffH <48490400+RandySheriffH@users.noreply.github.com>
Date: Thu, 19 Nov 2020 11:25:53 -0800
Subject: [PATCH] Address scalar expand to scalar case (#5854)

* address scalar expand to scalar case

* remove tensorrt unsupported type

* exclude tensorrt from scalar to scalar case

Co-authored-by: Randy Shuai <rashuai@microsoft.com>
---
 .../core/providers/cpu/tensor/expand.cc       |  6 ++++++
 .../test/providers/cpu/tensor/expand_test.cc  | 20 +++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/onnxruntime/core/providers/cpu/tensor/expand.cc b/onnxruntime/core/providers/cpu/tensor/expand.cc
index 68e494c117..4b9c3870c8 100644
--- a/onnxruntime/core/providers/cpu/tensor/expand.cc
+++ b/onnxruntime/core/providers/cpu/tensor/expand.cc
@@ -72,6 +72,12 @@ Status Expand<T>::Compute(OpKernelContext* context) const {
   auto* output_dims = output_shape.data();
   auto output_dims_size = static_cast<int64_t>(output_shape.size());
   auto max_dims_size = std::max(input_dims_size, output_dims_size);
+
+  if (0 == max_dims_size) {
+    *output_data = *input_data;
+    return Status::OK();
+  }
+
   std::unique_ptr<int64_t[]> input_dim_group{new int64_t[max_dims_size]};
   std::unique_ptr<int64_t[]> output_dim_group{new int64_t[max_dims_size]};
   std::unique_ptr<int64_t[]> expand_dim_size{new int64_t[max_dims_size]};
diff --git a/onnxruntime/test/providers/cpu/tensor/expand_test.cc b/onnxruntime/test/providers/cpu/tensor/expand_test.cc
index f3b45775aa..5b71f69d58 100644
--- a/onnxruntime/test/providers/cpu/tensor/expand_test.cc
+++ b/onnxruntime/test/providers/cpu/tensor/expand_test.cc
@@ -163,5 +163,25 @@ TEST(ExpandOpTest, Expand_2x2x1x2x1_float) {
   test.Run();
 }
 
+#ifndef USE_TENSORRT
+TEST(ExpandOpTest, Expand_scalar_float) {
+  OpTester test("Expand", 8);
+  test.AddInput<float>("data_0", {}, {3.0f});
+  test.AddInput<int64_t>("data_1", {0}, {});
+  test.AddOutput<float>("result", {}, {3.0f});
+  test.Run();
+}
+#endif
+
+TEST(ExpandOpTest, Expand_scalar_int32) {
+  OpTester test("Expand", 8);
+  test.AddInput<int32_t>("data_0", {}, {9});
+  test.AddInput<int64_t>("data_1", {3}, {2, 3, 4});
+  test.AddOutput<int32_t>("result", {2, 3, 4},
+                         {9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
+                          9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9});
+  test.Run();
+}
+
 }  //namespace test
 }  //namespace onnxruntime