Function RewriteIfNode

tensorflow/core/common_runtime/lower_if_op.cc:271–292

}  // namespace

Status RewriteIfNode(Node* n, Graph* g, bool keep_node_fetchable) {
  VLOG(2) << "Lower If node (keep_node_fetchable=" << keep_node_fetchable
          << "): " << SummarizeNode(*n);

  const AttrValue* then_attr = n->attrs().Find("then_branch");
  if (then_attr == nullptr) {
    return errors::InvalidArgument("Then branch function missing");
  }
  const AttrValue* else_attr = n->attrs().Find("else_branch");
  if (else_attr == nullptr) {
    return errors::InvalidArgument("Else branch function missing");
  }

  CondBuilder cb(n, then_attr->func(), else_attr->func(), keep_node_fetchable,
                 g);
  TF_RETURN_IF_ERROR(cb.CreatePivotNodes());
  TF_RETURN_IF_ERROR(cb.AddInputs());
  TF_RETURN_IF_ERROR(cb.AddOutputs());
  g->RemoveNode(n);

  return Status::OK();
}

}  // namespace tensorflow