From efd9b924824922e9f281e1859fbfecf963e176c1 Mon Sep 17 00:00:00 2001
From: Dmitri Smirnov <yuslepukhin@users.noreply.github.com>
Date: Mon, 13 Apr 2020 16:24:35 -0700
Subject: [PATCH] Handle Scalars in TernaryOps and Where. (#3509)

Handle Scalars in TernaryOps and Where.
---
 onnxruntime/core/providers/cpu/tensor/utils.h |  5 +++
 .../core/providers/cuda/tensor/where.cc       | 42 +++++++------------
 .../providers/cpu/tensor/where_op_test.cc     | 15 +++++++
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/tensor/utils.h b/onnxruntime/core/providers/cpu/tensor/utils.h
index c905fa32b8..b8ad186e9c 100644
--- a/onnxruntime/core/providers/cpu/tensor/utils.h
+++ b/onnxruntime/core/providers/cpu/tensor/utils.h
@@ -25,6 +25,11 @@ struct TensorPitches : std::vector<int64_t> {
     if (gsl::narrow_cast<ptrdiff_t>(padded_rank) < 0)
       return false;
 
+    // Guard against Scalars
+    if (pitch_rank == 0) {
+      return true;
+    }
+
     *(p.rbegin()) = 1;  // The innermost axis is 1 (single values)
     if (tensor_rank > 1) {
       for (size_t i = tensor_rank - 1; i-- > 0;) {
diff --git a/onnxruntime/core/providers/cuda/tensor/where.cc b/onnxruntime/core/providers/cuda/tensor/where.cc
index 3fc0b9272e..e2dd66e3cf 100644
--- a/onnxruntime/core/providers/cuda/tensor/where.cc
+++ b/onnxruntime/core/providers/cuda/tensor/where.cc
@@ -94,40 +94,30 @@ struct TernaryElementwisePreparation {
 
     output_rank_or_simple_broadcast = out_rank;
 
-    if (a_shape != output_shape) {
-      TensorPitches a_pitches(a_shape.GetDims());
-      a_padded_strides.size_ = out_rank;
-      auto offset = out_rank - a_rank;
-      for (auto i = offset; i < out_rank; ++i) {
-        // the stride for broadcast dimension is kept as 0
-        if (a_shape.GetDims()[i - offset] != 1) {
-          a_padded_strides[i] = a_pitches[i];
+    auto padder = [out_rank](int32_t rank, const TensorShape& shape, TArray<int64_t>& padded_strides) {
+      padded_strides.size_ = out_rank;
+      if (rank > 0) {
+        TensorPitches pitches(shape.GetDims());
+        auto offset = out_rank - rank;
+        for (auto i = offset; i < out_rank; ++i) {
+          // the stride for broadcast dimension is kept as 0
+          if (shape.GetDims()[i - offset] != 1) {
+            padded_strides[i] = pitches[i - offset];
+          }
         }
       }
+    };
+
+    if (a_shape != output_shape) {
+      padder(a_rank, a_shape, a_padded_strides);
     }
 
     if (b_shape != output_shape) {
-      TensorPitches b_pitches(b_shape.GetDims());
-      b_padded_strides.size_ = out_rank;
-      auto offset = out_rank - b_rank;
-      for (auto i = offset; i < out_rank; ++i) {
-        // the stride for broadcast dimension is kept as 0
-        if (b_shape.GetDims()[i - offset] != 1) {
-          b_padded_strides[i] = b_pitches[i];
-        }
-      }
+      padder(b_rank, b_shape, b_padded_strides);
     }
 
     if (c_shape != output_shape) {
-      TensorPitches c_pitches(c_shape.GetDims());
-      c_padded_strides.size_ = out_rank;
-      auto offset = out_rank - c_rank;
-      for (auto i = offset; i < out_rank; ++i) {
-        // the stride for broadcast dimension is kept as 0
-        if (c_shape.GetDims()[i - offset] != 1) {
-          c_padded_strides[i] = c_pitches[i];
-        }
-      }
+      padder(c_rank, c_shape, c_padded_strides);
     }
 
     TensorPitches output_pitches(output_shape.GetDims());
diff --git a/onnxruntime/test/providers/cpu/tensor/where_op_test.cc b/onnxruntime/test/providers/cpu/tensor/where_op_test.cc
index e953761df0..cc941cd12d 100644
--- a/onnxruntime/test/providers/cpu/tensor/where_op_test.cc
+++ b/onnxruntime/test/providers/cpu/tensor/where_op_test.cc
@@ -122,5 +122,20 @@ TEST(WhereOpTest, BroadcastDimWithZero) {
   // exclude NGraph as this isn't handled by that EP
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider});
 }
+
+TEST(WhereOpTest, BroadcastWithScalar) {
+  OpTester test{kOpName, kOpVersion};
+
+  test.AddInput<bool>("condition", {3}, {true, false, true});
+  test.AddInput<int64_t>("X", {1, 3}, {1, 2, 3});
+  test.AddInput<int64_t>("Y", {}, {1});
+
+  test.AddOutput<int64_t>("output", {1, 3}, {1, 1, 3});
+
+  // exclude NGraph as this isn't handled by that EP
+  //test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider});
+  test.Run();
+}
+
 }  // namespace test
 }  // namespace onnxruntime