[Op] Fix an issue with SelectV2 gradient broadcasting.

The Where3 op used previously doesn't have the same broadcasting
semantics as SelectV2, which caused an error when computing gradients.

Cherry-pick from TensorFlow 42fad1e

Author: liutongxuan

tensorflow/cc/gradients/math_grad.cc

@@ -1221,7 +1221,7 @@ Status SelectV2Grad(const Scope& scope, const Operation& op,
   auto y = op.input(2);
 
   auto zeros = ZerosLike(scope, grad_inputs[0]);
-  auto gx = Where3(scope, c, grad_inputs[0], zeros);
+  auto gx = SelectV2(scope, c, grad_inputs[0], zeros);
   auto x_shape = Shape(scope, x);
   auto output_shape = Shape(scope, op.output(0));
 
@@ -1231,7 +1231,7 @@ Status SelectV2Grad(const Scope& scope, const Operation& op,
       ReduceSum(scope, gx, /*axis=*/reduce_x.r0, ReduceSum::KeepDims(true));
   auto gx_sum_reshape = Reshape(scope, gx_sum, x_shape);
 
-  auto gy = Where3(scope, c, zeros, grad_inputs[0]);
+  auto gy = SelectV2(scope, c, zeros, grad_inputs[0]);
   auto y_shape = Shape(scope, y);
 
   // Reduce away broadcasted leading dims.

tensorflow/cc/gradients/math_grad_test.cc

@@ -1016,6 +1016,15 @@ TEST_F(NaryGradTest, SelectV2_Broadcast) {
   RunTest({x, y}, {x_shape, y_shape}, {z}, {x_shape});
 }
 
+TEST_F(NaryGradTest, SelectV2_Broadcast2) {
+  TensorShape x_shape({2, 3});
+  auto cond = Const<bool>(scope_, {{false}, {true}});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto z = SelectV2(scope_, cond, x, y);
+  RunTest({x, y}, {x_shape, x_shape}, {z}, {x_shape});
+}
+
 TEST_F(NaryGradTest, Atan2Grad) {
   TensorShape shape({3, 2, 5});
   auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));