[Op] Add C++ gradient op for Select.

Cherry-pick from TensorFlow a0dab8a

Author: liutongxuan

tensorflow/cc/gradients/math_grad.cc

@@ -1173,6 +1173,25 @@ Status CastGrad(const Scope& scope, const Operation& op,
 }
 REGISTER_GRADIENT_OP("Cast", CastGrad);
 
+Status SelectGrad(const Scope& scope, const Operation& op,
+                  const std::vector<Output>& grad_inputs,
+                  std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 3) {
+    return errors::InvalidArgument("Select requires 3 arguments");
+  }
+  if (grad_inputs.size() != 1) {
+    return errors::InvalidArgument("Select grad requires 1 grad input");
+  }
+
+  auto c = op.input(0);
+  auto zeros = ZerosLike(scope, grad_inputs[0]);
+  grad_outputs->push_back(NoGradient());  // Condition
+  grad_outputs->push_back(Where3(scope, c, grad_inputs[0], zeros));
+  grad_outputs->push_back(Where3(scope, c, zeros, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Select", SelectGrad);
+
 }  // anonymous namespace
 }  // namespace ops
 }  // namespace tensorflow

tensorflow/cc/gradients/math_grad_test.cc

@@ -52,6 +52,7 @@ using ops::SegmentSum;
 using ops::SquaredDifference;
 using ops::Sub;
 using ops::Sum;
+using ops::Where3;
 
 // TODO(andydavis) Test gradient function against numeric gradients output.
 // TODO(andydavis) As more gradients are added move common test functions
@@ -985,5 +986,14 @@ TEST_F(NaryGradTest, CastGrad) {
   EXPECT_LT(max_error, 1e-3);
 }
 
+TEST_F(NaryGradTest, Select) {
+  TensorShape shape({1, 3});
+  auto cond = Const<bool>(scope_, {{false, true, true}});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto z = Where3(scope_, cond, x, y);
+  RunTest({x, y}, {shape, shape}, {z}, {shape});
+}
+
 }  // namespace
 }  // namespace tensorflow