message: Update shape function for `tf.where` / `SelectOp`
diff: @@ -758,6 +758,57 @@ REGISTER_OP("Select")
  return Status::OK();
  });
  +REGISTER_OP("SelectV2")
  + .Input("condition: bool")
  + .Input("t: T")
  + .Input("e: T")
  + .Output("output: T")
  + .Attr("T: type")
  + .SetShapeFn([](InferenceContext* c) {
  + auto* handle_data_1 = c->input_handle_shapes_and_types(1);
  + auto* handle_data_2 = c->input_handle_shapes_and_types(2);
  + // Merge handle shape and dtype if applicable.
  + if (handle_data_1 != nullptr && handle_data_2 != nullptr) {
  + const auto size = handle_data_1->size();
  + std::vector<shape_inference::ShapeAndType> merged_handle_data(size);
  + if (size != handle_data_2->size()) {
  + return errors::InvalidArgument(
  + "Trying to merge handles pointing to different numbers of "
  + "tensors.");
  + }
  +
  + for (int i = 0; i < size; ++i) {
  + const shape_inference::ShapeAndType& s1 = (*handle_data_1)[i];
  + const shape_inference::ShapeAndType& s2 = (*handle_data_2)[i];
  + if (s1.dtype != s2.dtype) {
  + // TODO(apassos) resolve this in the manner of b/32476923
  + return errors::InvalidArgument(
  + "Trying to merge handles pointing to different dtypes.");
  + }
  + merged_handle_data[i].dtype = s1.dtype;
  + TF_RETURN_IF_ERROR(
  + c->Merge(s1.shape, s2.shape, &merged_handle_data[i].shape));
  + }
  +
  + c->set_output_handle_shapes_and_types(0, merged_handle_data);
  + }
  +
  + // The inputs 'cond', 'then', and 'else' must be broadcastable.
  + // TODO (yongtang): Consolidate 3-ary broadcast instead of
  + // multiple 2-ary broadcast.
  + ShapeHandle cond = c->input(0);
  + ShapeHandle then = c->input(1);
  + ShapeHandle else_ = c->input(2);
  + ShapeHandle other;
  + TF_RETURN_IF_ERROR(
  + BroadcastBinaryOpOutputShapeFnHelper(c, then, else_, &other));
  + ShapeHandle output;
  + TF_RETURN_IF_ERROR(
  + BroadcastBinaryOpOutputShapeFnHelper(c, cond, other, &output));
  + c->set_output(0, output);
  + return Status::OK();
  + });
  +
  // --------------------------------------------------------------------------
  REGISTER_OP("MatMul")
  