chore: import upstream snapshot with attribution

tensorflow/cc/framework/cc_op_gen.cc

@@ -0,0 +1,446 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/cc_op_gen.h"
+
+#include <memory>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/strings/escaping.h"
+#include "tensorflow/cc/framework/cc_op_gen_util.h"
+#include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/op_def_util.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+namespace cc_op {
+namespace {
+
+const int kRightMargin = 79;
+
+std::string GetConstructorDecl(const OpInfo& op_info,
+                               absl::string_view op_name_prefix,
+                               bool include_attr) {
+  const std::string prefix = absl::StrCat(op_name_prefix, op_info.op_name, "(");
+  std::string c_decl;
+  for (int i = 0; i < op_info.arg_types.size(); ++i) {
+    if (i > 0) absl::StrAppend(&c_decl, ", ");
+    absl::StrAppend(&c_decl, op_info.arg_types[i], " ", op_info.arg_names[i]);
+  }
+  if (include_attr && op_info.has_optional_attrs) {
+    absl::StrAppend(&c_decl, ", const ", op_info.op_name, "::Attrs& attrs");
+  }
+  absl::StrAppend(&c_decl, ")");
+  return WordWrap(prefix, c_decl, kRightMargin);
+}
+
+void WriteClassDecl(const OpInfo& op_info, WritableFile* h) {
+  std::string class_decl = op_info.comment;
+  absl::StrAppend(&class_decl, "class ", op_info.op_name, " {\n");
+  absl::StrAppend(&class_decl, " public:\n");
+  if (op_info.has_optional_attrs) {
+    absl::StrAppend(&class_decl, op_info.GetOpAttrStruct());
+  }
+  absl::StrAppend(&class_decl, "  ",
+                  GetConstructorDecl(op_info, "", /* include_attr */ false),
+                  ";\n");
+  if (op_info.has_optional_attrs) {
+    absl::StrAppend(&class_decl, "  ",
+                    GetConstructorDecl(op_info, "", /* include_attr */ true),
+                    ";\n");
+  }
+  if (op_info.output_types.empty()) {
+    // Allow casting this class to Operation.
+    absl::StrAppend(&class_decl,
+                    "  operator ::tensorflow::Operation() const { "
+                    "return operation; }\n");
+  } else if (op_info.output_types.size() == 1) {
+    if (op_info.is_list_output[0]) {
+      // Write the subscript operator, allowing out[i] for the list-typed
+      // output.
+      absl::StrAppend(&class_decl,
+                      "  ::tensorflow::Output operator[](size_t index) "
+                      "const { return ",
+                      op_info.output_names[0], "[index]; }\n\n");
+
+    } else {
+      // Write type cast functions, allowing casting this class to Input and
+      // Output.
+      absl::StrAppend(&class_decl,
+                      "  operator ::tensorflow::Output() const { return ",
+                      op_info.output_names[0], "; }\n");
+      absl::StrAppend(&class_decl,
+                      "  operator ::tensorflow::Input() const { return ",
+                      op_info.output_names[0], "; }\n");
+      // Write node() to get the Node* directly.
+      absl::StrAppend(&class_decl,
+                      "  ::tensorflow::Node* node() const { return ",
+                      op_info.output_names[0], ".node(); }\n");
+    }
+  }
+  // Add the static functions to set optional attrs
+  if (op_info.has_optional_attrs) {
+    absl::StrAppend(&class_decl, "\n");
+    for (int i = 0; i < op_info.graph_op_def.attr_size(); ++i) {
+      const auto& attr(op_info.graph_op_def.attr(i));
+      const auto& api_def_attr(op_info.api_def.attr(i));
+      if ((op_info.inferred_input_attrs.find(attr.name()) !=
+           op_info.inferred_input_attrs.end()) ||
+          !api_def_attr.has_default_value()) {
+        continue;
+      }
+      const auto entry = AttrTypeName(attr.type());
+      const auto attr_type_name = entry.first;
+      const bool use_const = entry.second;
+      const std::string camel_case_name = ToCamelCase(api_def_attr.rename_to());
+      const std::string suffix =
+          (camel_case_name == op_info.op_name || camel_case_name == "Attrs")
+              ? "_"
+              : "";
+      const std::string attr_func_def = strings::StrCat(
+          camel_case_name, suffix, "(", use_const ? "const " : "",
+          attr_type_name, use_const ? "&" : "");
+      absl::StrAppend(&class_decl, "  static Attrs ", attr_func_def, " x) {\n");
+      absl::StrAppend(&class_decl, "    return Attrs().", camel_case_name,
+                      suffix, "(x);\n");
+      absl::StrAppend(&class_decl, "  }\n");
+    }
+  }
+
+  absl::StrAppend(&class_decl, "\n  Operation operation;\n");
+  for (int i = 0; i < op_info.output_types.size(); ++i) {
+    strings::StrAppend(&class_decl, "  ", op_info.output_types[i], " ",
+                       op_info.output_names[i], ";\n");
+  }
+
+  absl::StrAppend(&class_decl, "};\n");
+  if (!op_info.aliases.empty()) {
+    for (const auto& alias : op_info.aliases) {
+      strings::StrAppend(&class_decl, "typedef ", op_info.op_name, " ", alias,
+                         ";\n");
+    }
+  }
+  absl::StrAppend(&class_decl, "\n");
+  TF_CHECK_OK(h->Append(class_decl));
+}
+
+void GetOutput(const OpInfo& op_info, std::string* out) {
+  const std::string scope_str = op_info.arg_names[0];
+  std::string return_on_error =
+      absl::StrCat("if (!", scope_str, ".ok()) return;");
+
+  absl::StrAppend(out, "  this->operation = Operation(ret);\n");
+
+  // No outputs.
+  if (op_info.graph_op_def.output_arg_size() == 0) {
+    absl::StrAppend(out, "  return;\n");
+    return;
+  }
+  if (op_info.graph_op_def.output_arg_size() == 1) {
+    // One output, no need for NameRangeMap
+    if (op_info.is_list_output[0]) {
+      absl::StrAppend(out,
+                      "  for (int32 i = 0; i < ret->num_outputs(); ++i)\n");
+      absl::StrAppend(out, "    this->", op_info.output_names[0],
+                      ".push_back(Output(ret, i));\n");
+    } else {
+      absl::StrAppend(out, "  this->", op_info.output_names[0],
+                      " = Output(ret, 0);\n");
+    }
+    return;
+  }
+  absl::StrAppend(out, "  ::tensorflow::NameRangeMap _outputs_range;\n");
+  absl::StrAppend(out,
+                  "  ::tensorflow::Status _status_ = "
+                  "::tensorflow::NameRangesForNode(*ret, ret->op_def(), "
+                  "nullptr, &_outputs_range);\n");
+  strings::StrAppend(out, "  if (!_status_.ok()) {\n", "    ", scope_str,
+                     ".UpdateStatus(_status_);\n", "    return;\n");
+  absl::StrAppend(out, "  }\n\n");
+
+  for (int i = 0; i < op_info.graph_op_def.output_arg_size(); ++i) {
+    const std::string arg_range = absl::StrCat(
+        "_outputs_range[\"", op_info.graph_op_def.output_arg(i).name(), "\"]");
+    if (op_info.is_list_output[i]) {
+      strings::StrAppend(out, "  for (int32 i = ", arg_range, ".first; i < ",
+                         arg_range, ".second; ++i)\n");
+      absl::StrAppend(out, "    this->", op_info.output_names[i],
+                      ".push_back(Output(ret, i));\n");
+    } else {
+      strings::StrAppend(out, "  this->", op_info.output_names[i],
+                         " = Output(ret, ", arg_range, ".first);\n");
+    }
+  }
+}
+
+std::string GetConstructorBody(const OpInfo& op_info) {
+  const std::string scope_str = op_info.arg_names[0];
+
+  std::string body;
+  std::string return_on_error =
+      absl::StrCat("if (!", scope_str, ".ok()) return;");
+
+  absl::StrAppend(&body, "  ", return_on_error, "\n");
+
+  for (int i = 0; i < op_info.graph_op_def.input_arg_size(); ++i) {
+    const auto& arg(op_info.graph_op_def.input_arg(i));
+    const auto& api_def_arg(op_info.api_def.in_arg(i));
+    strings::StrAppend(
+        &body, "  auto _", api_def_arg.rename_to(), " = ::tensorflow::ops::",
+        ArgIsList(arg) ? "AsNodeOutList" : "AsNodeOut", "(", scope_str, ", ",
+        AvoidCPPKeywords(api_def_arg.rename_to()), ");\n");
+    absl::StrAppend(&body, "  ", return_on_error, "\n");
+  }
+
+  absl::StrAppend(&body, "  ::tensorflow::Node* ret;\n");
+  strings::StrAppend(&body, "  const auto unique_name = ", scope_str,
+                     ".GetUniqueNameForOp(\"", op_info.op_name, "\");\n");
+  absl::StrAppend(&body,
+                  "  auto builder = ::tensorflow::NodeBuilder(unique_name, \"",
+                  op_info.graph_op_def.name(), "\")\n");
+  const std::string spaces = "                     ";
+  for (int i = 0; i < op_info.api_def.in_arg_size(); ++i) {
+    const auto& arg(op_info.api_def.in_arg(i));
+    absl::StrAppend(&body, spaces, ".Input(_", arg.rename_to(), ")\n");
+  }
+  for (int i = 0; i < op_info.api_def.attr_size(); ++i) {
+    const auto& graph_attr(op_info.graph_op_def.attr(i));
+    const auto& api_def_attr(op_info.api_def.attr(i));
+    if (op_info.inferred_input_attrs.find(api_def_attr.name()) !=
+        op_info.inferred_input_attrs.end()) {
+      continue;
+    }
+    const std::string attr_name =
+        api_def_attr.has_default_value()
+            ? absl::StrCat("attrs.", api_def_attr.rename_to(), "_")
+            : AvoidCPPKeywords(api_def_attr.rename_to());
+    strings::StrAppend(&body, spaces, ".Attr(\"", graph_attr.name(), "\", ",
+                       attr_name, ")\n");
+  }
+  absl::StrAppend(&body, "  ;\n");
+  absl::StrAppend(&body, "  ", scope_str, ".UpdateBuilder(&builder);\n");
+  strings::StrAppend(&body, "  ", scope_str, ".UpdateStatus(builder.Finalize(",
+                     scope_str, ".graph(), &ret));\n");
+  absl::StrAppend(&body, "  ", return_on_error, "\n");
+  strings::StrAppend(&body, "  ", scope_str, ".UpdateStatus(", scope_str,
+                     ".DoShapeInference(ret));\n");
+
+  GetOutput(op_info, &body);
+  return body;
+}
+
+void WriteClassDef(const OpInfo& op_info, WritableFile* cc) {
+  std::string class_def;
+  absl::StrAppend(
+      &class_def,
+      GetConstructorDecl(op_info, absl::StrCat(op_info.op_name, "::"),
+                         /* include_attr */ true),
+      " {\n");
+  absl::StrAppend(&class_def, GetConstructorBody(op_info));
+  absl::StrAppend(&class_def, "}\n\n");
+
+  if (op_info.has_optional_attrs) {
+    absl::StrAppend(
+        &class_def,
+        GetConstructorDecl(op_info, absl::StrCat(op_info.op_name, "::"),
+                           /* include_attr */ false));
+    absl::StrAppend(&class_def, "\n  : ", op_info.op_name, "(");
+    int i = 0;
+    for (; i < op_info.arg_names.size(); ++i) {
+      if (i > 0) absl::StrAppend(&class_def, ", ");
+      absl::StrAppend(&class_def, op_info.arg_names[i]);
+    }
+    if (i > 0) absl::StrAppend(&class_def, ", ");
+    absl::StrAppend(&class_def, op_info.op_name, "::Attrs()");
+    absl::StrAppend(&class_def, ") {}\n\n");
+  }
+  TF_CHECK_OK(cc->Append(class_def));
+}
+
+void WriteCCOp(const OpDef& graph_op_def, const ApiDef& api_def,
+               const std::vector<std::string>& aliases, WritableFile* h,
+               WritableFile* cc) {
+  OpInfo op_info(graph_op_def, api_def, aliases);
+
+  WriteClassDecl(op_info, h);
+  WriteClassDef(op_info, cc);
+}
+
+void StartFiles(bool internal, const std::string& dot_h_fname, WritableFile* h,
+                WritableFile* cc, std::string* op_header_guard) {
+  const std::string header =
+      R"header(// This file is MACHINE GENERATED! Do not edit.
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+)header";
+
+  // TODO(keveman): Make namespaces configurable.
+  const std::string namespace_begin = internal ? R"namespace(
+namespace tensorflow {
+namespace ops {
+namespace internal {
+// NOTE: This namespace has internal TensorFlow details that
+// are not part of TensorFlow's public API.
+
+)namespace"
+                                               : R"namespace(
+namespace tensorflow {
+namespace ops {
+
+)namespace";
+
+  const std::string op_header = GetPath(dot_h_fname);
+  *op_header_guard = ToGuard(op_header);
+  const std::string cc_header = strings::StrCat(
+      R"include(// This file is MACHINE GENERATED! Do not edit.
+
+
+#include "tensorflow/cc/ops/const_op.h"
+)include",
+      "#include \"", op_header, "\"\n", namespace_begin);
+
+  const std::string filename = GetFilename(dot_h_fname);
+  const std::string doxygen = strings::StrCat(
+      "/// @defgroup ", filename, " ", ToTitle(filename), "\n", "/// @{\n\n");
+
+  TF_CHECK_OK(h->Append(
+      strings::StrCat("// This file is MACHINE GENERATED! Do not edit.\n\n"
+                      "#ifndef ",
+                      *op_header_guard,
+                      "\n"
+                      "#define ",
+                      *op_header_guard, "\n\n")));
+  TF_CHECK_OK(h->Append(header));
+  TF_CHECK_OK(h->Append(namespace_begin));
+  TF_CHECK_OK(h->Append(doxygen));
+  TF_CHECK_OK(cc->Append(cc_header));
+}
+
+void FinishFiles(bool internal, WritableFile* h, WritableFile* cc,
+                 const std::string& op_header_guard) {
+  const std::string footer = internal ? R"footer(}  // namespace internal
+}  // namespace ops
+}  // namespace tensorflow
+)footer"
+                                      :
+                                      R"footer(/// @}
+
+}  // namespace ops
+}  // namespace tensorflow
+)footer";
+
+  TF_CHECK_OK(h->Append(footer));
+  TF_CHECK_OK(
+      h->Append(absl::StrCat("\n#endif  ", "// ", op_header_guard, "\n")));
+  TF_CHECK_OK(cc->Append(footer));
+
+  TF_CHECK_OK(cc->Close());
+  TF_CHECK_OK(h->Close());
+}
+
+std::string MakeInternal(const std::string& fname) {
+  auto dot_pos = fname.rfind('.');
+  if (dot_pos == std::string::npos) {
+    return absl::StrCat(fname, "_internal");
+  } else {
+    return absl::StrCat(fname.substr(0, dot_pos), "_internal",
+                        fname.substr(dot_pos));
+  }
+}
+
+}  // namespace
+
+void WriteCCOps(const OpList& ops, const ApiDefMap& api_def_map,
+                const std::string& dot_h_fname,
+                const std::string& dot_cc_fname) {
+  Env* env = Env::Default();
+
+  // Write the initial boilerplate to the .h and .cc files.
+  std::unique_ptr<WritableFile> h = nullptr;
+  std::unique_ptr<WritableFile> cc = nullptr;
+  TF_CHECK_OK(env->NewWritableFile(dot_h_fname, &h));
+  TF_CHECK_OK(env->NewWritableFile(dot_cc_fname, &cc));
+  std::string op_header_guard;
+  StartFiles(false, dot_h_fname, h.get(), cc.get(), &op_header_guard);
+
+  // Create the internal versions of these files for the hidden ops.
+  std::unique_ptr<WritableFile> internal_h = nullptr;
+  std::unique_ptr<WritableFile> internal_cc = nullptr;
+  const std::string internal_dot_h_fname = MakeInternal(dot_h_fname);
+  TF_CHECK_OK(env->NewWritableFile(internal_dot_h_fname, &internal_h));
+  TF_CHECK_OK(env->NewWritableFile(MakeInternal(dot_cc_fname), &internal_cc));
+  std::string internal_op_header_guard;
+  StartFiles(true /* internal */, internal_dot_h_fname, internal_h.get(),
+             internal_cc.get(), &internal_op_header_guard);
+
+  for (const auto& graph_op_def : ops.op()) {
+    // Skip deprecated ops.
+    // TODO(josh11b): If needed, can put them into a "deprecated" namespace
+    // instead of skipping.
+    if (graph_op_def.has_deprecation() &&
+        graph_op_def.deprecation().version() <= TF_GRAPH_DEF_VERSION) {
+      continue;
+    }
+
+    // We use a hand-written wrapper for "Const", since the generated
+    // code depends on it.
+    if (graph_op_def.name() == "Const") continue;
+
+    const auto* api_def = api_def_map.GetApiDef(graph_op_def.name());
+
+    std::vector<std::string> aliases;
+    if (api_def->visibility() == ApiDef::SKIP) continue;
+    // First endpoint is canonical, the rest are aliases.
+    for (int endpoint_i = 1; endpoint_i < api_def->endpoint_size();
+         ++endpoint_i) {
+      aliases.push_back(api_def->endpoint(endpoint_i).name());
+    }
+    if (api_def->visibility() == ApiDef::HIDDEN) {
+      // Write hidden ops to _internal.h and _internal.cc.
+      WriteCCOp(graph_op_def, *api_def, aliases, internal_h.get(),
+                internal_cc.get());
+      continue;
+    }
+    // This isn't a hidden op, write it to the main files.
+    WriteCCOp(graph_op_def, *api_def, aliases, h.get(), cc.get());
+  }
+
+  FinishFiles(false, h.get(), cc.get(), op_header_guard);
+  FinishFiles(true /* internal */, internal_h.get(), internal_cc.get(),
+              internal_op_header_guard);
+}
+
+}  // namespace cc_op
+}  // namespace tensorflow

tensorflow/cc/framework/cc_op_gen.h

@@ -0,0 +1,35 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_
+#define TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_
+
+#include <string>
+
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace cc_op {
+/// Result is written to files dot_h and dot_cc.
+void WriteCCOps(const OpList& ops, const ApiDefMap& api_def_map,
+                const std::string& dot_h_fname,
+                const std::string& dot_cc_fname);
+
+}  // namespace cc_op
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_H_

tensorflow/cc/framework/cc_op_gen_main.cc

@@ -0,0 +1,70 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/cc/framework/cc_op_gen.h"
+#include "tensorflow/cc/framework/cc_op_gen_util.h"
+#include "xla/tsl/platform/status.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace cc_op {
+namespace {
+
+void PrintAllCCOps(const std::string& dot_h, const std::string& dot_cc,
+                   bool include_internal,
+                   const std::vector<std::string>& api_def_dirs) {
+  OpList ops;
+  absl::StatusOr<ApiDefMap> api_def_map =
+      LoadOpsAndApiDefs(ops, include_internal, api_def_dirs);
+  TF_CHECK_OK(api_def_map.status());
+  api_def_map->UpdateDocs();
+  WriteCCOps(ops, *api_def_map, dot_h, dot_cc);
+}
+
+}  // namespace
+}  // namespace cc_op
+}  // namespace tensorflow
+
+int main(int argc, char* argv[]) {
+  tensorflow::port::InitMain(argv[0], &argc, &argv);
+  if (argc != 5) {
+    for (int i = 1; i < argc; ++i) {
+      fprintf(stderr, "Arg %d = %s\n", i, argv[i]);
+    }
+    fprintf(stderr,
+            "Usage: %s out.h out.cc include_internal "
+            "api_def_dirs1,api_def_dir2 ...\n"
+            "  include_internal: 1 means include internal ops\n",
+            argv[0]);
+    exit(1);
+  }
+
+  bool include_internal = absl::string_view("1") == argv[3];
+  std::vector<std::string> api_def_dirs = tensorflow::str_util::Split(
+      argv[4], ",", tensorflow::str_util::SkipEmpty());
+  tensorflow::cc_op::PrintAllCCOps(argv[1], argv[2], include_internal,
+                                   api_def_dirs);
+  return 0;
+}

tensorflow/cc/framework/cc_op_gen_test.cc

@@ -0,0 +1,192 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/cc_op_gen.h"
+
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+constexpr char kBaseOpDef[] = R"(
+op {
+  name: "Foo"
+  input_arg {
+    name: "images"
+    description: "Images to process."
+  }
+  input_arg {
+    name: "dim"
+    description: "Description for dim."
+    type: DT_FLOAT
+  }
+  output_arg {
+    name: "output"
+    description: "Description for output."
+    type: DT_FLOAT
+  }
+  attr {
+    name: "T"
+    type: "type"
+    description: "Type for images"
+    allowed_values {
+      list {
+        type: DT_UINT8
+        type: DT_INT8
+      }
+    }
+    default_value {
+      i: 1
+    }
+  }
+  summary: "Summary for op Foo."
+  description: "Description for op Foo."
+}
+)";
+
+void ExpectHasSubstr(absl::string_view s, absl::string_view expected) {
+  EXPECT_TRUE(absl::StrContains(s, expected))
+      << "'" << s << "' does not contain '" << expected << "'";
+}
+
+void ExpectDoesNotHaveSubstr(absl::string_view s, absl::string_view expected) {
+  EXPECT_FALSE(absl::StrContains(s, expected))
+      << "'" << s << "' contains '" << expected << "'";
+}
+
+void ExpectSubstrOrder(const std::string& s, const std::string& before,
+                       const std::string& after) {
+  int before_pos = s.find(before);
+  int after_pos = s.find(after);
+  ASSERT_NE(std::string::npos, before_pos);
+  ASSERT_NE(std::string::npos, after_pos);
+  EXPECT_LT(before_pos, after_pos)
+      << before << " is not before " << after << " in " << s;
+}
+
+// Runs WriteCCOps and stores output in (internal_)cc_file_path and
+// (internal_)h_file_path.
+void GenerateCcOpFiles(Env* env, const OpList& ops,
+                       const ApiDefMap& api_def_map, std::string* h_file_text,
+                       std::string* internal_h_file_text) {
+  const std::string& tmpdir = testing::TmpDir();
+
+  const auto h_file_path = io::JoinPath(tmpdir, "test.h");
+  const auto cc_file_path = io::JoinPath(tmpdir, "test.cc");
+  const auto internal_h_file_path = io::JoinPath(tmpdir, "test_internal.h");
+  const auto internal_cc_file_path = io::JoinPath(tmpdir, "test_internal.cc");
+
+  cc_op::WriteCCOps(ops, api_def_map, h_file_path, cc_file_path);
+
+  TF_ASSERT_OK(ReadFileToString(env, h_file_path, h_file_text));
+  TF_ASSERT_OK(
+      ReadFileToString(env, internal_h_file_path, internal_h_file_text));
+}
+
+TEST(CcOpGenTest, TestVisibilityChangedToHidden) {
+  const std::string api_def = R"(
+op {
+  graph_op_name: "Foo"
+  visibility: HIDDEN
+}
+)";
+  Env* env = Env::Default();
+  OpList op_defs;
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+  ApiDefMap api_def_map(op_defs);
+
+  std::string h_file_text, internal_h_file_text;
+  // Without ApiDef
+  GenerateCcOpFiles(env, op_defs, api_def_map, &h_file_text,
+                    &internal_h_file_text);
+  ExpectHasSubstr(h_file_text, "class Foo");
+  ExpectDoesNotHaveSubstr(internal_h_file_text, "class Foo");
+
+  // With ApiDef
+  TF_ASSERT_OK(api_def_map.LoadApiDef(api_def));
+  GenerateCcOpFiles(env, op_defs, api_def_map, &h_file_text,
+                    &internal_h_file_text);
+  ExpectHasSubstr(internal_h_file_text, "class Foo");
+  ExpectDoesNotHaveSubstr(h_file_text, "class Foo");
+}
+
+TEST(CcOpGenTest, TestArgNameChanges) {
+  const std::string api_def = R"(
+op {
+  graph_op_name: "Foo"
+  arg_order: "dim"
+  arg_order: "images"
+}
+)";
+  Env* env = Env::Default();
+  OpList op_defs;
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+
+  ApiDefMap api_def_map(op_defs);
+  std::string cc_file_text, h_file_text;
+  std::string internal_cc_file_text, internal_h_file_text;
+  // Without ApiDef
+  GenerateCcOpFiles(env, op_defs, api_def_map, &h_file_text,
+                    &internal_h_file_text);
+  ExpectSubstrOrder(h_file_text, "Input images", "Input dim");
+
+  // With ApiDef
+  TF_ASSERT_OK(api_def_map.LoadApiDef(api_def));
+  GenerateCcOpFiles(env, op_defs, api_def_map, &h_file_text,
+                    &internal_h_file_text);
+  ExpectSubstrOrder(h_file_text, "Input dim", "Input images");
+}
+
+TEST(CcOpGenTest, TestEndpoints) {
+  const std::string api_def = R"(
+op {
+  graph_op_name: "Foo"
+  endpoint {
+    name: "Foo1"
+  }
+  endpoint {
+    name: "Foo2"
+  }
+}
+)";
+  Env* env = Env::Default();
+  OpList op_defs;
+  protobuf::TextFormat::ParseFromString(kBaseOpDef, &op_defs);  // NOLINT
+
+  ApiDefMap api_def_map(op_defs);
+  std::string cc_file_text, h_file_text;
+  std::string internal_cc_file_text, internal_h_file_text;
+  // Without ApiDef
+  GenerateCcOpFiles(env, op_defs, api_def_map, &h_file_text,
+                    &internal_h_file_text);
+  ExpectHasSubstr(h_file_text, "class Foo {");
+  ExpectDoesNotHaveSubstr(h_file_text, "class Foo1");
+  ExpectDoesNotHaveSubstr(h_file_text, "class Foo2");
+
+  // With ApiDef
+  TF_ASSERT_OK(api_def_map.LoadApiDef(api_def));
+  GenerateCcOpFiles(env, op_defs, api_def_map, &h_file_text,
+                    &internal_h_file_text);
+  ExpectHasSubstr(h_file_text, "class Foo1");
+  ExpectHasSubstr(h_file_text, "typedef Foo1 Foo2");
+  ExpectDoesNotHaveSubstr(h_file_text, "class Foo {");
+}
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/framework/cc_op_gen_util.cc

@@ -0,0 +1,764 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/cc_op_gen_util.h"
+
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/escaping.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/hash.h"
+#include "tensorflow/core/platform/numbers.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/str_util.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace cc_op {
+
+absl::StatusOr<ApiDefMap> LoadOpsAndApiDefs(
+    OpList& ops, bool include_internal,
+    const std::vector<std::string>& api_def_dirs) {
+  OpRegistry::Global()->Export(include_internal, &ops);
+  ApiDefMap api_def_map(ops);
+  if (!api_def_dirs.empty()) {
+    Env* env = Env::Default();
+    // Only load files that correspond to "ops".
+    for (const auto& op : ops.op()) {
+      for (const auto& api_def_dir : api_def_dirs) {
+        const std::string api_def_file_pattern =
+            io::JoinPath(api_def_dir, "api_def_" + op.name() + ".pbtxt");
+        if (env->FileExists(api_def_file_pattern).ok()) {
+          auto status = api_def_map.LoadFile(env, api_def_file_pattern);
+          if (!status.ok()) return status;
+        }
+      }
+    }
+  }
+  return api_def_map;
+}
+
+std::string GetPath(absl::string_view dot_h_fname) {
+  auto pos = dot_h_fname.find("/bin/");
+  std::string result(dot_h_fname);
+  if (pos != std::string::npos) {
+    // - 1 account for the terminating null character (\0) in "/genfiles/".
+    result = dot_h_fname.substr(pos + sizeof("/bin/") - 1);
+  } else {
+    pos = dot_h_fname.find("/genfiles/");
+    if (pos != std::string::npos) {
+      result = dot_h_fname.substr(pos + sizeof("/genfiles/") - 1);
+    }
+  }
+  if (result.size() > sizeof("external/") &&
+      result.compare(0, sizeof("external/") - 1, "external/") == 0) {
+    result = result.substr(sizeof("external/") - 1);
+    pos = result.find('/');
+    if (pos != std::string::npos) {
+      result = result.substr(pos + 1);
+    }
+  }
+  return result;
+}
+
+std::string GetFilename(absl::string_view path) {
+  size_t slash_pos = path.rfind('/');
+  if (slash_pos == path.npos) slash_pos = -1;
+  size_t dot_pos = path.rfind('.');
+  return std::string(path.substr(slash_pos + 1, dot_pos - (slash_pos + 1)));
+}
+
+std::string ToGuard(absl::string_view path) {
+  std::string guard;
+  guard.reserve(path.size() + 1);  // + 1 -> trailing _
+  for (const char c : path) {
+    if (absl::ascii_isupper(c)) {
+      guard += c;
+    } else if (absl::ascii_islower(c)) {
+      guard += absl::ascii_toupper(c);
+    } else {
+      guard += '_';
+    }
+  }
+  guard += '_';
+  return guard;
+}
+
+std::string ToTitle(absl::string_view name) {
+  std::string title(name);
+  for (int i = 0; i < title.size(); ++i) {
+    if (title[i] == '_') title[i] = ' ';
+  }
+  str_util::TitlecaseString(&title, " ");
+  return title;
+}
+
+std::string MakeComment(absl::string_view text, absl::string_view indent) {
+  std::string ret;
+  while (!text.empty()) {
+    int last_non_space = -1;
+    int newline;
+    for (newline = 0; newline < static_cast<int>(text.size()); ++newline) {
+      if (text[newline] == '\n') break;
+      if (text[newline] != ' ') last_non_space = newline;
+    }
+    if (last_non_space == -1) {
+      absl::StrAppend(&ret, indent, "///\n");
+    } else {
+      absl::StrAppend(&ret, indent, "/// ", text.substr(0, last_non_space + 1),
+                      "\n");
+    }
+    text.remove_prefix(newline + 1);
+  }
+  return ret;
+}
+
+std::string PrintString(absl::string_view str) {
+  return absl::StrCat("\"", absl::CEscape(str), "\"");
+}
+
+std::string PrintTensorShape(const TensorShapeProto& shape_proto) {
+  PartialTensorShape shape(shape_proto);
+  if (shape.IsIdenticalTo(PartialTensorShape())) {
+    return "::tensorflow::PartialTensorShape() /* unknown */";
+  }
+  std::string ret = "{";
+  for (int d = 0; d < shape.dims(); ++d) {
+    if (d > 0) absl::StrAppend(&ret, ", ");
+    absl::StrAppend(&ret, shape.dim_size(d));
+  }
+  absl::StrAppend(&ret, "}");
+  return ret;
+}
+
+std::string PrintTensor(const TensorProto& tensor_proto) {
+  Tensor t(tensor_proto.dtype());
+  CHECK(t.FromProto(tensor_proto));
+  const int64_t num_elts = t.NumElements();
+  switch (t.dtype()) {
+    case DT_FLOAT:
+      return PrintArray(num_elts, t.flat<float>().data());
+    case DT_DOUBLE:
+      return PrintArray(num_elts, t.flat<double>().data());
+    case DT_INT32:
+      return PrintArray(num_elts, t.flat<int32_t>().data());
+    case DT_UINT8:
+    case DT_QUINT8:
+      return PrintArray(num_elts, t.flat<uint8_t>().data());
+    case DT_UINT16:
+    case DT_QUINT16:
+      return PrintArray(num_elts, t.flat<uint16_t>().data());
+    case DT_INT16:
+    case DT_QINT16:
+      return PrintArray(num_elts, t.flat<int16_t>().data());
+    case DT_INT8:
+    case DT_QINT8:
+      return PrintArray(num_elts, t.flat<int8_t>().data());
+    case DT_INT64:
+      return PrintArray(num_elts, t.flat<int64_t>().data());
+    case DT_BOOL:
+      return PrintArray(num_elts, t.flat<bool>().data());
+    case DT_STRING: {
+      std::string ret;
+      for (int64_t i = 0; i < num_elts; ++i) {
+        if (i > 0) absl::StrAppend(&ret, " ");
+        absl::StrAppend(&ret, absl::CEscape(t.flat<tstring>()(i)));
+      }
+      return ret;
+    }
+    default: {
+      LOG(FATAL) << "Not handling type " << DataType_Name(t.dtype());
+      return std::string();
+    }
+  }
+}
+
+std::string PrintTensorProto(const TensorProto& proto) {
+  return absl::StrCat("Input::Initializer(", "{", PrintTensor(proto), "}, ",
+                      PrintTensorShape(proto.tensor_shape()),
+                      ").AsTensorProto()");
+}
+
+std::string PrintAttrValue(const std::string& op, const AttrValue& attr_value) {
+  switch (attr_value.value_case()) {
+    case AttrValue::kS:
+      return PrintString(attr_value.s());
+    case AttrValue::kI:
+      return absl::StrCat(attr_value.i());
+    case AttrValue::kF: {
+      const float f = attr_value.f();
+      if (std::isinf(f)) {
+        return absl::StrCat(f < 0.0f ? "-" : "+",
+                            "std::numeric_limits<float>::infinity()");
+      } else {
+        return absl::StrCat(strings::LegacyPrecision(attr_value.f()),
+                            floorf(f) == f ? ".0" : "", "f");
+      }
+    }
+    case AttrValue::kB:
+      return attr_value.b() ? "true" : "false";
+    case AttrValue::kType:
+      return DataType_Name(attr_value.type());
+    case AttrValue::kShape:
+      return PrintTensorShape(attr_value.shape());
+    case AttrValue::kTensor:
+      return PrintTensorProto(attr_value.tensor());
+    case AttrValue::kList: {
+      std::string ret = "{";
+      if (attr_value.list().s_size() > 0) {
+        for (int i = 0; i < attr_value.list().s_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          absl::StrAppend(&ret, PrintString(attr_value.list().s(i)));
+        }
+      } else if (attr_value.list().i_size() > 0) {
+        for (int i = 0; i < attr_value.list().i_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          absl::StrAppend(&ret, attr_value.list().i(i));
+        }
+      } else if (attr_value.list().f_size() > 0) {
+        for (int i = 0; i < attr_value.list().f_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          const float f = attr_value.list().f(i);
+          absl::StrAppend(&ret, strings::LegacyPrecision(f),
+                          floorf(f) == f ? ".0" : "", "f");
+        }
+      } else if (attr_value.list().b_size() > 0) {
+        for (int i = 0; i < attr_value.list().b_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          absl::StrAppend(&ret, attr_value.list().b(i) ? "true" : "false");
+        }
+      } else if (attr_value.list().type_size() > 0) {
+        for (int i = 0; i < attr_value.list().type_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          absl::StrAppend(&ret, DataType_Name(attr_value.list().type(i)));
+        }
+      } else if (attr_value.list().shape_size() > 0) {
+        for (int i = 0; i < attr_value.list().shape_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          absl::StrAppend(&ret, PrintTensorShape(attr_value.list().shape(i)));
+        }
+      } else if (attr_value.list().tensor_size() > 0) {
+        for (int i = 0; i < attr_value.list().tensor_size(); ++i) {
+          if (i > 0) absl::StrAppend(&ret, ", ");
+          absl::StrAppend(&ret, PrintTensorProto(attr_value.list().tensor(i)));
+        }
+      }
+      absl::StrAppend(&ret, "}");
+      return ret;
+    }
+    default:
+      LOG(FATAL) << "Unsupported Attr type: " << op << " "
+                 << attr_value.value_case();
+  }
+  return "<Unknown AttrValue type>";  // Prevent missing return warning
+}
+
+bool IsEmptyList(const AttrValue::ListValue& list) {
+  return list.s_size() == 0 && list.i_size() == 0 && list.f_size() == 0 &&
+         list.b_size() == 0 && list.type_size() == 0 &&
+         list.shape_size() == 0 && list.tensor_size() == 0;
+}
+
+std::string ToCamelCase(absl::string_view str) {
+  std::string result;
+  const char joiner = '_';
+  size_t i = 0;
+  bool cap = true;
+  while (i < str.size()) {
+    const char c = str[i++];
+    if (c == '>') {
+      cap = true;
+    } else if (c == joiner) {
+      cap = true;
+    } else if (cap) {
+      result += absl::ascii_toupper(c);
+      cap = false;
+    } else {
+      result += c;
+    }
+  }
+  return result;
+}
+
+std::string SeparateNamespaces(absl::string_view str) {
+  std::string result;
+  const char joiner = '_';
+  size_t i = 0;
+  while (i < str.size()) {
+    const char c = str[i++];
+    if (c == '>') {
+      result += joiner;
+    } else {
+      result += c;
+    }
+  }
+  return result;
+}
+
+std::pair<absl::string_view, bool> AttrTypeName(absl::string_view attr_type) {
+  static const auto* attr_type_map = new std::unordered_map<
+      absl::string_view, std::pair<absl::string_view, bool>, StringPieceHasher>{
+      {"string", {"StringPiece", false}},
+      {"list(string)", {"gtl::ArraySlice<::tensorflow::tstring>", true}},
+      {"int", {"int64", false}},
+      {"list(int)", {"gtl::ArraySlice<int>", true}},
+      {"float", {"float", false}},
+      {"list(float)", {"gtl::ArraySlice<float>", true}},
+      {"bool", {"bool", false}},
+      {"list(bool)", {"gtl::ArraySlice<bool>", true}},
+      {"type", {"DataType", false}},
+      {"list(type)", {"DataTypeSlice", true}},
+      {"shape", {"PartialTensorShape", false}},
+      {"list(shape)", {"gtl::ArraySlice<PartialTensorShape>", true}},
+      {"tensor", {"TensorProto", true}},
+      {"list(tensor)", {"gtl::ArraySlice<TensorProto>", true}},
+      {"func", {"NameAttrList", true}},
+      {"list(func)", {"gtl::ArraySlice<NameAttrList>", true}},
+  };
+
+  auto entry = attr_type_map->find(attr_type);
+  if (entry == attr_type_map->end()) {
+    LOG(FATAL) << "Unsupported Attr type: " << attr_type;
+    return {"", false};
+  }
+  return entry->second;
+}
+
+absl::string_view ListElementTypeName(absl::string_view attr_type) {
+  static const auto* attr_list_type_map = new absl::flat_hash_map<
+      absl::string_view, absl::string_view, StringPieceHasher>{
+      {"list(string)", "string"},      {"list(int)", "int"},
+      {"list(float)", "float"},        {"list(bool)", "bool"},
+      {"list(type)", "DataType"},      {"list(shape)", "PartialTensorShape"},
+      {"list(tensor)", "TensorProto"},
+  };
+
+  auto entry = attr_list_type_map->find(attr_type);
+  if (entry == attr_list_type_map->end()) {
+    LOG(FATAL) << "Unsupported or non-list Attr type: " << attr_type;
+    return "";
+  }
+  return entry->second;
+}
+
+bool IsCPPKeyword(absl::string_view name) {
+  static const absl::flat_hash_set<absl::string_view, StringPieceHasher>*
+      // Keywords obtained from http://en.cppreference.com/w/cpp/keyword
+      kCPPReserved = new absl::flat_hash_set<absl::string_view,
+                                             StringPieceHasher>{
+          "alignas",
+          "alignof",
+          "and",
+          "and_eq",
+          "asm",
+          "atomic_cancel",
+          "atomic_commit",
+          "atomic_noexcept",
+          "auto",
+          "bitand",
+          "bitor",
+          "bool",
+          "break",
+          "case",
+          "catch",
+          "char",
+          "char16_t",
+          "char32_t",
+          "class",
+          "compl",
+          "concept",
+          "const",
+          "const_cast",
+          "constexpr",
+          "continue",
+          "decltype",
+          "default",
+          "delete",
+          "do",
+          "double",
+          "dynamic_cast",
+          "else",
+          "enum",
+          "explicit",
+          "export",
+          "extern",
+          "false",
+          "final",
+          "float",
+          "for",
+          "friend",
+          "goto",
+          "if",
+          "import",
+          "inline",
+          "int",
+          "long",
+          "module",
+          "mutable",
+          "namespace",
+          "new",
+          "noexcept",
+          "not",
+          "not_eq",
+          "nullptr",
+          "operator",
+          "or",
+          "or_eq",
+          "override",
+          "private",
+          "protected",
+          "public",
+          "register",
+          "reinterpret_cast",
+          "requires",
+          "return",
+          "short",
+          "signed",
+          "sizeof",
+          "static",
+          "static_assert",
+          "static_cast",
+          "struct",
+          "switch",
+          "synchronized",
+          "template",
+          "this",
+          "thread_local",
+          "throw",
+          "true",
+          "try",
+          "typedef",
+          "typeid",
+          "typename",
+          "union",
+          "unsigned",
+          "using",
+          "virtual",
+          "void",
+          "volatile",
+          "wchar_t",
+          "while",
+          "xor",
+          "xor_eq",
+
+          // The following are not C++ keywords, but names of local variables
+          // and parameters used in the op constructor. Treating them as
+          // keywords, so that other parameter names don't conflict with these.
+          "builder",
+          "node",
+          "ret",
+          "scope",
+          "unique_name",
+      };
+  return kCPPReserved->count(name) > 0;
+}
+
+std::string AvoidCPPKeywords(absl::string_view name) {
+  if (IsCPPKeyword(name)) {
+    return absl::StrCat(name, "_");
+  }
+  return std::string(name);
+}
+
+void InferArgAttributes(
+    const OpDef::ArgDef& arg,
+    std::unordered_map<std::string, std::string>* inferred_attrs) {
+  if (!arg.type_attr().empty()) {
+    gtl::InsertIfNotPresent(inferred_attrs, arg.type_attr(), arg.name());
+  } else if (!arg.type_list_attr().empty()) {
+    gtl::InsertIfNotPresent(inferred_attrs, arg.type_list_attr(), arg.name());
+  }
+  if (!arg.number_attr().empty()) {
+    gtl::InsertIfNotPresent(inferred_attrs, arg.number_attr(), arg.name());
+  }
+}
+
+void InferOpAttributes(
+    const OpDef& op_def,
+    std::unordered_map<std::string, std::string>* inferred_input_attrs) {
+  for (int i = 0; i < op_def.input_arg_size(); ++i) {
+    const auto& arg(op_def.input_arg(i));
+    InferArgAttributes(arg, inferred_input_attrs);
+  }
+}
+
+bool ArgIsList(const OpDef::ArgDef& arg) {
+  return !arg.type_list_attr().empty() || !arg.number_attr().empty();
+}
+
+bool HasOptionalAttrs(
+    const ApiDef& api_def,
+    const std::unordered_map<std::string, std::string>& inferred_input_attrs) {
+  for (int i = 0; i < api_def.attr_size(); ++i) {
+    const auto& attr(api_def.attr(i));
+    if ((inferred_input_attrs.find(attr.name()) ==
+         inferred_input_attrs.end()) &&
+        attr.has_default_value()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def,
+               const std::vector<std::string>& aliases)
+    : graph_op_def(graph_op_def), api_def(api_def), aliases(aliases) {
+  op_name = SeparateNamespaces(api_def.endpoint(0).name());
+  InferOpAttributes(graph_op_def, &inferred_input_attrs);
+  has_optional_attrs = HasOptionalAttrs(api_def, inferred_input_attrs);
+  arg_types.push_back("const ::tensorflow::Scope&");
+  arg_names.push_back("scope");
+
+  if (graph_op_def.has_deprecation()) {
+    if (!api_def.summary().empty()) {
+      comment = absl::StrCat(api_def.summary(), "\n");
+    }
+    absl::StrAppend(&comment, "DEPRECATED at GraphDef version ",
+                    graph_op_def.deprecation().version(), ":\n",
+                    graph_op_def.deprecation().explanation(), ".\n");
+  } else if (api_def.summary().empty()) {
+    comment = "TODO: add doc.\n";
+  } else {
+    comment = absl::StrCat(api_def.summary(), "\n");
+  }
+  if (!api_def.description().empty()) {
+    absl::StrAppend(&comment, "\n", api_def.description(), "\n");
+  }
+  absl::StrAppend(&comment, "\nArgs:\n* scope: A Scope object\n");
+
+  // Process inputs
+  for (int i = 0; i < api_def.arg_order_size(); ++i) {
+    const auto& arg = *FindInputArg(api_def.arg_order(i), graph_op_def);
+    const auto& api_def_arg = *FindInputArg(api_def.arg_order(i), api_def);
+    arg_types.push_back(
+        absl::StrCat("::tensorflow::", ArgIsList(arg) ? "InputList" : "Input"));
+    arg_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to()));
+
+    // TODO(keveman): Include input type information.
+    absl::string_view description = api_def_arg.description();
+    if (!description.empty()) {
+      ConsumeEquals(&description);
+      absl::StrAppend(&comment, "* ", AvoidCPPKeywords(api_def_arg.rename_to()),
+                      ": ", api_def_arg.description(), "\n");
+    }
+  }
+
+  // Process attrs
+  std::string required_attrs_comment;
+  std::string optional_attrs_comment;
+  for (int i = 0; i < graph_op_def.attr_size(); ++i) {
+    // ApiDef attributes must be in the same order as in OpDef since
+    // we initialize ApiDef based on OpDef.
+    const auto& attr(graph_op_def.attr(i));
+    const auto& api_def_attr(api_def.attr(i));
+    CHECK_EQ(attr.name(), api_def_attr.name());
+    // Skip inferred arguments
+    if (inferred_input_attrs.count(attr.name()) > 0) continue;
+
+    const auto entry = AttrTypeName(attr.type());
+    const auto attr_type_name = entry.first;
+    const bool use_const = entry.second;
+    std::string attr_name = AvoidCPPKeywords(api_def_attr.rename_to());
+
+    std::string attr_comment;
+    if (!api_def_attr.description().empty()) {
+      // TODO(keveman): Word wrap and indent this, to handle multi-line
+      // descriptions.
+      absl::StrAppend(&attr_comment, "* ", attr_name, ": ",
+                      api_def_attr.description(), "\n");
+    }
+    if (api_def_attr.has_default_value()) {
+      absl::StrAppend(&optional_attrs_comment, attr_comment);
+    } else {
+      absl::StrAppend(&required_attrs_comment, attr_comment);
+      arg_types.push_back(absl::StrCat(use_const ? "const " : "",
+                                       attr_type_name, use_const ? "&" : ""));
+      arg_names.push_back(attr_name);
+    }
+  }
+
+  absl::StrAppend(&comment, required_attrs_comment);
+
+  if (!optional_attrs_comment.empty()) {
+    absl::StrAppend(&comment, "\nOptional attributes (see `Attrs`):\n");
+    absl::StrAppend(&comment, optional_attrs_comment);
+  }
+
+  // Process outputs
+  for (int i = 0; i < graph_op_def.output_arg_size(); ++i) {
+    // ApiDef arguments must be in the same order as in OpDef since
+    // we initialize ApiDef based on OpDef.
+    const auto& arg = graph_op_def.output_arg(i);
+    const auto& api_def_arg(api_def.out_arg(i));
+    CHECK_EQ(arg.name(), api_def_arg.name());
+
+    bool is_list = ArgIsList(arg);
+    output_types.push_back(
+        absl::StrCat("::tensorflow::", is_list ? "OutputList" : "Output"));
+    output_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to()));
+    is_list_output.push_back(is_list);
+  }
+
+  absl::StrAppend(&comment, "\nReturns:\n");
+  if (graph_op_def.output_arg_size() == 0) {  // No outputs.
+    absl::StrAppend(&comment, "* the created `Operation`\n");
+  } else if (graph_op_def.output_arg_size() == 1) {  // One output
+    if (is_list_output[0]) {
+      absl::StrAppend(&comment, "* `OutputList`: ");
+    } else {
+      absl::StrAppend(&comment, "* `Output`: ");
+    }
+    if (api_def.out_arg(0).description().empty()) {
+      absl::StrAppend(&comment, "The ", api_def.out_arg(0).name(),
+                      " tensor.\n");
+    } else {
+      // TODO(josh11b): Word wrap this.
+      absl::StrAppend(&comment, api_def.out_arg(0).description(), "\n");
+    }
+  } else {  // Multiple outputs.
+    for (int i = 0; i < graph_op_def.output_arg_size(); ++i) {
+      if (is_list_output[i]) {
+        absl::StrAppend(&comment, "* `OutputList`");
+      } else {
+        absl::StrAppend(&comment, "* `Output`");
+      }
+      absl::StrAppend(&comment, " ", output_names[i]);
+      if (api_def.out_arg(i).description().empty()) {
+        absl::StrAppend(&comment, "\n");
+      } else {
+        // TODO(josh11b): Word wrap this.
+        absl::StrAppend(&comment, ": ", api_def.out_arg(i).description(), "\n");
+      }
+    }
+  }
+
+  if (!aliases.empty()) {
+    absl::StrAppend(&comment, "\nAliases:\n");
+    for (const auto& alias : aliases) {
+      absl::StrAppend(&comment, "* ", alias, "\n");
+    }
+  }
+  comment = MakeComment(comment, "");
+}
+
+std::string OpInfo::GetOpAttrStruct() const {
+  std::string struct_fields;
+  std::string setters;
+  std::string defaults_static_storage;
+
+  for (int i = 0; i < graph_op_def.attr_size(); ++i) {
+    const auto& attr(graph_op_def.attr(i));
+    const auto& api_def_attr(api_def.attr(i));
+    // If attr will be inferred or it doesn't have a default value, don't
+    // add it to the struct.
+    if ((inferred_input_attrs.find(attr.name()) !=
+         inferred_input_attrs.end()) ||
+        !api_def_attr.has_default_value()) {
+      continue;
+    }
+    const auto entry = AttrTypeName(attr.type());
+    const auto attr_type_name = entry.first;
+    const bool use_const = entry.second;
+    const std::string camel_case_name = ToCamelCase(api_def_attr.rename_to());
+    const std::string suffix =
+        (camel_case_name == op_name || camel_case_name == "Attrs") ? "_" : "";
+    const std::string attr_func_def =
+        absl::StrCat(camel_case_name, suffix, "(", use_const ? "const " : "",
+                     attr_type_name, use_const ? "&" : "");
+
+    std::string attr_comment;
+    if (!api_def_attr.description().empty()) {
+      absl::StrAppend(&attr_comment, api_def_attr.description(), "\n\n");
+    }
+    absl::StrAppend(&attr_comment, "Defaults to ",
+                    SummarizeAttrValue(api_def_attr.default_value()), "\n");
+    attr_comment = MakeComment(attr_comment, "    ");
+
+    absl::StrAppend(&setters, attr_comment);
+    absl::StrAppend(&setters, "    TF_MUST_USE_RESULT Attrs ", attr_func_def,
+                    " x) {\n");
+    absl::StrAppend(&setters, "      Attrs ret = *this;\n");
+    absl::StrAppend(&setters, "      ret.", api_def_attr.rename_to(),
+                    "_ = x;\n");
+    absl::StrAppend(&setters, "      return ret;\n    }\n\n");
+
+    std::string field_initiliazer;
+    auto& default_value = api_def_attr.default_value();
+    if (default_value.value_case() == AttrValue::kList &&
+        !IsEmptyList(default_value.list())) {
+      // Non-empty lists need static storage for their defaults. Define a
+      // function with static local variable that stores the array.
+      absl::StrAppend(&defaults_static_storage, "    static ", attr_type_name,
+                      " Default_", api_def_attr.rename_to(), "() {\n");
+      absl::StrAppend(
+          &defaults_static_storage, "      static const ",
+          ListElementTypeName(attr.type()), " kStorage[] = ",
+          PrintAttrValue(graph_op_def.name(), api_def_attr.default_value()),
+          ";\n");
+      absl::StrAppend(&defaults_static_storage, "      return ", attr_type_name,
+                      "(kStorage);\n    }\n");
+      // Set the field_initializer to call the defined function.
+      absl::StrAppend(&field_initiliazer, "Default_", api_def_attr.rename_to(),
+                      "()");
+    } else {
+      field_initiliazer =
+          PrintAttrValue(graph_op_def.name(), api_def_attr.default_value());
+    }
+    absl::StrAppend(&struct_fields, "    ", attr_type_name, " ",
+                    api_def_attr.rename_to(), "_ = ", field_initiliazer, ";\n");
+  }
+
+  if (struct_fields.empty()) {
+    return "";
+  }
+
+  std::string attrs_comment =
+      absl::StrCat("Optional attribute setters for ", op_name, "\n");
+  std::string struct_decl = MakeComment(attrs_comment, "  ");
+  absl::StrAppend(&struct_decl, "  struct Attrs {\n");
+  absl::StrAppend(&struct_decl, setters, struct_fields);
+  if (!defaults_static_storage.empty()) {
+    absl::StrAppend(&struct_decl, "  private:\n", defaults_static_storage);
+  }
+  absl::StrAppend(&struct_decl, "  };\n");
+
+  return struct_decl;
+}
+
+}  // namespace cc_op
+}  // namespace tensorflow

tensorflow/cc/framework/cc_op_gen_util.h

@@ -0,0 +1,154 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_UTIL_H_
+#define TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_UTIL_H_
+
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/op_def_util.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/platform/numbers.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace cc_op {
+
+absl::StatusOr<ApiDefMap> LoadOpsAndApiDefs(
+    OpList& ops, bool include_internal,
+    const std::vector<std::string>& api_def_dirs);
+
+// Converts:
+//   bazel-out/.../(bin|genfiles)/(external/YYY/)?XX
+// to: XX.
+std::string GetPath(absl::string_view dot_h_fname);
+
+// Converts: some/path/to/file.xx
+// to: file
+// (note that suffix is removed)
+std::string GetFilename(absl::string_view path);
+
+// Converts:
+//   cc/ops/gen_foo_ops.h
+// to:
+//   CC_OPS_GEN_FOO_OPS_H_
+std::string ToGuard(absl::string_view path);
+
+// Converts: some_name_xyz
+// to: Some Name Xyz
+std::string ToTitle(absl::string_view name);
+
+// Change:     Into:
+//   ABC         /// ABC
+//               ///
+//   DEF         /// DEF
+std::string MakeComment(absl::string_view text, absl::string_view indent);
+
+std::string PrintString(absl::string_view str);
+
+std::string PrintTensorShape(const TensorShapeProto& shape_proto);
+
+template <typename T>
+std::string PrintArray(int64_t num_elts, const T* array) {
+  std::string ret;
+  for (int64_t i = 0; i < num_elts; ++i) {
+    if (i > 0) absl::StrAppend(&ret, ", ");
+    absl::StrAppend(&ret, strings::LegacyPrecision(array[i]));
+  }
+  return ret;
+}
+
+std::string PrintTensor(const TensorProto& tensor_proto);
+
+std::string PrintTensorProto(const TensorProto& proto);
+
+std::string PrintAttrValue(absl::string_view, const AttrValue& attr_value);
+
+bool IsEmptyList(const AttrValue::ListValue& list);
+
+std::string ToCamelCase(absl::string_view str);
+
+std::string SeparateNamespaces(absl::string_view str);
+
+// Returns a <string, bool> pair. The string is the C++ type name to be used for
+// attr_type when defining an object of that type. The bool is a flag to
+// indicate whether to treat the type as const when accepting the C++ type as an
+// argument to a function.
+std::pair<absl::string_view, bool> AttrTypeName(absl::string_view attr_type);
+
+absl::string_view ListElementTypeName(absl::string_view attr_type);
+
+bool IsCPPKeyword(absl::string_view name);
+
+std::string AvoidCPPKeywords(absl::string_view name);
+
+void InferArgAttributes(
+    const OpDef::ArgDef& arg,
+    std::unordered_map<std::string, std::string>* inferred_attrs);
+
+void InferOpAttributes(
+    const OpDef& op_def,
+    std::unordered_map<std::string, std::string>* inferred_input_attrs);
+
+bool ArgIsList(const OpDef::ArgDef& arg);
+
+bool HasOptionalAttrs(
+    const ApiDef& api_def,
+    const std::unordered_map<std::string, std::string>& inferred_input_attrs);
+
+struct OpInfo {
+  // graph_op_def: The OpDef used by the runtime, has the names that
+  //   must be used when calling NodeBuilder.
+  // interface_op_def: The OpDef used in the interface in the generated
+  //   code, with possibly overridden names and defaults.
+  OpInfo(const OpDef& graph_op_def, const ApiDef& api_def,
+         const std::vector<std::string>& aliases);
+  OpInfo(const OpDef& graph_op_def, const ApiDef& api_def);
+  std::string GetOpAttrStruct() const;
+  std::string GetConstructorDecl(absl::string_view op_name_prefix,
+                                 bool include_attr) const;
+
+  std::string op_name;
+  std::vector<std::string> arg_types;
+  std::vector<std::string> arg_names;
+  std::vector<std::string> output_types;
+  std::vector<std::string> output_names;
+  std::vector<bool> is_list_output;
+  bool has_optional_attrs;
+  std::string comment;
+
+  const OpDef& graph_op_def;
+  const ApiDef& api_def;
+  const std::vector<std::string>& aliases;
+  // Map from type attribute to corresponding original argument name.
+  std::unordered_map<std::string, std::string> inferred_input_attrs;
+};
+
+}  // namespace cc_op
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_UTIL_H_

tensorflow/cc/framework/cc_ops_test.cc

@@ -0,0 +1,253 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/cc/ops/test_op.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+Output Linear(const Scope& scope, Input x, Input w, Input b) {
+  auto cop_scopes = scope.GetCompositeOpScopes("linear");
+  auto m = MatMul(cop_scopes.child, x, w);
+  return BiasAdd(cop_scopes.last, m, b);
+}
+
+void GetColocationConstraints(const Output& tensor,
+                              std::vector<std::string>* constraints) {
+  constraints->clear();
+  TF_EXPECT_OK(GetNodeAttr(tensor.op().node()->attrs(), kColocationAttrName,
+                           constraints));
+}
+
+TEST(CCOpTest, Basic) {
+  Scope root = Scope::NewRootScope();
+  auto c = Const(root, {{1, 1}});
+  // NOTE: The recommended style for constructing ops is
+  // auto v = OpConstructor(t0, t1, ..);
+  // Since the wrappers are implemented as one class per op, the following
+  // style is also possible :
+  // PrimitiveOp p(t0, t1, ...);
+  // It's being used here ONLY to ensure that, that style is tested.
+  MatMul m(root, c, {{41}, {1}});
+  TF_EXPECT_OK(root.status());
+  Tensor out;
+  test::GetTensor(root, m, &out);
+  test::ExpectTensorEqual<int>(out, test::AsTensor<int>({42}, {1, 1}));
+}
+
+TEST(CCOpTest, Attrs) {
+  Scope root = Scope::NewRootScope();
+  auto m = MatMul(root, {{1}, {1}}, {{41}, {1}}, MatMul::TransposeA(true));
+  TF_EXPECT_OK(root.status());
+  Tensor out;
+  test::GetTensor(root, m, &out);
+  test::ExpectTensorEqual<int>(out, test::AsTensor<int>({42}, {1, 1}));
+}
+
+TEST(CCOpTest, SplitConcat) {
+  Scope root = Scope::NewRootScope();
+  Split p(root, 0, {{1}, {2}}, 2);
+  auto c = Concat(root, {p[0], p[1]}, 0);
+  TF_EXPECT_OK(root.status());
+  Tensor out;
+  test::GetTensor(root, c, &out);
+  test::ExpectTensorEqual<int>(out, test::AsTensor<int>({1, 2}, {2, 1}));
+}
+
+TEST(CCOpTest, CompositeOp) {
+  Scope root = Scope::NewRootScope();
+  auto l = Linear(root.WithOpName("layer0"), {{10.0f, -3.0f}},
+                  {{.8f, .5f}, {.1f, .6f}}, {-8.0f, 31.0f});
+  TF_EXPECT_OK(root.status());
+  EXPECT_EQ(l.node()->name(), "layer0");
+  Tensor out;
+  test::GetTensor(root, l, &out);
+  test::ExpectClose(out, test::AsTensor<float>({-0.3, 34.2}, {1, 2}));
+}
+
+TEST(CCOpTest, MultiOutput) {
+  Scope root = Scope::NewRootScope();
+  auto u = Unique(root, {1, 2, 2, 4, 3, 2});
+  std::vector<Tensor> outputs;
+  test::GetTensors(root, {u.y, u.idx}, &outputs);
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({1, 2, 4, 3}));
+  test::ExpectTensorEqual<int>(outputs[1],
+                               test::AsTensor<int>({0, 1, 1, 2, 3, 1}));
+}
+
+TEST(CCOpTest, ExampleTrainer) {
+  Scope root = Scope::NewRootScope();
+  // a = [3 2; -1 0]
+  auto a = Const(root, {{3.f, 2.f}, {-1.f, 0.f}});
+  // x = [1.0; 1.0]
+  auto x = Const(root.WithOpName("x"), {{1.f}, {1.f}});
+  // y = a * x
+  auto y = MatMul(root.WithOpName("y"), a, x);
+  // y2 = y.^2
+  auto y2 = Square(root, y);
+  // y2_sum = sum(y2)
+  auto y2_sum = Sum(root, y2, 0);
+  // y_norm = sqrt(y2_sum)
+  auto y_norm = Sqrt(root, y2_sum);
+  // y_normalized = y ./ y_norm
+  auto y_normalized = Div(root.WithOpName("y_normalized"), y, y_norm);
+  Tensor out;
+  test::GetTensor(root, y_normalized, &out);
+  test::ExpectTensorNear<float>(
+      out, test::AsTensor<float>({0.98058069, -0.19611613}, {2, 1}), 1e-5);
+}
+
+TEST(CCOpTest, ThrowAwayOp) {
+  Scope root = Scope::NewRootScope();
+  ThrowAway1(root, 1, 2.3f, 1, 1, 1, ThrowAway1::Builder(42));
+  ThrowAway2(root, ThrowAway2::ThrowAway2_(3).Scope(1));
+  TF_EXPECT_OK(root.status());
+}
+
+TEST(CCOpTest, ControlDeps) {
+  Scope root = Scope::NewRootScope();
+  auto v = Variable(root, {}, DT_FLOAT);
+  auto assign = Assign(root, v, 41.0f);
+  Scope with_control_deps = root.WithControlDependencies(assign);
+  auto add = Add(with_control_deps, v, 1.0f);
+  Scope no_control_deps = with_control_deps.WithNoControlDependencies();
+  auto sub = Sub(no_control_deps, 3.0f, 2.0f);
+  auto is_inited =
+      IsVariableInitialized(no_control_deps.WithControlDependencies(sub), v);
+
+  TF_EXPECT_OK(root.status());
+
+  std::vector<Tensor> out;
+
+  test::GetTensors(root, {add}, &out);
+  test::ExpectTensorNear<float>(out[0], test::AsTensor<float>({42.0f}, {}),
+                                1e-5);
+
+  out.clear();
+  // Note : GetTensors creates a new session, so 'v' is uninitialized.
+  // sub should have no control deps, so it should not cause the assign to run.
+  // Hence is_inited should be false.
+  test::GetTensors(root, {sub, is_inited}, &out);
+  test::ExpectTensorNear<float>(out[0], test::AsTensor<float>({1.0f}, {}),
+                                1e-5);
+  test::ExpectTensorEqual<bool>(out[1], test::AsTensor<bool>({false}, {}));
+}
+
+TEST(CCOpTest, KernelLabel) {
+  Scope root = Scope::NewRootScope();
+  auto add = Add(root.WithKernelLabel("AddWithKernelLabel"), 1.0f, 2.0f);
+  TF_EXPECT_OK(root.status());
+  AttrSlice attrs = add.z.op().node()->attrs();
+  const auto* kernel_attr = attrs.Find("_kernel");
+  ASSERT_TRUE(kernel_attr);
+  TF_EXPECT_OK(AttrValueHasType(*kernel_attr, "string"));
+  EXPECT_EQ(kernel_attr->s(), "AddWithKernelLabel");
+}
+
+TEST(CCOpTest, ColocateWith) {
+  Scope root = Scope::NewRootScope();
+  auto c1 = Const(root.WithOpName("c1"), 1);
+  auto c2 = Const(root.WithOpName("c2").ColocateWith(c1), 2);
+  std::vector<std::string> constraints;
+  GetColocationConstraints(c2, &constraints);
+  EXPECT_EQ(constraints[0], "loc:@c1");
+
+  auto c3 = Const(root.WithOpName("c3").ColocateWith(c2), 3);
+  GetColocationConstraints(c3, &constraints);
+  EXPECT_EQ(constraints[0], "loc:@c1");
+
+  auto a = Const(root.WithOpName("a"), 4);
+  auto c4 = Const(root.WithOpName("c4").ColocateWith(a), 5);
+  GetColocationConstraints(c4, &constraints);
+  EXPECT_EQ(constraints[0], "loc:@a");
+
+  auto c5 = Const(root.WithOpName("c5").ColocateWith(c3).ColocateWith(c4), 6);
+  GetColocationConstraints(c5, &constraints);
+  EXPECT_EQ(constraints[0], "loc:@a");
+  EXPECT_EQ(constraints[1], "loc:@c1");
+
+  Scope with_colocate = root.ColocateWith(c3).ColocateWith(c4);
+  auto c6 = Const(with_colocate.WithOpName("c6").ClearColocation(), 7);
+  EXPECT_FALSE(c6.op().node()->attrs().Find("_class"));
+}
+
+TEST(CCOpTest, TemplatedConst) {
+  Scope root = Scope::NewRootScope();
+  auto c1 = ops::Const<float>(root, {{3, 2}, {-1, 0}});
+  TF_EXPECT_OK(root.status());
+
+  Tensor out;
+  test::GetTensor(root, c1, &out);
+  test::ExpectTensorEqual<float>(
+      out, test::AsTensor<float>({3.f, 2.f, -1.f, 0.f}, {2, 2}));
+
+  auto c2 = ops::Const<tstring>(root, {{"this"}, {"is"}, {"a"}, {"constant"}});
+  test::GetTensor(root, c2, &out);
+  test::ExpectTensorEqual<tstring>(
+      out, test::AsTensor<tstring>({"this", "is", "a", "constant"}, {4, 1}));
+}
+
+TEST(CCOpTest, EmptyConst) {
+  Scope root = Scope::NewRootScope();
+
+  auto c1 = ops::Const(root, {});
+  TF_CHECK_OK(root.status());
+
+  Tensor out;
+  test::GetTensor(root, c1, &out);
+  test::ExpectTensorEqual<float>(out, Tensor(DT_FLOAT, {0}));
+
+  auto c2 = ops::Const(root, {{}});
+  TF_CHECK_OK(root.status());
+  test::GetTensor(root, c2, &out);
+  test::ExpectTensorEqual<float>(out, Tensor(DT_FLOAT, {1, 0}));
+
+  auto c3 = ops::Const(root, {{{}, {}}});
+  TF_CHECK_OK(root.status());
+  test::GetTensor(root, c3, &out);
+  test::ExpectTensorEqual<float>(out, Tensor(DT_FLOAT, {1, 2, 0}));
+
+  auto c4 = ops::Const<int>(root, {{{}}});
+  TF_CHECK_OK(root.status());
+  test::GetTensor(root, c4, &out);
+  test::ExpectTensorEqual<int>(out, Tensor(DT_INT32, {1, 1, 0}));
+
+  ops::Const(root, {{}, {{}}});
+  EXPECT_FALSE(root.status().ok());
+}
+
+TEST(CCOpTest, InvalidFinalize) {
+  Scope root = Scope::NewRootScope();
+  auto read_up_to = ops::ReaderReadUpTo(root, Variable(root, {}, DT_STRING),
+                                        Variable(root, {}, DT_STRING),
+                                        static_cast<int32_t>(2));
+  EXPECT_FALSE(root.status().ok());
+  auto err_msg = std::string(root.status().message());
+  EXPECT_NE(err_msg.find("'num_records' passed int32 expected int64"),
+            std::string::npos);
+}
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/framework/fuzzing/BUILD

@@ -0,0 +1,71 @@
+# TODO(unda): describe this package.
+
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_copts")
+load(
+    "//tensorflow/cc/framework/fuzzing:op_fuzzing.bzl",
+    "tf_gen_op_wrappers_fuzz",
+)
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+cc_library(
+    name = "cc_op_fuzz_gen_main",
+    srcs = [
+        "cc_op_fuzz_gen.cc",
+        "cc_op_fuzz_gen.h",
+        "cc_op_fuzz_gen_main.cc",
+    ],
+    copts = tf_copts(),
+    data = [
+        "//tensorflow/core/api_def:base_api_def",
+    ],
+    deps = [
+        "//tensorflow/cc:cc_op_gen_util",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:op_gen_lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:hash",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/tsl/platform:status",
+    ],
+)
+
+# copybara:uncomment_begin
+# tf_gen_op_wrappers_fuzz(
+#     name = "array_ops_fuzz",
+#     api_def_srcs = ["//tensorflow/core/api_def:base_api_def"],
+#     kernel_deps = [
+#         "//tensorflow/c/kernels:bitcast_op",
+#         "//tensorflow/core/kernels:array",
+#         "//tensorflow/core/kernels:check_numerics_op",
+#         "//tensorflow/core/kernels:fake_quant_ops",
+#         "//tensorflow/core/kernels:quantized_ops",
+#         "//tensorflow/core/kernels:scatter_nd_op",
+#         "//tensorflow/core/kernels/image:extract_image_patches_op",
+#         "//tensorflow/core/kernels/image:extract_volume_patches_op",
+#         "//tensorflow/core/kernels/image:mirror_pad_op",
+#         "//tensorflow/core/kernels/linalg:matrix_band_part_op",
+#         "//tensorflow/core/kernels/linalg:matrix_diag_op",
+#         "//tensorflow/core/kernels/linalg:matrix_set_diag_op",
+#     ],
+#     op_def_src = "//tensorflow/core/ops:array_ops_op_lib",
+# )
+# copybara:uncomment_end
+
+bzl_library(
+    name = "op_fuzzing_bzl",
+    srcs = ["op_fuzzing.bzl"],
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow:tensorflow_bzl",
+        "//tensorflow/core/platform:rules_cc_bzl",
+    ],
+)

tensorflow/cc/framework/fuzzing/cc_op_fuzz_gen.cc

@@ -0,0 +1,332 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/fuzzing/cc_op_fuzz_gen.h"
+
+#include <iostream>
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
+#include "absl/strings/substitute.h"
+#include "tensorflow/cc/framework/cc_op_gen_util.h"
+#include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_def_util.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/hash.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+namespace cc_op {
+namespace {
+
+std::string DefaultValue(OpDef_AttrDef attr) {
+  static const auto* attr_default_value_map =
+      new absl::flat_hash_map<absl::string_view, absl::string_view,
+                              StringPieceHasher>{
+          {"int", "0"},
+          {"string", "\"\""},
+          {"list(int)", "{ 0, 1 }"},
+          {"list(float)", "{0.0, 1.0}"},
+          {"type", "DT_UINT8"},
+          {"shape",
+           "mediapipe::ParseTextProtoOrDie<TensorShapeProto>("
+           "\"dim:[] unknown_rank:true\")"}};
+
+  if (attr.has_minimum()) {
+    if (attr.type() == "int") {
+      return absl::StrCat(attr.minimum());
+    } else if (attr.type() == "list(int)") {
+      std::vector<int> v(attr.minimum());
+      for (int i = 0; i < v.size(); ++i) v[i] = i;
+      std::string s = absl::StrCat("{", absl::StrJoin(v, ","), "}");
+      return s;
+    }
+  }
+  if (attr.has_allowed_values()) {
+    if (!attr.allowed_values().list().s().empty()) {
+      return absl::StrCat("\"", attr.allowed_values().list().s(0), "\"");
+    } else if (!attr.allowed_values().list().type().empty()) {
+      return DataType_Name(attr.allowed_values().list().type(0));
+    }
+  }
+  auto entry = attr_default_value_map->find(attr.type());
+  if (entry == attr_default_value_map->end()) {
+    LOG(ERROR) << "Unsupported Attr type: " << attr.type();
+    return "";
+  }
+  return std::string(entry->second);
+}
+
+std::string WriteClassFuzzDef(const OpInfo& op_info) {
+  std::string class_signature_str = absl::Substitute(
+      "class Fuzz$0 : public FuzzSession<$1> {\n", op_info.op_name,
+      absl::StrJoin(op_info.graph_op_def.input_arg(), ", ",
+                    [](std::string* out, const auto arg) {
+                      absl::StrAppend(out, "Tensor");
+                      if (ArgIsList(arg)) absl::StrAppend(out, ", Tensor");
+                    }));
+
+  std::string build_graph_body = absl::StrCat(
+      absl::StrJoin(
+          op_info.graph_op_def.input_arg(), "",
+          [op_info](std::string* out, const OpDef_ArgDef arg) {
+            std::string type = "DT_UINT8";
+
+            if (arg.type() != DT_INVALID) {
+              type = DataType_Name(arg.type());
+            } else if (!arg.type_attr().empty()) {
+              OpDef_AttrDef attr =
+                  *FindAttr(arg.type_attr(), op_info.graph_op_def);
+              if (attr.has_default_value() &&
+                  attr.default_value().value_case() == AttrValue::kType) {
+                type = DataType_Name(attr.default_value().type());
+              } else if (attr.has_allowed_values() &&
+                         attr.allowed_values().value_case() ==
+                             AttrValue::kList &&
+                         !attr.allowed_values().list().type().empty()) {
+                type = DataType_Name(attr.allowed_values().list().type(0));
+              }
+            }
+            if (ArgIsList(arg)) {
+              strings::StrAppend(
+                  out, "    Input ", arg.name(),
+                  "_0 = ", "tensorflow::ops::Placeholder(scope.WithOpName(\"",
+                  arg.name(), "\"), ", type, ");\n");
+              strings::StrAppend(
+                  out, "    Input ", arg.name(),
+                  "_1 = ", "tensorflow::ops::Placeholder(scope.WithOpName(\"",
+                  arg.name(), "\"), ", type, ");\n");
+              absl::StrAppend(
+                  out, absl::Substitute("    InputList $0({$0_0, $0_1});\n",
+                                        arg.name()));
+            } else {
+              strings::StrAppend(
+                  out, "    auto ", arg.name(), " = ",
+                  "tensorflow::ops::Placeholder(scope.WithOpName(\"",
+                  arg.name(), "\"), ", type, ");\n");
+            }
+          }),
+      absl::StrJoin(op_info.graph_op_def.attr(), "",
+                    [op_info](std::string* out, const OpDef_AttrDef attr) {
+                      if (op_info.inferred_input_attrs.count(attr.name()) ==
+                              0 &&
+                          !attr.has_default_value()) {
+                        strings::StrAppend(out, "    auto ", attr.name(), " = ",
+                                           DefaultValue(attr), ";\n");
+                      }
+                    }));
+
+  std::string constructor_call_str = absl::Substitute(
+      "    tensorflow::ops::$0(scope.WithOpName(\"output\")$1);\n",
+      op_info.op_name,
+      absl::StrCat(
+          op_info.api_def.arg_order().empty()
+              ? absl::StrJoin(op_info.api_def.in_arg(), "",
+                              [](std::string* out, const auto api_def_arg) {
+                                strings::StrAppend(out, ", ",
+                                                   api_def_arg.name());
+                              })
+              : absl::StrJoin(op_info.api_def.arg_order(), "",
+                              [](std::string* out, const auto name) {
+                                strings::StrAppend(out, ", ", name);
+                              }),
+          absl::StrJoin(op_info.graph_op_def.attr(), "",
+                        [op_info](std::string* out, const OpDef_AttrDef attr) {
+                          if (op_info.inferred_input_attrs.count(attr.name()) ==
+                                  0 &&
+                              !attr.has_default_value()) {
+                            absl::StrAppend(out, ", ", attr.name());
+                          }
+                        })));
+
+  std::string fuzz_impl_signature_str = absl::Substitute(
+      "  void FuzzImpl($0) final {\n",
+      absl::StrJoin(
+          op_info.graph_op_def.input_arg(), ", ",
+          [](std::string* out, const auto arg) {
+            strings::StrAppend(out, "const Tensor& ", arg.name(), "_0");
+            if (ArgIsList(arg))
+              strings::StrAppend(out, ", const Tensor& ", arg.name(), "_1");
+          }));
+
+  std::string run_inputs_str = absl::Substitute(
+      "    RunInputs({$0});\n",
+      absl::StrJoin(op_info.graph_op_def.input_arg(), ", ",
+                    [](std::string* out, const auto arg) {
+                      if (ArgIsList(arg)) {
+                        strings::StrAppend(
+                            out, "{\"", arg.name(), "\", ", arg.name(), "_0}, ",
+                            "{\"", arg.name(), "\", ", arg.name(), "_1}");
+                      } else {
+                        strings::StrAppend(out, "{\"", arg.name(), "\", ",
+                                           arg.name(), "_0}");
+                      }
+                    }));
+
+  std::string fuzz_class_def = strings::StrCat(
+      class_signature_str, "  void BuildGraph(const Scope& scope) override {\n",
+      build_graph_body, constructor_call_str, "  }\n", fuzz_impl_signature_str,
+      run_inputs_str, "  }\n", "};\n");
+
+  return fuzz_class_def;
+}
+
+std::string WriteFuzzTest(const OpInfo& op_info) {
+  return absl::Substitute(
+      "FUZZ_TEST_F(Fuzz$0, Fuzz).WithDomains($1);\n", op_info.op_name,
+      absl::StrJoin(op_info.graph_op_def.input_arg(), ", ",
+                    [](std::string* out, const auto arg) {
+                      absl::StrAppend(out, "AnyTensor()");
+                      if (ArgIsList(arg)) absl::StrAppend(out, ", AnyTensor()");
+                    }));
+}
+
+std::string FuzzerFileStart() {
+  const std::string fuzz_namespace_begin = R"namespace(
+namespace tensorflow {
+namespace fuzzing {
+
+)namespace";
+
+  const std::string fuzz_header =
+      absl::StrCat(R"include(// This file is MACHINE GENERATED! Do not edit.
+
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/security/fuzzing/cc/fuzz_session.h"
+#include "third_party/mediapipe/framework/port/parse_text_proto.h"
+)include",
+                   fuzz_namespace_begin);
+
+  return fuzz_header;
+}
+
+std::string FuzzerFileEnd() {
+  const std::string fuzz_footer = R"footer(
+}  // namespace fuzzing
+}  // namespace tensorflow
+)footer";
+
+  return fuzz_footer;
+}
+
+}  // namespace
+
+bool OpFuzzingIsOk(const OpInfo& op_info) {
+  // Skip deprecated ops.
+  if (op_info.graph_op_def.has_deprecation() &&
+      op_info.graph_op_def.deprecation().version() <= TF_GRAPH_DEF_VERSION) {
+    std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+              << " is deprecated.\n";
+    return false;
+  }
+
+  // TODO(unda, b/249347507): should we hide fuzzers for hidden ops?
+  if (op_info.api_def.visibility() == ApiDef::HIDDEN) {
+    std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+              << " is hidden.\n";
+    return false;
+  }
+
+  if (op_info.api_def.visibility() == ApiDef::SKIP) {
+    std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+              << " is skipped.\n";
+    return false;
+  }
+
+  // TODO(unda) : zero input ops
+  std::set<std::string> zero_input_ops = {"Placeholder", "ImmutableConst"};
+  if (zero_input_ops.find(op_info.op_name) != zero_input_ops.end()) {
+    std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+              << " takes zero inputs.\n";
+    return false;
+  }
+
+  // TODO(unda, 253431636): constrained kernel
+  std::set<std::string> constrained_kernel = {"Diag",
+                                              "DiagPart",
+                                              "GatherNd",
+                                              "GatherV2",
+                                              "QuantizeAndDequantizeV2",
+                                              "QuantizeAndDequantizeV3",
+                                              "QuantizeAndDequantizeV4",
+                                              "QuantizeAndDequantizeV4Grad",
+                                              "QuantizedConcat",
+                                              "QuantizedInstanceNorm",
+                                              "QuantizedReshape",
+                                              "ScatterNd",
+                                              "TensorScatterUpdate"};
+
+  // TODO(unda, b/253431636): constrained kernel
+  if (constrained_kernel.find(op_info.op_name) != constrained_kernel.end()) {
+    std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+              << " has a constrained kernel.\n";
+    return false;
+  }
+
+  for (int i = 0; i < op_info.graph_op_def.input_arg_size(); ++i) {
+    const auto& arg(op_info.graph_op_def.input_arg(i));
+    // TODO(unda, b/249298521): deal with inputs that are required to be refs
+    if (arg.is_ref()) {
+      std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+                << " requires a ref argument.\n";
+      return false;
+    }
+  }
+
+  std::set<std::string> unhandled_attr_types = {
+      "list(type)",   "func",         "float",      "bool",
+      "tensor",       "list(string)", "list(bool)", "list(shape)",
+      "list(tensor)", "list(attr)"};
+  for (int i = 0; i < op_info.graph_op_def.attr_size(); ++i) {
+    const auto& attr(op_info.graph_op_def.attr(i));
+    const auto& api_def_attr(op_info.api_def.attr(i));
+    // Skip inferred arguments
+    if (op_info.inferred_input_attrs.count(attr.name()) > 0) continue;
+    // Skip if it has default value (TODO(unda, b/249345399): add our custom
+    // values)
+    if (api_def_attr.has_default_value()) continue;
+    // TODO(unda, b/253432797): handle unimplemented input attribute types
+    if (unhandled_attr_types.find(attr.type()) != unhandled_attr_types.end()) {
+      std::cout << "NOT fuzzing: " << op_info.graph_op_def.name()
+                << " requires an unhandled attr type (" << attr.type()
+                << ").\n";
+      return false;
+    }
+  }
+  std::cout << "fuzzing: " << op_info.graph_op_def.name() << "\n";
+  return true;
+}
+
+std::string WriteSingleFuzzer(const OpInfo& op_info, bool is_fuzzable) {
+  return absl::StrCat(
+      FuzzerFileStart(), is_fuzzable ? WriteClassFuzzDef(op_info) : "",
+      is_fuzzable ? WriteFuzzTest(op_info) : "", FuzzerFileEnd());
+}
+
+}  // namespace cc_op
+}  // namespace tensorflow

tensorflow/cc/framework/fuzzing/cc_op_fuzz_gen.h

@@ -0,0 +1,36 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_FUZZING_CC_OP_FUZZ_GEN_H_
+#define TENSORFLOW_CC_FRAMEWORK_FUZZING_CC_OP_FUZZ_GEN_H_
+
+#include "tensorflow/cc/framework/cc_op_gen_util.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace cc_op {
+
+// String with single fuzzer file content.
+std::string WriteSingleFuzzer(const OpInfo& op_info, bool is_fuzzable);
+
+// Do we have all we need to create a fuzzer
+bool OpFuzzingIsOk(const OpInfo& op_info);
+
+}  // namespace cc_op
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_FUZZING_CC_OP_FUZZ_GEN_H_

tensorflow/cc/framework/fuzzing/cc_op_fuzz_gen_main.cc

@@ -0,0 +1,98 @@
+// Main executable to generate op fuzzers
+
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/cc/framework/cc_op_gen_util.h"
+#include "tensorflow/cc/framework/fuzzing/cc_op_fuzz_gen.h"
+#include "xla/tsl/platform/status.h"
+#include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/str_util.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace cc_op {
+namespace {
+
+void WriteAllFuzzers(std::string root_location,
+                     std::vector<std::string> api_def_dirs,
+                     std::vector<std::string> op_names) {
+  OpList ops;
+  absl::StatusOr<ApiDefMap> api_def_map =
+      LoadOpsAndApiDefs(ops, false, api_def_dirs);
+
+  TF_CHECK_OK(api_def_map.status());
+
+  Env* env = Env::Default();
+  absl::Status status;
+  std::unique_ptr<WritableFile> fuzz_file = nullptr;
+  for (const OpDef& op_def : ops.op()) {
+    if (std::find(op_names.begin(), op_names.end(), op_def.name()) ==
+        op_names.end())
+      continue;
+
+    const ApiDef* api_def = api_def_map->GetApiDef(op_def.name());
+    if (api_def == nullptr) {
+      continue;
+    }
+
+    OpInfo op_info(op_def, *api_def, std::vector<std::string>());
+    status.Update(env->NewWritableFile(
+        root_location + "/" + op_def.name() + "_fuzz.cc", &fuzz_file));
+    status.Update(
+        fuzz_file->Append(WriteSingleFuzzer(op_info, OpFuzzingIsOk(op_info))));
+    status.Update(fuzz_file->Close());
+  }
+  TF_CHECK_OK(status);
+}
+
+}  // namespace
+}  // namespace cc_op
+}  // namespace tensorflow
+
+int main(int argc, char* argv[]) {
+  tensorflow::port::InitMain(argv[0], &argc, &argv);
+  if (argc != 4) {
+    for (int i = 1; i < argc; ++i) {
+      fprintf(stderr, "Arg %d = %s\n", i, argv[i]);
+    }
+    fprintf(stderr, "Usage: %s location api_def1,api_def2 op1,op2,op3\n",
+            argv[0]);
+    exit(1);
+  }
+  for (int i = 1; i < argc; ++i) {
+    fprintf(stdout, "Arg %d = %s\n", i, argv[i]);
+  }
+  std::vector<std::string> api_def_srcs = tensorflow::str_util::Split(
+      argv[2], ",", tensorflow::str_util::SkipEmpty());
+  std::vector<std::string> op_names = tensorflow::str_util::Split(
+      argv[3], ",", tensorflow::str_util::SkipEmpty());
+  tensorflow::cc_op::WriteAllFuzzers(argv[1], api_def_srcs, op_names);
+  return 0;
+}

tensorflow/cc/framework/fuzzing/op_fuzzing.bzl

@@ -0,0 +1,171 @@
+"""Functions for automatically generating fuzzers."""
+
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_not_windows",
+    "lrt_if_needed",
+    "tf_cc_binary",
+    "tf_copts",
+)
+load(
+    "//tensorflow/core/platform:rules_cc.bzl",
+    "cc_test",
+)
+
+def tf_gen_op_wrappers_fuzz(
+        name,
+        op_def_src,
+        api_def_srcs = [],
+        kernel_deps = []):
+    """
+    Generates fuzzers for several groups of ops.
+
+    For each one we need the corresponding OpDef, ApiDef and KernelDef,
+    since they all can contain constraints for the inputs.
+
+    Args:
+        name: the name of the fuzz artifact
+        op_def_src: op definitions
+        api_def_srcs: api definitions
+        kernel_deps: op kernel dependencies
+    """
+
+    # Create tool to generate .cc fuzzer files.
+    tf_cc_binary(
+        name = "op_fuzz_gen_tool",
+        copts = tf_copts(),
+        linkopts = if_not_windows(["-lm", "-Wl,-ldl"]) + lrt_if_needed(),
+        linkstatic = 1,  # Faster to link this one-time-use binary dynamically
+        deps = [
+            "//tensorflow/cc/framework/fuzzing:cc_op_fuzz_gen_main",
+            op_def_src,
+        ] + kernel_deps,
+    )
+
+    # Add relevant locations to look for api_defs.
+    api_def_src_locations = ",".join(["$$(dirname $$(echo $(locations " + api_def_src + ") | cut -d\" \" -f1))" for api_def_src in api_def_srcs])
+
+    out_fuzz_files = [op_name + "_fuzz.cc" for op_name in op_names]
+    native.genrule(
+        name = name + "_genrule",
+        outs = out_fuzz_files,
+        srcs = api_def_srcs,
+        tools = [":op_fuzz_gen_tool"],
+        cmd = ("$(location :op_fuzz_gen_tool) " +
+               " $$(dirname $(location " + out_fuzz_files[0] + "))" +
+               " " + api_def_src_locations + " " + (",".join(op_names))),
+    )
+
+    for op_name in op_names:
+        cc_test(
+            name = op_name.lower() + "_fuzz",
+            srcs = [op_name + "_fuzz.cc"],
+            deps = kernel_deps +
+                   [
+                       "//tensorflow/security/fuzzing/cc:fuzz_session",
+                       "@com_google_googletest//:gtest_main",
+                       "@com_google_fuzztest//fuzztest",
+                       "//tensorflow/cc:cc_ops",
+                       "//third_party/mediapipe/framework/port:parse_text_proto",
+                   ],
+        )
+
+op_names = [
+    "BatchMatrixBandPart",
+    "BatchMatrixDiag",
+    "BatchMatrixDiagPart",
+    "BatchMatrixSetDiag",
+    "BatchToSpace",
+    "BatchToSpaceND",
+    "Bitcast",
+    "BroadcastArgs",
+    "BroadcastTo",
+    "CheckNumerics",
+    "ConcatV2",
+    "ConjugateTranspose",
+    "DebugGradientIdentity",
+    "DeepCopy",
+    "DepthToSpace",
+    "Dequantize",
+    "EditDistance",
+    "Empty",
+    "EnsureShape",
+    "ExpandDims",
+    "ExtractImagePatches",
+    "ExtractVolumePatches",
+    "FakeQuantWithMinMaxArgs",
+    "FakeQuantWithMinMaxArgsGradient",
+    "FakeQuantWithMinMaxVars",
+    "FakeQuantWithMinMaxVarsGradient",
+    "FakeQuantWithMinMaxVarsPerChannel",
+    "FakeQuantWithMinMaxVarsPerChannelGradient",
+    "Fill",
+    "Fingerprint",
+    "Gather",
+    "GuaranteeConst",
+    "Identity",
+    "IdentityN",
+    "InplaceAdd",
+    "InplaceSub",
+    "InplaceUpdate",
+    "InvertPermutation",
+    "ListDiff",
+    "MatrixBandPart",
+    "MatrixDiag",
+    "MatrixDiagPart",
+    "MatrixDiagPartV2",
+    "MatrixDiagPartV3",
+    "MatrixDiagV2",
+    "MatrixDiagV3",
+    "MatrixSetDiag",
+    "MatrixSetDiagV2",
+    "MatrixSetDiagV3",
+    "MirrorPad",
+    "OneHot",
+    "OnesLike",
+    "Pack",
+    "Pad",
+    "PadV2",
+    "ParallelConcat",
+    "PlaceholderV2",
+    "PlaceholderWithDefault",
+    "PreventGradient",
+    "QuantizeAndDequantize",
+    "QuantizeV2",
+    "Rank",
+    "Reshape",
+    "ResourceStridedSliceAssign",
+    "ReverseSequence",
+    "ReverseV2",
+    "ScatterNdNonAliasingAdd",
+    "Shape",
+    "ShapeN",
+    "Size",
+    "Slice",
+    "Snapshot",
+    "SpaceToBatch",
+    "SpaceToBatchND",
+    "SpaceToDepth",
+    "Split",
+    "SplitV",
+    "Squeeze",
+    "StopGradient",
+    "StridedSlice",
+    "StridedSliceGrad",
+    "TensorScatterAdd",
+    "TensorScatterMax",
+    "TensorScatterMin",
+    "TensorScatterSub",
+    "TensorStridedSliceUpdate",
+    "Tile",
+    "TileGrad",
+    "Transpose",
+    "Unique",
+    "UniqueV2",
+    "UniqueWithCounts",
+    "UniqueWithCountsV2",
+    "Unpack",
+    "UnravelIndex",
+    "Where",
+    "ZerosLike",
+]

tensorflow/cc/framework/grad_op_registry.cc

@@ -0,0 +1,49 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/grad_op_registry.h"
+
+namespace tensorflow {
+namespace ops {
+
+// static
+GradOpRegistry* GradOpRegistry::Global() {
+  static GradOpRegistry* grad_op_registry = new GradOpRegistry;
+  return grad_op_registry;
+}
+
+bool GradOpRegistry::Register(const std::string& op, GradFunc func) {
+  CHECK(registry_.insert({op, func}).second) << "Existing gradient for " << op;
+  return true;
+}
+
+absl::Status GradOpRegistry::Lookup(const std::string& op,
+                                    GradFunc* func) const {
+  auto iter = registry_.find(op);
+  if (iter == registry_.end()) {
+    const std::string error_msg =
+        "No gradient defined for op: " + op +
+        ". Please see "
+        "https://www.tensorflow.org/code/"
+        "tensorflow/cc/gradients/README.md"
+        " for instructions on how to add C++ gradients.";
+    return absl::NotFoundError(error_msg);
+  }
+  *func = iter->second;
+  return absl::OkStatus();
+}
+
+}  // end namespace ops
+}  // namespace tensorflow

tensorflow/cc/framework/grad_op_registry.h

@@ -0,0 +1,77 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_
+#define TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+namespace ops {
+
+/// GradFunc is the signature for all gradient functions in GradOpRegistry.
+/// Implementations should add operations to compute the gradient outputs of
+/// 'op' (returned in 'grad_outputs') using 'scope' and 'grad_inputs'.
+typedef absl::Status (*GradFunc)(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs);
+
+/// GradOpRegistry maintains a static registry of gradient functions.
+/// Gradient functions are indexed in the registry by the forward op name (i.e.
+/// "MatMul" -> MatMulGrad func).
+class GradOpRegistry {
+ public:
+  /// Registers 'func' as the gradient function for 'op'.
+  /// Returns true if registration was successful, check fails otherwise.
+  bool Register(const std::string& op, GradFunc func);
+
+  /// Sets 'func' to the gradient function for 'op' and returns Status OK if
+  /// the gradient function for 'op' exists in the registry.
+  /// Note that 'func' can be null for ops that have registered no-gradient with
+  /// the registry.
+  /// Returns error status otherwise.
+  absl::Status Lookup(const std::string& op, GradFunc* func) const;
+
+  /// Returns a pointer to the global gradient function registry.
+  static GradOpRegistry* Global();
+
+ private:
+  std::unordered_map<std::string, GradFunc> registry_;
+};
+
+}  // namespace ops
+
+// Macros used to define gradient functions for ops.
+#define REGISTER_GRADIENT_OP(name, fn) \
+  REGISTER_GRADIENT_OP_UNIQ_HELPER(__COUNTER__, name, fn)
+
+#define REGISTER_NO_GRADIENT_OP(name) \
+  REGISTER_GRADIENT_OP_UNIQ_HELPER(__COUNTER__, name, nullptr)
+
+#define REGISTER_GRADIENT_OP_UNIQ_HELPER(ctr, name, fn) \
+  REGISTER_GRADIENT_OP_UNIQ(ctr, name, fn)
+
+#define REGISTER_GRADIENT_OP_UNIQ(ctr, name, fn) \
+  static bool unused_ret_val_##ctr =             \
+      ::tensorflow::ops::GradOpRegistry::Global()->Register(name, fn)
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_GRAD_OP_REGISTRY_H_

tensorflow/cc/framework/gradient_checker.cc

@@ -0,0 +1,451 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/gradient_checker.h"
+
+#include <algorithm>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_util.h"
+#include "tensorflow/core/framework/type_traits.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace {
+
+// TODO(andydavis) Support returning relative error (as opposed to max error)
+// between theoretical and numerical jacobians:
+//   fabs(jac_t - jac_n) / max(fabs(jac_t), fabs(jac_n))
+
+// TODO(andydavis) Vectorize and/or multi-thread Jacobian computations if
+// performance becomes an issue.
+
+// BaseUnitsForType provides a list of typed unit values for each basis in the
+// requested type.
+// When T is real,
+// BaseUnitsForType<T>::values() is just a single-entry vector [1]
+// When T is complex,
+// BaseUnitsForType<T>::values() is a two-entry vector [1, i] - the unit
+// values in each of its two bases.
+template <typename T>
+struct BaseUnitsForType {};  // Specializations below
+
+// Template specialization for BaseUnitsForType
+#define SET_BASE_UNITS_FOR_TYPE(TYPE, INIT)                         \
+  template <>                                                       \
+  struct BaseUnitsForType<TYPE> {                                   \
+    static const std::vector<TYPE>& values() {                      \
+      static std::vector<TYPE>* units = new std::vector<TYPE> INIT; \
+      return *units;                                                \
+    }                                                               \
+  }
+
+SET_BASE_UNITS_FOR_TYPE(float, {1});
+SET_BASE_UNITS_FOR_TYPE(double, {1});
+SET_BASE_UNITS_FOR_TYPE(complex64, ({{1, 0}, {0, 1}}));
+SET_BASE_UNITS_FOR_TYPE(complex128, ({{1, 0}, {0, 1}}));
+
+// SetJacobian sets the jacobian value at the provided row and column from a
+// tensor entry with type T.
+// When T is real, this is a simple assignment that casts the entry into the
+// jacobian type.
+// When T is complex, it assigns the real and complex values to successive rows
+// or columns in the matrix depending on the expand_by_row parameter
+template <typename T, typename JAC_T>
+typename std::enable_if<std::is_floating_point<T>::value>::type SetJacobian(
+    typename TTypes<JAC_T>::Matrix* jacobian, const int row, const int col,
+    const T& value, const bool expand_by_row) {
+  (*jacobian)(row, col) = JAC_T{value};
+}
+
+template <typename T, typename JAC_T>
+typename std::enable_if<is_complex<T>::value>::type SetJacobian(
+    typename TTypes<JAC_T>::Matrix* jacobian, const int row, const int col,
+    const T& value, const bool expand_by_row) {
+  (*jacobian)(row, col) = JAC_T{value.real()};
+  if (expand_by_row) {
+    (*jacobian)(row + 1, col) = JAC_T{value.imag()};
+  } else {
+    (*jacobian)(row, col + 1) = JAC_T{value.imag()};
+  }
+}
+
+// JacobianStride<T>::value holds the number of Jacobian elements needed to
+// represent one element of the given type.
+// When T is real the stride is 1, and when T is complex the stride is 2.
+template <typename T>
+struct JacobianStride {};  // Specializations below
+
+#define SET_JACOBIAN_STRIDE(TYPE, VALUE) \
+  template <>                            \
+  struct JacobianStride<TYPE> {          \
+    static constexpr int value = VALUE;  \
+  }
+
+SET_JACOBIAN_STRIDE(float, 1);
+SET_JACOBIAN_STRIDE(double, 1);
+SET_JACOBIAN_STRIDE(complex64, 2);
+SET_JACOBIAN_STRIDE(complex128, 2);
+
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeTheoreticalJacobianTranspose(
+    const Scope& scope, const OutputList& xs,
+    const std::vector<TensorShape>& x_shapes,
+    const std::vector<Tensor>& x_datas, const OutputList& ys,
+    const std::vector<TensorShape>& y_shapes,
+    std::vector<Tensor>* jacobian_ts) {
+  size_t y_num = y_shapes.size();
+  size_t x_num = x_shapes.size();
+  // Call AddSymbolicGradients to get 'dxs' (we will feed 'dys').
+  OutputList dys;
+  dys.reserve(y_shapes.size());
+  for (const auto& y_shape : y_shapes) {
+    // TODO(suharshs): This currently assumes that all y's are the same type.
+    dys.push_back(
+        ops::Cast(scope, ops::Const(scope, 1.0, y_shape), ys[0].type()));
+  }
+  OutputList dxs;
+  TF_RETURN_IF_ERROR(AddSymbolicGradients(scope, ys, xs, dys, &dxs));
+
+  // Initialize 'dy_data' to zeros.
+  std::vector<Tensor> dy_datas(y_num);
+  for (int i = 0; i < y_num; i++) {
+    dy_datas[i] = Tensor(ys[i].type(), y_shapes[i]);
+    auto dy_data_flat = dy_datas[i].flat<Y_T>();
+    dy_data_flat.setZero();
+  }
+
+  // Create the feed list.
+  ClientSession::FeedType feed_list;
+  for (int i = 0; i < x_num; i++) {
+    feed_list.insert({xs[i], x_datas[i]});
+  }
+  for (int i = 0; i < y_num; i++) {
+    feed_list.insert({dys[i], dy_datas[i]});
+  }
+
+  // x_stride and y_stride are used to calculate the correct jacobian row and
+  // column position for a pair of elements at positions r, c within the x and y
+  // tensors respectively.
+  const int x_stride = JacobianStride<X_T>::value;
+  const int y_stride = JacobianStride<Y_T>::value;
+  ClientSession session(scope);
+  for (int y_idx = 0; y_idx < y_num; y_idx++) {
+    auto dy_data_flat = dy_datas[y_idx].flat<Y_T>();
+    const int64_t dy_size = y_shapes[y_idx].num_elements();
+
+    // Compute the theoretical Jacobians one row at a time by back propagating
+    // '1.0' (or '1' and 'i' if y is complex) for each element of 'dy', while
+    // holding all other elements of 'dy' at zero.
+    for (int c = 0; c < dy_size; ++c) {
+      int unit_dimension = 0;
+      for (Y_T unit : BaseUnitsForType<Y_T>::values()) {
+        dy_data_flat(c) = unit;
+
+        std::vector<Tensor> dxout;
+        TF_RETURN_IF_ERROR(session.Run(feed_list, dxs, &dxout));
+
+        for (int x_idx = 0; x_idx < x_num; x_idx++) {
+          if (x_shapes[x_idx] != dxout[x_idx].shape()) {
+            return absl::InternalError(
+                absl::StrCat("Gradient for input ", x_idx, " expected shape ",
+                             x_shapes[x_idx].DebugString(), " but was ",
+                             dxout[x_idx].shape().DebugString()));
+          }
+          const int64_t x_size = x_shapes[x_idx].num_elements();
+          auto jacobian = (*jacobian_ts)[x_idx * y_num + y_idx].matrix<JAC_T>();
+          auto dx_flat = dxout[x_idx].flat<X_T>();
+          for (int r = 0; r < x_size; ++r) {
+            SetJacobian<X_T, JAC_T>(&jacobian, r * x_stride,
+                                    c * y_stride + unit_dimension, dx_flat(r),
+                                    true /* expand_by_row=true */);
+          }
+        }
+
+        dy_data_flat(c) = Y_T{0};
+        unit_dimension++;
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status EvaluateGraph(ClientSession* session, const OutputList& xs,
+                           const OutputList& ys, std::vector<Tensor>* x_datas,
+                           std::vector<Tensor>* y_datas) {
+  // Create the feed list.
+  ClientSession::FeedType feed_list;
+  for (int i = 0; i < x_datas->size(); i++) {
+    feed_list.insert({xs[i], (*x_datas)[i]});
+  }
+
+  TF_RETURN_IF_ERROR(session->Run(feed_list, ys, y_datas));
+  for (int y_idx = 0; y_idx < y_datas->size(); y_idx++) {
+    for (int x_idx = 0; x_idx < x_datas->size(); x_idx++) {
+      Tensor y_data = (*y_datas)[y_idx];
+      if (y_data.SharesBufferWith((*x_datas)[x_idx])) {
+        // Create copies of outputs that share a buffer with any inputs since
+        // the underlying buffer of the input Tensors are not copied for some
+        // operations (i.e. Identity), which can lead to incorrect results for
+        // the centered difference calculation.
+        (*y_datas)[y_idx] = tensor::DeepCopy(y_data);
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeNumericJacobianTranspose(
+    const Scope& scope, const OutputList& xs,
+    const std::vector<TensorShape>& x_shapes, const OutputList& ys,
+    const std::vector<TensorShape>& y_shapes, const JAC_T delta,
+    std::vector<Tensor>* x_datas, std::vector<Tensor>* jacobian_ts) {
+  size_t y_num = y_shapes.size();
+  size_t x_num = x_shapes.size();
+  // x_stride and y_stride are used to calculate the correct jacobian row and
+  // column position for a pair of elements at positions r, c within the x and y
+  // tensors respectively.
+  const int x_stride = JacobianStride<X_T>::value;
+  const int y_stride = JacobianStride<Y_T>::value;
+
+  // Check that the output tensors have the same shape as the expected shapes.
+  auto check_shapes = [&y_shapes](const std::vector<Tensor>& y_tensors) {
+    for (int y_idx = 0; y_idx < y_shapes.size(); y_idx++) {
+      if (y_tensors[y_idx].shape().num_elements() !=
+          y_shapes[y_idx].num_elements()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("Gradient for output ", y_idx, " expected shape ",
+                         y_shapes[y_idx].DebugString(), " but was ",
+                         y_tensors[y_idx].shape().DebugString()));
+      }
+    }
+    return absl::OkStatus();
+  };
+
+  ClientSession session(scope);
+  for (int x_idx = 0; x_idx < x_num; x_idx++) {
+    auto x_data_flat = (*x_datas)[x_idx].flat<X_T>();
+    const int64_t x_size = x_shapes[x_idx].num_elements();
+
+    // Compute the numeric Jacobian one column at a time by perturbing each
+    // element of 'x_data' (positively and negatively) by 'delta', and
+    // updating the jacobian with the centered difference. When x_data is
+    // complex-valued, we perturb its real and complex parts separately.
+    for (int r = 0; r < x_size; ++r) {
+      int unit_dimension = 0;
+      for (X_T unit : BaseUnitsForType<X_T>::values()) {
+        X_T x_delta = unit * X_T{delta};
+        // Store current value of 'x' at 'r'.
+        X_T v = x_data_flat(r);
+        // Evaluate at positive delta.
+        x_data_flat(r) = v + x_delta;
+        std::vector<Tensor> y_pos;
+        TF_RETURN_IF_ERROR(EvaluateGraph(&session, xs, ys, x_datas, &y_pos));
+        TF_RETURN_IF_ERROR(check_shapes(y_pos));
+        // Evaluate at negative delta.
+        x_data_flat(r) = v - x_delta;
+        std::vector<Tensor> y_neg;
+        TF_RETURN_IF_ERROR(EvaluateGraph(&session, xs, ys, x_datas, &y_neg));
+        TF_RETURN_IF_ERROR(check_shapes(y_neg));
+
+        for (int y_idx = 0; y_idx < y_num; y_idx++) {
+          // Compute element-wise centered difference and store in each
+          // Jacobian.
+          auto y_pos_flat = y_pos[y_idx].flat<Y_T>();
+          auto y_neg_flat = y_neg[y_idx].flat<Y_T>();
+          const int64_t y_size = y_shapes[y_idx].num_elements();
+          const Y_T scale = 2 * delta;
+          auto jacobian = (*jacobian_ts)[x_idx * y_num + y_idx].matrix<JAC_T>();
+          for (int c = 0; c < y_size; ++c) {
+            SetJacobian<Y_T, JAC_T>(&jacobian, r * x_stride + unit_dimension,
+                                    c * y_stride,
+                                    (y_pos_flat(c) - y_neg_flat(c)) / scale,
+                                    false /* expand_by_row=false */);
+          }
+        }
+        // Restore pre-perturbation value.
+        x_data_flat(r) = v;
+        unit_dimension++;
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+// The Jacobian is always a real-valued matrix.
+// Given y = f(x) for tensors y and x, it contains the derivatives dy_i/dx_j for
+// every pair y_i in y and x_j in x.  Note that the Jacobian is defined directly
+// over the elements of tensors y and x, and doesn't depend on their shapes.
+//
+// If x = (x_1, x_2, ..., x_m) and y = (y_1, y_2, .., y_n) the matrix evaluated
+// is actually the Jacobian transpose, defined as this mxn matrix:
+// dy_1/d_x1 dy_2/dx_1 ... dy_n/dx_1
+// dy_1/dx_2 dy_2/dx_2 ... dy_n/dx_2
+//     .
+//     .
+//     .
+// dy_1/dx_m dy_2/dx_m ... dy_n/dx_m
+//
+// If x or y is complex, each complex entry is "expanded" into a real and
+// imaginary entry, and the Jacobian is organized as above on the expanded list.
+// e.g.
+// [y1, y2] = Square([x1, x2]) where x and y are complex.
+// Writing
+// x = [x1_real, x1_imag, x2_real, x2_imag]
+// y = [y1_real, y1_imag, y2_real, y2_imag]
+// the Jacobian transpose is
+// the 4x4 matrix:
+// dy1_real/dx1_real dy1_imag/dx1_real dy2_real/dx1_real dy2_imag/dx1_real
+// dy1_real/dx1_imag dy1_imag/dx1_imag dy2_real/dx1_imag dy2_imag/dx1_imag
+// dy1_real/dx2_real dy1_imag/dx2_real dy2_real/dx2_real dy2_imag/dx2_real
+// dy1_real/dx2_imag dy1_imag/dx2_imag dy2_real/dx2_imag dy2_imag/dx2_imag
+template <typename X_T, typename Y_T, typename JAC_T>
+void InitJacobians(const OutputList& xs,
+                   const std::vector<TensorShape>& x_shapes,
+                   const std::vector<TensorShape>& y_shapes,
+                   std::vector<Tensor>* jacobians) {
+  const size_t y_num = y_shapes.size();
+  const size_t x_num = x_shapes.size();
+  const DataType jacobian_type = DataTypeToEnum<JAC_T>::v();
+
+  jacobians->resize(y_num * x_num);
+  for (int x_idx = 0; x_idx < x_num; x_idx++) {
+    // The number of rows is the number of elements in the x tensor multiplied
+    // by the number of Jacobian entries needed to represent each x type.
+    const int64_t x_size =
+        x_shapes[x_idx].num_elements() * JacobianStride<X_T>::value;
+    for (int y_idx = 0; y_idx < y_num; y_idx++) {
+      // The number of columns is the number of elements in the y tensor
+      // multiplied by the number of Jacobian entries needed to represent each
+      // y type.
+      const int64_t y_size =
+          y_shapes[y_idx].num_elements() * JacobianStride<Y_T>::value;
+      Tensor jacobian_t(jacobian_type, {x_size, y_size});
+      auto jacobian_t_flat = jacobian_t.flat<JAC_T>();
+      jacobian_t_flat.setZero();
+      (*jacobians)[x_idx * y_num + y_idx] = std::move(jacobian_t);
+    }
+  }
+}
+
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeGradientErrorInternal(
+    const Scope& scope, const OutputList& xs,
+    const std::vector<TensorShape>& x_shapes, const OutputList& ys,
+    const std::vector<TensorShape>& y_shapes, std::vector<Tensor>* x_datas,
+    JAC_T* max_error) {
+  // Initialize theoretical Jacobians to zeros.
+  std::vector<Tensor> jacobian_ts;
+  InitJacobians<X_T, Y_T, JAC_T>(xs, x_shapes, y_shapes, &jacobian_ts);
+
+  // Compute theoretical Jacobian.
+  TF_RETURN_IF_ERROR((ComputeTheoreticalJacobianTranspose<X_T, Y_T, JAC_T>(
+      scope, xs, x_shapes, *x_datas, ys, y_shapes, &jacobian_ts)));
+
+  // Initialize numeric Jacobian to zeros.
+  std::vector<Tensor> jacobian_ns;
+  InitJacobians<X_T, Y_T, JAC_T>(xs, x_shapes, y_shapes, &jacobian_ns);
+
+  // Compute numeric Jacobian.
+  TF_RETURN_IF_ERROR((ComputeNumericJacobianTranspose<X_T, Y_T, JAC_T>(
+      scope, xs, x_shapes, ys, y_shapes, JAC_T{1e-3f}, x_datas, &jacobian_ns)));
+
+  for (int i = 0; i < jacobian_ts.size(); i++) {
+    // Compute the maximum error between theoretical and numeric Jacobians.
+    *max_error = 0.0;
+    auto jac_t = jacobian_ts[i].matrix<JAC_T>();
+    auto jac_n = jacobian_ns[i].matrix<JAC_T>();
+    for (int r = 0; r < jacobian_ts[i].dim_size(0); ++r) {
+      for (int c = 0; c < jacobian_ts[i].dim_size(1); ++c) {
+        auto cur_error = std::fabs(jac_t(r, c) - jac_n(r, c));
+        // Treat any NaN as max_error and immediately return.
+        // (Note that std::max may ignore NaN arguments.)
+        if (std::isnan(cur_error)) {
+          *max_error = cur_error;
+          return absl::OkStatus();
+        }
+        *max_error = std::max(*max_error, cur_error);
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeGradientError(const Scope& scope, const OutputList& xs,
+                                  const std::vector<TensorShape>& x_shapes,
+                                  const OutputList& ys,
+                                  const std::vector<TensorShape>& y_shapes,
+                                  JAC_T* max_error) {
+  if (xs.size() != x_shapes.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("xs(size ", xs.size(), ") and x_shapes(size ",
+                     x_shapes.size(), ") must be the same size."));
+  }
+  if (ys.size() != y_shapes.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("ys(size ", ys.size(), ") and y_shapes(size ",
+                     y_shapes.size(), ") must be the same size."));
+  }
+  // Initialize 'x_datas' to random values.
+  std::vector<Tensor> x_datas(x_shapes.size());
+  for (int i = 0; i < x_shapes.size(); i++) {
+    x_datas[i] = Tensor(xs[i].type(), x_shapes[i]);
+    auto x_data_flat = x_datas[i].flat<X_T>();
+    x_data_flat.setRandom();
+  }
+  // Compute gradient error.
+  return ComputeGradientErrorInternal<X_T, Y_T, JAC_T>(
+      scope, xs, x_shapes, ys, y_shapes, &x_datas, max_error);
+}
+
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeGradientError(const Scope& scope, const Output& x,
+                                  const Tensor& x_init_value, const Output& y,
+                                  const TensorShape& y_shape,
+                                  JAC_T* max_error) {
+  // Initialize 'x_data' from 'x_init_value'.
+  std::vector<Tensor> x_datas(1, Tensor(x_init_value));
+  // Compute gradient error.
+  return ComputeGradientErrorInternal<X_T, Y_T, JAC_T>(
+      scope, {x}, {x_datas[0].shape()}, {y}, {y_shape}, &x_datas, max_error);
+}
+
+#define INSTANTIATE_GRAD_ERR_TYPE(X_T, Y_T, JAC_T)                     \
+  template Status ComputeGradientError<X_T, Y_T, JAC_T>(               \
+      const Scope& scope, const OutputList& xs,                        \
+      const std::vector<TensorShape>& x_shapes, const OutputList& ys,  \
+      const std::vector<TensorShape>& y_shapes, JAC_T* max_error);     \
+  template Status ComputeGradientError<X_T, Y_T, JAC_T>(               \
+      const Scope& scope, const Output& x, const Tensor& x_init_value, \
+      const Output& y, const TensorShape& y_shape, JAC_T* max_error);
+
+INSTANTIATE_GRAD_ERR_TYPE(float, float, float);
+INSTANTIATE_GRAD_ERR_TYPE(double, float, double);
+INSTANTIATE_GRAD_ERR_TYPE(double, double, double);
+INSTANTIATE_GRAD_ERR_TYPE(complex64, float, float);
+INSTANTIATE_GRAD_ERR_TYPE(float, complex64, float);
+INSTANTIATE_GRAD_ERR_TYPE(complex64, complex64, float);
+INSTANTIATE_GRAD_ERR_TYPE(complex128, complex128, double);
+
+}  // namespace tensorflow

tensorflow/cc/framework/gradient_checker.h

@@ -0,0 +1,66 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_
+#define TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+
+/// Returns in 'max_error' the maximum element-wise error for dy/dx between the
+/// computed and numeric Jacobian matrices where 'xs' and 'ys' are tensors.
+/// X_T and Y_T are the c++ types for the x and y tensors, and JAC_T is a
+/// real-valued type to store the Jacobian derivatives dy/dx.
+/// This function adds operations to the graph associated with 'scope'.
+///
+/// Examples:
+/// if y = Square(x), where x (and so y) are DT_FLOAT,
+/// <X_T, Y_T, JAC_T> should be <float, float, float>
+///
+/// if y = Square(x), where x (and so y) are DT_DOUBLE,
+/// <X_T, Y_T, JAC_T> should be <double, double, double>
+///
+/// if y = Square(x), where x (and so y) are DT_COMPLEX64,
+/// <X_T, Y_T, JAC_T> should be <complex64, complex64, float>
+/// Note that JAC_T is always real-valued, and should be an appropriate
+/// precision to host the partial derivatives for dy/dx
+///
+/// if y = ComplexAbs(x) where x is DT_COMPLEX64 (so y is DT_FLOAT)
+/// <X_T, Y_T, JAC_T> should be <complex64, float, float>
+///
+/// if y = Complex(x, x) where x is DT_FLOAT (so y is DT_COMPLEX64)
+/// <X_T, Y_T, JAC_T> should be <float, complex64, float>
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeGradientError(const Scope& scope, const OutputList& xs,
+                                  const std::vector<TensorShape>& x_shapes,
+                                  const OutputList& ys,
+                                  const std::vector<TensorShape>& y_shapes,
+                                  JAC_T* max_error);
+
+/// Overload of ComputeGradientError which takes an initial value for 'x'.
+template <typename X_T, typename Y_T, typename JAC_T>
+absl::Status ComputeGradientError(const Scope& scope, const Output& x,
+                                  const Tensor& x_init_value, const Output& y,
+                                  const TensorShape& y_shape, JAC_T* max_error);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_GRADIENT_CHECKER_H_

tensorflow/cc/framework/gradient_checker_test.cc

@@ -0,0 +1,206 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/gradient_checker.h"
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/equal_graph_def.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Complex;
+using ops::Const;
+using ops::Div;
+using ops::MatMul;
+using ops::Placeholder;
+using ops::Real;
+using ops::Split;
+using ops::Square;
+using ops::Stack;
+using ops::Sub;
+using ops::Unstack;
+
+TEST(GradientCheckerTest, BasicFloat) {
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Square(scope, x);
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-4);
+}
+
+TEST(GradientCheckerTest, BasicDouble) {
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_DOUBLE, Placeholder::Shape(shape));
+  auto y = Square(scope, x);
+  double max_error;
+  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-10);
+}
+
+TEST(GradientCheckerTest, BasicComplex64) {
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_COMPLEX64, Placeholder::Shape(shape));
+  auto y = Square(scope, x);
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<complex64, complex64, float>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-4);
+}
+
+TEST(GradientCheckerTest, BasicComplex128) {
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_COMPLEX128, Placeholder::Shape(shape));
+  auto y = Square(scope, x);
+  double max_error;
+  TF_ASSERT_OK((ComputeGradientError<complex128, complex128, double>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-10);
+}
+
+TEST(GradientCheckerTest, FloatToComplex64) {
+  // Test an op whose inputs are real and outputs are complex
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Complex(scope, x, x);
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<float, complex64, float>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-4);
+}
+
+TEST(GradientCheckerTest, Complex64ToFloat) {
+  // Test an op whose inputs are complex and outputs are real
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_COMPLEX64, Placeholder::Shape(shape));
+  auto y = Real(scope, x);
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<complex64, float, float>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-4);
+}
+
+// When calculating gradients that are undefined, test we get NaN
+// as the computed error rather than 0.
+TEST(GradientCheckerTest, BasicNan) {
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({2, 4, 3});
+  auto x = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+  // y = x/(x-x) should always return NaN
+  auto y = Div(scope, x, Sub(scope, x, x));
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+      scope, {x}, {shape}, {y}, {shape}, &max_error)));
+  EXPECT_TRUE(std::isnan(max_error));
+}
+
+TEST(GradientCheckerTest, MatMulGrad) {
+  Scope scope = Scope::NewRootScope();
+
+  TensorShape x_shape({4, 3});
+  TensorShape y_shape({3, 2});
+  TensorShape z_shape({4, 2});
+
+  auto x = Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape));
+  auto y = Const(scope, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, y_shape);
+  auto z = MatMul(scope, x, y);
+  double max_error;
+  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
+      scope, {x}, {x_shape}, {z}, {z_shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-10);
+}
+
+TEST(GradientCheckerTest, SplitGrad) {
+  // Split is an op with single inputs and multiple outputs.
+  Scope scope = Scope::NewRootScope();
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape));
+  // Split along the second dimension.
+  auto split_dim = Const(scope, 1, {});
+  auto y = Split(scope, split_dim, x, /* num_split */ 2);
+  TensorShape y_shape = TensorShape({5, 1});
+  double max_error;
+  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
+      scope, {x}, {x_shape}, y.output, {y_shape, y_shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-10);
+}
+
+TEST(GradientCheckerTest, StackGrad) {
+  // Stack is an op with multiple inputs and a single output.
+  Scope scope = Scope::NewRootScope();
+  TensorShape x_shape({1, 2, 3});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape)));
+  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape)));
+  auto y = Stack(scope, xs, Stack::Axis(0));
+  TensorShape y_shape({2, 1, 2, 3});
+  double max_error;
+  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
+      scope, xs, {x_shape, x_shape}, {y}, {y_shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-10);
+}
+
+TEST(GradientCheckerTest, StackUnstackGrad) {
+  // Chaining a Stack op to an Unstack op allows us to test the gradient checker
+  // in a multiple input/output scenario.
+  Scope scope = Scope::NewRootScope();
+  TensorShape shape({1, 2, 3});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(shape)));
+  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(shape)));
+  auto tmp = Stack(scope, xs, Stack::Axis(0));
+  auto y = Unstack(scope, tmp, 2, Unstack::Axis(0));
+  double max_error;
+  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
+      scope, xs, {shape, shape}, y.output, {shape, shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-10);
+}
+
+TEST(GradientCheckerTest, ShapeMismatchError) {
+  Scope scope = Scope::NewRootScope();
+  auto x = Placeholder(scope, DT_FLOAT);
+  auto y = ops::Concat(scope, std::vector<Output>{x, x}, ops::Const(scope, 0));
+
+  TensorShape x_shape({1});
+  TensorShape y_shape({4});  // Intentionally wrong, evaluate returns 2.
+
+  float max_error;
+  auto status = ComputeGradientError<float, float, float>(
+      scope, {x}, {x_shape}, {y}, {y_shape}, &max_error);
+  EXPECT_FALSE(status.ok());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_NE(
+      std::string(status.message()).find("expected shape [4] but was [2]"),
+      std::string::npos);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/framework/gradients.cc

@@ -0,0 +1,589 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/gradients.h"
+
+#include <deque>
+#include <map>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/while_gradients.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/graph/algorithm.h"
+#include "tensorflow/core/graph/while_context.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/platform/macros.h"
+
+namespace tensorflow {
+namespace {
+
+struct OutputHash {
+  uint64_t operator()(const Output& x) const { return x.hash(); }
+};
+
+struct OutputEq {
+  bool operator()(const Output& x, const Output& y) const {
+    return (x.node() == y.node()) && (x.index() == y.index());
+  }
+};
+
+class SymbolicGradientBuilder {
+ public:
+  SymbolicGradientBuilder(const Scope& scope,
+                          const ops::GradOpRegistry* registry,
+                          const std::vector<Output>& outputs,
+                          const std::vector<Output>& inputs,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs);
+
+  absl::Status AddGradients();
+
+  static Output NoGradient() { return Output(nullptr, -1); }
+
+ private:
+  absl::Status Initialize();
+
+  // For each forward edge from `src` to `dst` in the initial/forward graph:
+  // propagates gradients `dst_grad` backwards along the edge from `src`
+  // to `dst` in the graph. This will add `dst_grad` to the list of pending
+  // gradients for the node associated with `src`.
+  absl::Status BackpropAlongEdge(const Output& dst_grad, const Output& src);
+
+  // Adds a node to the graph (returned in `grad`) that sums the in-bound
+  // gradients to `src` (if there are more than one).
+  absl::Status SumGradients(const Output& src, Output* grad);
+
+  // Returns true if `opname` is registered in `registry_` with no gradient
+  // function, false otherwise.
+  bool IsPrimitiveOpWithNoGrad(const std::string& opname);
+
+  // Call the gradient function for `op`, storing the result in `grad_outputs`.
+  absl::Status CallGradFunction(const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs);
+
+  // Returns a list mapping whether each node in the graph is reachable
+  // from outputs_. Keyed by node id.
+  std::vector<bool> GetReachableNodes();
+
+  // Creates the gradient subgraph for a while loop (or just stores
+  // `summed_grads` if not all incoming gradients are available yet). All exit
+  // nodes (which are the first nodes of a loop encountered in the backwards
+  // pass) are passed to this function rather than processed normally.
+  // `summed_grads` is the sum of `exit_node`s gradients.
+  absl::Status ProcessWhileLoop(Node* exit_node, const Output& summed_grads);
+
+  // Gets the set of node ids at which to stop backprop. These are all elements
+  // of `outputs_` that do not get transitively consumed by other `outputs_`.
+  // Used to identify nodes at which to stop backprop.
+  std::unordered_set<int> GetStopBackpropNodes(
+      const std::vector<bool>& reachable_nodes,
+      const std::unordered_set<int>& output_nodes) const;
+
+  const Scope& scope_;
+  const ops::GradOpRegistry* registry_;
+  const std::vector<Output>& outputs_;
+  const std::vector<Output>& inputs_;
+  const std::vector<Output>& grad_inputs_;
+  std::vector<Output>* grad_outputs_;
+
+  // A vector of output endpoints which represents backpropagated gradients.
+  typedef std::vector<Output> BackproppedGradients;
+
+  // backprops_ is a map from a node output to its accumulated
+  // gradients.  When a node output has accumulated all its
+  // gradients, we add a node which sums them up.
+  std::unordered_map<Output, BackproppedGradients, OutputHash, OutputEq>
+      backprops_;
+
+  // pending[i] is count-down counter for i-th node's expected
+  // backprops.  When pending[i] becomes zero, we collected all
+  // backprop gradients for all outputs of the ith-node.
+  std::vector<int> pending_;
+
+  // `ready` keeps track of nodes that have been completely
+  // backpropped. Initially, for every output in `outputs_`, we add initial
+  // gradients from `grad_inputs_`.
+  std::deque<Node*> ready_;
+
+  // The set of node ids in `inputs_`. Used to identify nodes at backprop
+  // frontier. Maps from Output -> index into `grad_outputs_`.
+  std::unordered_map<Output, int, OutputHash, OutputEq> input_nodes_;
+
+  // For each while loop in the graph, collects the summed gradients for each of
+  // the loop's exit nodes. Note that unlike backprops_, this map contains the
+  // output of SumGradients(), not the input (i.e. each exit node may have
+  // multiple incoming gradients, but we only store the combined Output here).
+  std::map<WhileContext*, std::map<Node*, Output>> while_backprops_;
+
+  SymbolicGradientBuilder(const SymbolicGradientBuilder&) = delete;
+  void operator=(const SymbolicGradientBuilder&) = delete;
+};
+
+SymbolicGradientBuilder::SymbolicGradientBuilder(
+    const Scope& scope, const ops::GradOpRegistry* registry,
+    const std::vector<Output>& outputs, const std::vector<Output>& inputs,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs)
+    : scope_(scope),
+      registry_(registry),
+      outputs_(outputs),
+      inputs_(inputs),
+      grad_inputs_(grad_inputs),
+      grad_outputs_(grad_outputs) {}
+
+absl::Status SymbolicGradientBuilder::BackpropAlongEdge(const Output& dst_grad,
+                                                        const Output& src) {
+  if (src.node() == nullptr) {
+    return absl::InternalError("Attempted to backprop along an invalid edge.");
+  }
+  auto iter = backprops_.find(src);
+  if (iter != backprops_.end()) {
+    auto* grads = &iter->second;
+    grads->push_back(dst_grad);
+    if (--pending_[src.node()->id()] == 0) {
+      ready_.push_back(src.node());
+    }
+  }
+  return absl::OkStatus();
+}
+
+std::vector<bool> SymbolicGradientBuilder::GetReachableNodes() {
+  std::vector<bool> reachable_nodes(scope_.graph()->num_node_ids(), false);
+  std::deque<Node*> queue;
+  for (const Output& out : outputs_) {
+    if (!reachable_nodes[out.node()->id()]) {
+      queue.push_back(out.node());
+      reachable_nodes[out.node()->id()] = true;
+    }
+  }
+
+  while (!queue.empty()) {
+    Node* n = queue.front();
+    queue.pop_front();
+    for (const Edge* e : n->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      if (!reachable_nodes[e->src()->id()]) {
+        queue.push_back(e->src());
+        reachable_nodes[e->src()->id()] = true;
+      }
+    }
+  }
+  return reachable_nodes;
+}
+
+std::unordered_set<int> SymbolicGradientBuilder::GetStopBackpropNodes(
+    const std::vector<bool>& reachable_nodes,
+    const std::unordered_set<int>& output_nodes) const {
+  // Output nodes that get transitively consumed by other `outputs_` are stored
+  // in `internal_outputs`.
+  std::unordered_set<int> internal_outputs;
+  std::unordered_set<Node*> visited;
+  // Initialize `queue` for BFS traversal. Nodes in `queue` hold upcoming nodes
+  // along with the last Node in `output_` encountered along that path. If no
+  // `output_` node was encountered, pair.second will be nullptr.
+  std::deque<std::pair<Node*, Node*>> queue;
+  for (const Output& nout : inputs_) {
+    auto const& pair = visited.insert(nout.node());
+    if (pair.second) {
+      queue.push_back(std::make_pair(nout.node(), static_cast<Node*>(nullptr)));
+    }
+  }
+  // BFS from nodes in 'inputs_' along out edges for the entire graph. Internal
+  // output nodes are recorded during the traversal. All nodes that are output
+  // nodes but not internal output nodes are considered the frontier of the
+  // output nodes, and thus our stop backprop nodes.
+  while (!queue.empty()) {
+    std::pair<Node*, Node*> p = queue.front();
+    Node* n = p.first;
+    queue.pop_front();
+    for (const Edge* e : n->out_edges()) {
+      // If a node is not reachable from outputs_, we can stop.
+      if (e->IsControlEdge() || !reachable_nodes[e->dst()->id()]) continue;
+
+      auto const& pair = visited.insert(e->dst());
+      if (pair.second) {
+        int node_id = e->dst()->id();
+        Node* last_output_node = p.second;
+        if (output_nodes.find(node_id) != output_nodes.end()) {
+          // We reached an output node.
+          if (last_output_node != nullptr) {
+            // If we had already found an output node on this path so we mark
+            // it as an internal output.
+            internal_outputs.insert(last_output_node->id());
+          }
+          // Mark this newly found output node to insert in the queue.
+          last_output_node = e->dst();
+        }
+        queue.push_back(std::make_pair(e->dst(), last_output_node));
+      }
+    }
+  }
+  // Finally, we set stop_backprop_nodes to all output_nodes that aren't also
+  // internal_outputs.
+  std::unordered_set<int> stop_backprop_nodes;
+  for (int output_node : output_nodes) {
+    if (internal_outputs.find(output_node) == internal_outputs.end()) {
+      stop_backprop_nodes.insert(output_node);
+    }
+  }
+  return stop_backprop_nodes;
+}
+
+absl::Status SymbolicGradientBuilder::Initialize() {
+  if (outputs_.size() != grad_inputs_.size()) {
+    return absl::InvalidArgumentError(
+        "Must specify a gradient input for each output.");
+  }
+  std::vector<bool> reachable_nodes = GetReachableNodes();
+  for (const Output& input : inputs_) {
+    if (!reachable_nodes[input.node()->id()]) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Cannot compute the partial derivative for node '",
+                       input.node()->name(),
+                       "' as it's unreachable from the output node(s)."));
+    }
+  }
+  grad_outputs_->clear();
+  grad_outputs_->resize(inputs_.size());
+
+  std::unordered_set<int> output_nodes;
+  output_nodes.reserve(outputs_.size());
+  for (size_t i = 0; i < outputs_.size(); ++i) {
+    output_nodes.insert(outputs_[i].node()->id());
+  }
+
+  std::unordered_set<int> stop_backprop_nodes =
+      GetStopBackpropNodes(reachable_nodes, output_nodes);
+
+  // Populate `input_nodes_` from Outputs in `inputs_`.
+  input_nodes_.reserve(inputs_.size());
+  for (size_t i = 0; i < inputs_.size(); ++i) {
+    input_nodes_.insert({inputs_[i], i});
+  }
+
+  // TODO(andydavis) Consider a more efficient data structure for `pending_` to
+  // handle computing gradients over small subgraphs from a very large graph.
+  pending_.resize(scope_.graph()->num_node_ids(), 0);
+  {
+    backprops_.clear();
+    std::unordered_set<Node*> visited;
+    std::deque<Node*> queue;
+    for (const Output& nout : inputs_) {
+      auto const& pair = visited.insert(nout.node());
+      if (pair.second) {
+        queue.push_back(nout.node());
+      }
+    }
+
+    // Going forward to figure out which endpoints need backprop-ed.
+    // A node's endpoints need to be backprop-ed only if one of the
+    // arg node can reach the node via data edges.
+    while (!queue.empty()) {
+      Node* n = queue.front();
+      queue.pop_front();
+      for (int i = 0; i < n->num_outputs(); ++i) {
+        backprops_[{n, i}].clear();
+      }
+      int num_expected_backprops = 0;
+      if (stop_backprop_nodes.find(n->id()) == stop_backprop_nodes.end()) {
+        // Internal node: continue BFS along connected outputs.
+        for (const Edge* e : n->out_edges()) {
+          // If a node is not reachable from outputs_,
+          // we don't expect it to receive a backpropagated gradient.
+          // It will not be counted in num_expected_backprops.
+          if (e->IsControlEdge() || !reachable_nodes[e->dst()->id()]) continue;
+          auto const& pair = visited.insert(e->dst());
+          if (pair.second) {
+            queue.push_back(e->dst());
+          }
+          ++num_expected_backprops;
+        }
+      }
+      if (output_nodes.find(n->id()) != output_nodes.end()) {
+        // Output node: update `num_expected_backprops` for each Output in
+        // `outputs_` that references `n`.
+        for (const Output& output : outputs_) {
+          if (output.node() == n) {
+            ++num_expected_backprops;
+          }
+        }
+      }
+      pending_[n->id()] = num_expected_backprops;
+    }
+  }
+
+  {
+    // Initialize backprop with `grad_inputs_`.
+    const size_t num_dy = grad_inputs_.size();
+    for (size_t i = 0; i < num_dy; ++i) {
+      TF_RETURN_IF_ERROR(BackpropAlongEdge(grad_inputs_[i], outputs_[i]));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status SymbolicGradientBuilder::SumGradients(const Output& src,
+                                                   Output* grad) {
+  auto iter = backprops_.find(src);
+  if (iter == backprops_.end()) {
+    return absl::InternalError(absl::StrCat(
+        "Unable to find backprop list for node.id ", src.node()->name()));
+  }
+  const auto& grads = iter->second;
+  // Filter any backpropped 'NoGradient' Outputs from 'grads' (if needed).
+  // Return any valid backpropped gradients that remain after filtering,
+  // or 'NoGradient' otherwise.
+  std::vector<Output> grads_to_keep;
+  for (const Output& o : grads) {
+    if (o == NoGradient()) continue;
+    grads_to_keep.push_back(o);
+  }
+
+  if (grads_to_keep.empty()) {
+    // Nothing propagated back. Return 'NoGradient'.
+    *grad = NoGradient();
+  } else if (grads_to_keep.size() == 1) {
+    // Just one backprop edge.
+    *grad = grads_to_keep[0];
+  } else {
+    // Otherwise, adds backprop-ed gradients.
+    // TODO(andydavis) Use a better accumulator here.
+    *grad = ops::AddN(scope_, grads_to_keep);
+  }
+
+  return absl::OkStatus();
+}
+
+bool SymbolicGradientBuilder::IsPrimitiveOpWithNoGrad(
+    const std::string& opname) {
+  ops::GradFunc grad_fn;
+  absl::Status s = registry_->Lookup(opname, &grad_fn);
+  return s.ok() && (grad_fn == nullptr);
+}
+
+absl::Status SymbolicGradientBuilder::CallGradFunction(
+    const Operation& op, const std::vector<Output>& grad_inputs,
+    std::vector<Output>* grad_outputs) {
+  ops::GradFunc grad_fn;
+  TF_RETURN_IF_ERROR(registry_->Lookup(op.node()->type_string(), &grad_fn));
+  TF_RETURN_IF_ERROR(grad_fn(scope_, op, grad_inputs, grad_outputs));
+  TF_RETURN_IF_ERROR(scope_.status());
+  return absl::OkStatus();
+}
+
+absl::Status SymbolicGradientBuilder::ProcessWhileLoop(
+    Node* exit_node, const Output& summed_grads) {
+  // TODO(skyewm): detect second-order gradient and return bad status
+  // TODO(skyewm): handle (or at least detect) nested while loops
+
+  // TODO(skyewm): handle NoGradient in while loop
+  if (summed_grads == NoGradient()) {
+    return absl::UnimplementedError(
+        "Missing gradient into while loop not yet implemented");
+  }
+
+  DCHECK(exit_node->IsExit());
+  WhileContext* while_ctx = exit_node->while_ctx();
+  DCHECK(while_ctx != nullptr);
+
+  // Record 'summed_grads' as the backprop input associated with 'exit_node'
+  std::map<Node*, Output>& backprops = while_backprops_[while_ctx];
+  DCHECK(backprops.find(exit_node) == backprops.end());
+  backprops[exit_node] = summed_grads;
+
+  // Wait until we have all exit nodes' backprops collected before processing
+  // the while loop.
+  // TODO(skyewm): what if not all the exit nodes are reachable?
+  if (backprops.size() < while_ctx->exit_nodes().size())
+    return absl::OkStatus();
+
+  // We've seen all the exit nodes for this loop and have collected all the
+  // backprops. Create the gradient graph for the while loop.
+  Scope while_scope =
+      scope_.NewSubScope(absl::StrCat(while_ctx->frame_name(), "_grad"));
+  std::vector<Output> dy;
+  for (Node* n : while_ctx->exit_nodes()) dy.push_back(backprops[n]);
+  std::vector<Output> dx;
+  TF_RETURN_IF_ERROR(AddWhileLoopGradient(while_ctx, while_scope, dy, &dx));
+
+  // Backprop along the in edges to the while loop (i.e. the inputs to the enter
+  // nodes)
+  DCHECK_EQ(dx.size(), while_ctx->enter_nodes().size());
+  for (int i = 0, end = dx.size(); i < end; ++i) {
+    Node* enter_node = while_ctx->enter_nodes()[i];
+    for (const Edge* e : enter_node->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      TF_RETURN_IF_ERROR(BackpropAlongEdge(dx[i], {e->src(), e->src_output()}));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status SymbolicGradientBuilder::AddGradients() {
+  // Initialize backprops.
+  TF_RETURN_IF_ERROR(Initialize());
+
+  // Backward propagation.
+  std::vector<Output> dy;
+  while (!ready_.empty()) {
+    // n has collected all gradients.
+    Node* n = ready_.front();
+    ready_.pop_front();
+
+    // dy[i] is the sum of i-th output's backpropped gradients.
+    const int num_y = n->num_outputs();
+    dy.clear();
+    dy.resize(num_y, {nullptr, 0});
+    std::vector<int> no_grad_dy_indices;
+    for (int i = 0; i < num_y; ++i) {
+      TF_RETURN_IF_ERROR(SumGradients({n, i}, &dy[i]));
+      if (dy[i] == NoGradient()) {
+        no_grad_dy_indices.push_back(i);
+      }
+      auto iter = input_nodes_.find({n, i});
+      if (iter != input_nodes_.end()) {
+        // Return gradients for Output in 'grad_outputs_'.
+        (*grad_outputs_)[iter->second] = dy[i];
+      }
+    }
+
+    // Stop backprop if none of the inputs to `n` are in `backprops_'.
+    bool stop_node = true;
+    for (const Edge* e : n->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      if (backprops_.find({e->src(), e->src_output()}) != backprops_.end()) {
+        stop_node = false;
+        break;
+      }
+    }
+
+    if (stop_node) {
+      continue;
+    }
+
+    // Special case: if we find an exit node, process the associated while loop.
+    // Note that ProcessWhileLoop() calls BackpropAlongEdge() if necessary
+    // (which updates ready_), and we skip all the regular processing below
+    // after calling it.
+    if (n->IsExit()) {
+      DCHECK_EQ(dy.size(), 1);
+      TF_RETURN_IF_ERROR(ProcessWhileLoop(n, dy[0]));
+      continue;
+    }
+    // All loop-specific control flow ops should have been handled above
+    DCHECK(!n->IsEnter() && !n->IsNextIteration()) << n->DebugString();
+
+    const int num_no_grad = no_grad_dy_indices.size();
+    if (IsPrimitiveOpWithNoGrad(n->type_string()) || num_no_grad == num_y) {
+      // No grad defined for this op, or all outputs returned 'NoGradient':
+      // Backprop 'NoGradient' along the in edges.
+      for (const Edge* e : n->in_edges()) {
+        if (e->IsControlEdge()) continue;
+        TF_RETURN_IF_ERROR(
+            BackpropAlongEdge(NoGradient(), {e->src(), e->src_output()}));
+      }
+      continue;
+    }
+
+    if (num_no_grad > 0 && num_no_grad < num_y) {
+      // The outputs of 'n' returned a mixture of valid gradients and
+      // 'NoGradient'. Therefore, we need to add 'ZerosLike' nodes for each
+      // 'NoGradient' output before we call the gradient function for 'n'.
+      // TODO(andydavis) If static shapes are known, replace 'ZerosLike' with
+      // zero-filled Constant node of appropriate shape.
+      for (const int dy_index : no_grad_dy_indices) {
+        dy[dy_index] = ops::ZerosLike(scope_, Output(n, dy_index));
+      }
+    }
+
+    // TODO(andydavis) Add option to encapsulate grad function in
+    // SymbolicGradientOp (as opposed to inlining into the graph).
+    std::vector<Output> dx;
+    TF_RETURN_IF_ERROR(CallGradFunction(Operation(n), dy, &dx));
+
+    // Backprop along the in edges.
+    // TODO(andydavis) Find cleaner way to map each grad output returned by
+    // gradient function to the src node/output to which it should be
+    // backpropped. Maybe grad functions can return a vector of Output pairs to
+    // make this association explicit.
+    for (const Edge* e : n->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      size_t dx_index = e->dst_input();
+      if (dx_index >= dx.size()) {
+        return absl::InternalError(absl::StrCat(
+            "Invalid gradient output index: ", dx_index, " size: ", dx.size()));
+      }
+      TF_RETURN_IF_ERROR(
+          BackpropAlongEdge(dx[dx_index], {e->src(), e->src_output()}));
+    }
+  }
+
+  // Check if any input nodes still have pending gradients and have not been
+  // processed yet. This happens if not all outputs of a node are in 'inputs_'.
+  std::unordered_map<Node*, int> requested_grads;
+  for (const Output& nout : inputs_) {
+    if (pending_[nout.node()->id()] > 0) {
+      DCHECK_GT(nout.node()->num_outputs(), 1);
+      int idx = input_nodes_[nout];
+      DCHECK(((*grad_outputs_)[idx].node() == nullptr));
+      TF_RETURN_IF_ERROR(SumGradients(nout, &(*grad_outputs_)[idx]));
+      ++requested_grads[nout.node()];
+    }
+  }
+  for (const auto& p : requested_grads) {
+    int num_requested_inputs = p.first->num_outputs() - pending_[p.first->id()];
+    CHECK_EQ(num_requested_inputs, p.second);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status AddSymbolicGradients(const Scope& scope,
+                                  const std::vector<Output>& outputs,
+                                  const std::vector<Output>& inputs,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs) {
+  SymbolicGradientBuilder builder(scope, ops::GradOpRegistry::Global(), outputs,
+                                  inputs, grad_inputs, grad_outputs);
+  return builder.AddGradients();
+}
+
+absl::Status AddSymbolicGradients(const Scope& scope,
+                                  const std::vector<Output>& outputs,
+                                  const std::vector<Output>& inputs,
+                                  std::vector<Output>* grad_outputs) {
+  std::vector<Output> grad_inputs;
+  grad_inputs.reserve(outputs.size());
+  for (const Output& output : outputs) {
+    grad_inputs.emplace_back(ops::OnesLike(scope, output));
+  }
+  return AddSymbolicGradients(scope, outputs, inputs, grad_inputs,
+                              grad_outputs);
+}
+
+Output NoGradient() { return SymbolicGradientBuilder::NoGradient(); }
+
+}  // end namespace tensorflow

tensorflow/cc/framework/gradients.h

@@ -0,0 +1,54 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_
+#define TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_
+
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+
+/// NOTE: This API is a work in progress and will likely be changing frequently.
+///
+/// Given initial gradients 'grad_inputs' (which represent the symbolic partial
+/// derivatives of some loss function 'L' w.r.t 'outputs'), adds gradient nodes
+/// to the graph associated with 'scope', which compute (and return in
+/// 'grad_outputs') the symbolic partial derivatives of 'L' w.r.t 'inputs'.
+absl::Status AddSymbolicGradients(const Scope& scope,
+                                  const std::vector<Output>& outputs,
+                                  const std::vector<Output>& inputs,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs);
+
+// Same as above, but uses 'OnesLike' for all shapes in
+// 'outputs' as grad_inputs.
+absl::Status AddSymbolicGradients(const Scope& scope,
+                                  const std::vector<Output>& outputs,
+                                  const std::vector<Output>& inputs,
+                                  std::vector<Output>* grad_outputs);
+
+/// Returns a sentinel Output that represents 'no gradient' (i.e. no gradient
+/// flows along some graph edge during backpropagation).
+/// Can be returned in 'grad_outputs' by an invocation of 'AddSymbolicGradients'
+/// (note that gradient flow through an Output can be stopped through the use of
+/// the StopGradient node).
+Output NoGradient();
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_GRADIENTS_H_

tensorflow/cc/framework/gradients_test.cc

@@ -0,0 +1,722 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/equal_graph_def.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Assign;
+using ops::Const;
+using ops::Identity;
+using ops::MatMul;
+using ops::OnesLike;
+using ops::Placeholder;
+using ops::Square;
+using ops::Stack;
+using ops::StopGradient;
+using ops::Unstack;
+using ops::Variable;
+
+// TODO(andydavis) Add more unit tests once more gradient functions are ported.
+class GradientsTest : public ::testing::Test {
+ protected:
+  GradientsTest()
+      : scope_expected_(Scope::NewRootScope()),
+        scope_test_(Scope::NewRootScope()) {}
+
+  void CompareTestAndExpectedGraphs() {
+    GraphDef gdef_test;
+    TF_ASSERT_OK(scope_test_.ToGraphDef(&gdef_test));
+    GraphDef gdef_exp;
+    TF_ASSERT_OK(scope_expected_.ToGraphDef(&gdef_exp));
+    TF_EXPECT_GRAPH_EQ(gdef_exp, gdef_test);
+  }
+
+  Scope scope_expected_;
+  Scope scope_test_;
+};
+
+// Example:
+//      ^             ^
+//    dy|           dx|        (MatMul Gradient Graph)
+//      |             |
+//   MatMul_1      MatMul_2
+//   ^   ^          ^    ^
+//   |   |----------|    |
+//   |        ^          |
+//   |      dz|          |
+//   |        |          |
+//   |     Const_3       |
+//   |                   |
+//   |        ^          |
+//   |       z|          |     (MatMul Forward Graph)
+//   |        |          |
+//   |      MatMul_0     |
+//   |     /        \    |
+//   |    ^          ^   |
+//   |    |          |   |
+//   |---x|         y|---|
+//        |          |
+//        |          |
+//      Const_0   Const_1
+//
+
+TEST_F(GradientsTest, OneMatMul) {
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto x = Const(scope, {{1.0, 2.0}, {3.0, 4.0}});
+    auto y = Const(scope, {{1.0, 0.0}, {0.0, 1.0}});
+    auto z = MatMul(scope, x, y);
+    TF_ASSERT_OK(scope.status());
+    CHECK_NOTNULL(z.node());
+
+    if (expected) {
+      // Construct backward graph.
+      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      auto dx = MatMul(scope, dz, y, MatMul::TransposeB(true));
+      auto dy = MatMul(scope, x, dz, MatMul::TransposeA(true));
+    } else {
+      // Call AddSymbolicGradients.
+      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(
+          AddSymbolicGradients(scope, {z}, {x, y}, {dz}, &grad_outputs));
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, OneMatMul_InferGradInputs) {
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto x = Const(scope, {{1.0, 2.0}, {3.0, 4.0}});
+    auto y = Const(scope, {{1.0, 0.0}, {0.0, 1.0}});
+    auto z = MatMul(scope, x, y);
+    TF_ASSERT_OK(scope.status());
+    CHECK_NOTNULL(z.node());
+
+    if (expected) {
+      // Construct backward graph.
+      // The gradients function adds a OnesLike to create a dz of ones with the
+      // shape of z.
+      auto dz = OnesLike(scope, z);
+      auto dx = MatMul(scope, dz, y, MatMul::TransposeB(true));
+      auto dy = MatMul(scope, x, dz, MatMul::TransposeA(true));
+    } else {
+      // Call AddSymbolicGradients.
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(AddSymbolicGradients(scope, {z}, {x, y}, &grad_outputs));
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, TwoMatMuls_Chained) {
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto u = Const(scope, {{1.0, 2.0}, {3.0, 4.0}});
+    auto v = Const(scope, {{1.0, 0.0}, {0.0, 1.0}});
+    auto x = MatMul(scope, u, v);
+
+    auto y = Const(scope, {{1.0, 0.0}, {0.0, 1.0}});
+    auto z = MatMul(scope, x, y);
+
+    TF_ASSERT_OK(scope.status());
+    CHECK_NOTNULL(z.node());
+
+    if (expected) {
+      // Construct backward graph.
+      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      auto dx = MatMul(scope, dz, y, MatMul::TransposeB(true));
+      auto dy = MatMul(scope, x, dz, MatMul::TransposeA(true));
+
+      auto du = MatMul(scope, dx, v, MatMul::TransposeB(true));
+      auto dv = MatMul(scope, u, dx, MatMul::TransposeA(true));
+    } else {
+      // Call AddSymbolicGradients.
+      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(
+          AddSymbolicGradients(scope, {z}, {u, v}, {dz}, &grad_outputs));
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, TwoMatMuls_Independent) {
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto t = Const(scope, {{1.0, 2.0}, {3.0, 4.0}});
+    auto u = Const(scope, {{1.0, 0.0}, {0.0, 1.0}});
+    auto v = MatMul(scope, t, u);
+    TF_ASSERT_OK(scope.status());
+    CHECK_NOTNULL(v.node());
+
+    auto x = Const(scope, {{5.0, 6.0}, {7.0, 8.0}});
+    auto y = Const(scope, {{1.0, 0.0}, {0.0, 1.0}});
+    auto z = MatMul(scope, x, y);
+    TF_ASSERT_OK(scope.status());
+    CHECK_NOTNULL(z.node());
+
+    if (expected) {
+      // Construct backward graph.
+      auto dv = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      auto dt = MatMul(scope, dv, u, MatMul::TransposeB(true));
+      auto du = MatMul(scope, t, dv, MatMul::TransposeA(true));
+
+      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      auto dx = MatMul(scope, dz, y, MatMul::TransposeB(true));
+      auto dy = MatMul(scope, x, dz, MatMul::TransposeA(true));
+    } else {
+      // Call AddSymbolicGradients.
+      auto dv = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(AddSymbolicGradients(scope, {v, z}, {t, u, x, y}, {dv, dz},
+                                        &grad_outputs));
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, StackUnstack_Chained) {
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto a = Const(scope, 1, {4, 2});
+    auto b = Const(scope, 2, {4, 2});
+    auto c = Const(scope, 3, {4, 2});
+
+    auto pack = Stack(scope, {a, b, c});
+    auto unpack = Unstack(scope, pack.output, 3);
+    TF_ASSERT_OK(scope.status());
+
+    // Construct grad inputs.
+    auto dx = Const(scope, 4, {4, 2});
+    auto dy = Const(scope, 5, {4, 2});
+    auto dz = Const(scope, 6, {4, 2});
+
+    if (expected) {
+      // Construct backward graph.
+      auto unpack_grad = Stack(scope, {dx, dy, dz});
+      auto pack_grad = Unstack(scope, unpack_grad.output, 3);
+    } else {
+      // Call AddSymbolicGradients.
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(AddSymbolicGradients(scope, unpack.output, {a, b, c},
+                                        {dx, dy, dz}, &grad_outputs));
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, StackUnstack_StopBackprop) {
+  // Tests that backprop stops before calculating gradients for Stack (because
+  // only gradients w.r.t the output of Stack are requested).
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto a = Const(scope, 1, {4, 2});
+    auto b = Const(scope, 2, {4, 2});
+    auto c = Const(scope, 3, {4, 2});
+
+    auto pack = Stack(scope, {a, b, c});
+    auto unpack = Unstack(scope, pack.output, 3);
+    TF_ASSERT_OK(scope.status());
+
+    // Construct grad inputs.
+    auto dx = Const(scope, 4, {4, 2});
+    auto dy = Const(scope, 5, {4, 2});
+    auto dz = Const(scope, 6, {4, 2});
+
+    if (expected) {
+      // Construct backward graph.
+      // NOTE: We should only expect the grad function for unpack in the
+      // gradients graph, based on the requested grad outputs.
+      auto unpack_grad = Stack(scope, {dx, dy, dz});
+    } else {
+      // Call AddSymbolicGradients.
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(AddSymbolicGradients(scope, unpack.output, {pack},
+                                        {dx, dy, dz}, &grad_outputs));
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, StackUnstack_SubsetOfUnstackOutputs) {
+  // Constructs an unstack with three outputs, and takes the gradient with
+  // respect to only two of the outputs. Tests that the output gradients are
+  // computed.
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto c = Const(scope, 1, {3, 4, 2});
+    auto unpack = Unstack(scope, c, 3);
+    auto x = Identity(scope, unpack.output[0]);
+    auto y = Identity(scope, unpack.output[1]);
+    auto z = Identity(scope, unpack.output[2]);
+    TF_ASSERT_OK(scope.status());
+
+    // Construct grad inputs.
+    auto dy = Const(scope, 4, {4, 2});
+    auto dz = Const(scope, 5, {4, 2});
+
+    if (expected) {
+      // Construct backward graph.
+      auto g1 = Identity(scope, dy);
+      auto g2 = Identity(scope, dz);
+    } else {
+      // Call AddSymbolicGradients.
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(AddSymbolicGradients(scope, {y, z},
+                                        {unpack.output[1], unpack.output[2]},
+                                        {dy, dz}, &grad_outputs));
+      ASSERT_EQ(grad_outputs.size(), 2);
+      EXPECT_TRUE(grad_outputs[0].node() != nullptr);
+      EXPECT_TRUE(grad_outputs[1].node() != nullptr);
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
+TEST_F(GradientsTest, DependentGradOutputs) {
+  // Tests that dependent gradients (in this case the gradients w.r.t to the
+  // output and one input of MatMul) are computed properly.
+
+  // Create two chained MatMul ops.
+  auto u = Const(scope_test_, {{2}});
+  auto v = Const(scope_test_, {{3}});
+  auto x = MatMul(scope_test_, u, v);
+
+  auto y = Const(scope_test_, {{4}});
+  auto z = MatMul(scope_test_, x, y);
+
+  TF_ASSERT_OK(scope_test_.status());
+  CHECK_NOTNULL(z.node());
+
+  // Call AddSymbolicGradients with '5' as initial gradients for 'dz'.
+  // The gradient w.r.t to 'v' (returned in grad_outputs[0]) is dependent on
+  // the gradient w.r.t. to 'x' (returned in grad_outputs[1]).
+  auto dz = Const(scope_test_, {{5}});
+  std::vector<Output> grad_outputs;
+  TF_ASSERT_OK(
+      AddSymbolicGradients(scope_test_, {z}, {v, x}, {dz}, &grad_outputs));
+
+  std::vector<Tensor> outputs;
+  test::GetTensors(scope_test_, {grad_outputs[0], grad_outputs[1]}, &outputs);
+
+  // The gradients w.r.t to 'dz' are passed into AddSymbolicGradients as '5'.
+  // Since z = MatMul(x, y), the gradients w.r.t 'x' are computed as:
+  //   'dx' = 5 * 'y' = 5 * 4 = 20.
+  // Since x = MatMul(u, v), the gradients w.r.t. 'v' are computed as:
+  //   'dv' = 'dx' * 'u' = 20 * 2 = 40.
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({40}, {1, 1}));
+  test::ExpectTensorEqual<int>(outputs[1], test::AsTensor<int>({20}, {1, 1}));
+}
+
+TEST_F(GradientsTest, MultipleNodeOutputGrads) {
+  // Tests that gradients for multiple outputs of the same node are returned.
+  auto x = Const(scope_test_, 1, {3, 4, 2});
+  auto unpack = Unstack(scope_test_, x, 3);
+  auto pack = Stack(scope_test_, unpack.output);
+
+  // clang-format off
+  auto dx = Const(scope_test_, {40, 41, 42, 43, 44, 45, 46, 47,
+                                50, 51, 52, 53, 55, 55, 56, 57,
+                                60, 61, 62, 63, 66, 66, 66, 67},
+                               {3, 4, 2});
+  // clang-format on
+
+  std::vector<Output> grad_outputs;
+  TF_ASSERT_OK(AddSymbolicGradients(scope_test_, {pack}, unpack.output, {dx},
+                                    &grad_outputs));
+
+  std::vector<Tensor> outputs;
+  test::GetTensors(scope_test_,
+                   {grad_outputs[0], grad_outputs[1], grad_outputs[2]},
+                   &outputs);
+
+  test::ExpectTensorEqual<int>(
+      outputs[0],
+      test::AsTensor<int>({40, 41, 42, 43, 44, 45, 46, 47}, {4, 2}));
+  test::ExpectTensorEqual<int>(
+      outputs[1],
+      test::AsTensor<int>({50, 51, 52, 53, 55, 55, 56, 57}, {4, 2}));
+  test::ExpectTensorEqual<int>(
+      outputs[2],
+      test::AsTensor<int>({60, 61, 62, 63, 66, 66, 66, 67}, {4, 2}));
+}
+
+TEST_F(GradientsTest, UnreachableEdgeGradOneOutput) {
+  auto x = Variable(scope_test_, {2, 3}, DT_DOUBLE);
+  auto x_const = Const(scope_test_, {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}});
+  auto x_assign = Assign(scope_test_, x, x_const);
+
+  auto y = Variable(scope_test_, {3, 1}, DT_DOUBLE);
+  auto y_const = Const(scope_test_, {{1.0}, {2.0}, {3.0}});
+  auto y_assign = Assign(scope_test_, y, y_const);
+
+  auto m = MatMul(scope_test_, x, y);
+
+  auto z = Variable(scope_test_, {1, 3}, DT_DOUBLE);
+  auto z_const = Const(scope_test_, {{9.0, 10.0, 11.0}});
+  auto z_assign = Assign(scope_test_, z, z_const);
+
+  auto diff_m = Const(scope_test_, {{0.5}, {0.5}});
+
+  std::vector<Output> grad_outputs;
+  TF_ASSERT_OK(
+      AddSymbolicGradients(scope_test_, {m}, {y}, {diff_m}, &grad_outputs));
+
+  std::vector<Tensor> outputs;
+  test::GetTensors(scope_test_, {x_assign, y_assign, z_assign},
+                   {grad_outputs[0]}, &outputs);
+  // dz/dy = xT * diff_m
+  test::ExpectTensorNear<double>(
+      outputs[0], test::AsTensor<double>({2.5, 3.5, 4.5}, {3, 1}), 1e-5);
+}
+
+TEST_F(GradientsTest, UnreachableEdgeGradTwoOutputs) {
+  auto x = Variable(scope_test_, {2, 3}, DT_DOUBLE);
+  auto x_const = Const(scope_test_, {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}});
+  auto x_assign = Assign(scope_test_, x, x_const);
+
+  auto y = Variable(scope_test_, {3, 1}, DT_DOUBLE);
+  auto y_const = Const(scope_test_, {{1.0}, {2.0}, {3.0}});
+  auto y_assign = Assign(scope_test_, y, y_const);
+
+  auto m1 = MatMul(scope_test_, x, y);
+
+  auto z = Variable(scope_test_, {1, 3}, DT_DOUBLE);
+  auto z_const = Const(scope_test_, {{9.0, 10.0, 11.0}});
+  auto z_assign = Assign(scope_test_, z, z_const);
+
+  auto m2 = MatMul(scope_test_, y, z);
+
+  auto dm1 = Const(scope_test_, {{0.5}, {0.5}});
+  auto dm2 =
+      Const(scope_test_, {{0.5, 0.5, 0.5}, {0.6, 0.7, 0.8}, {0.6, 0.7, 0.9}});
+
+  std::vector<Output> grad_outputs;
+  TF_ASSERT_OK(AddSymbolicGradients(scope_test_, {m1, m2}, {y}, {dm1, dm2},
+                                    &grad_outputs));
+
+  std::vector<Tensor> outputs;
+  test::GetTensors(scope_test_, {x_assign, y_assign, z_assign},
+                   {grad_outputs[0]}, &outputs);
+
+  // The gradients from m1 and m2 will be summed to compute the gradient
+  // w.r.t y:
+  // dz/dy = xT * dm1 + dm2 * zT
+  test::ExpectTensorNear<double>(
+      outputs[0], test::AsTensor<double>({17.5, 24.7, 26.8}, {3, 1}), 1e-5);
+}
+
+TEST_F(GradientsTest, UnreachableInput) {
+  auto x = Const(scope_test_, {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}});
+  auto y = Const(scope_test_, {{1.0}, {2.0}, {3.0}});
+  auto z = Const(scope_test_.WithOpName("z"), {{9.0, 10.0, 11.0}});
+
+  auto m1 = MatMul(scope_test_, x, y);
+  auto m2 = MatMul(scope_test_, y, z);
+  auto dm1 = Const(scope_test_, {{0.5}, {0.5}});
+
+  // From m1, z is unreachable, so an error status should be returned.
+  //   m2  m1
+  //   |   |
+  //   *   *
+  //  / \ / \
+  // z   y   x
+  std::vector<Output> grad_outputs;
+  absl::Status status =
+      AddSymbolicGradients(scope_test_, {m1}, {z}, {dm1}, &grad_outputs);
+  EXPECT_EQ(status.code(), error::INVALID_ARGUMENT);
+  EXPECT_EQ(status.message(),
+            "Cannot compute the partial derivative"
+            " for node 'z' as it's unreachable from the output node(s).");
+}
+
+TEST_F(GradientsTest, DependentOutputs) {
+  auto x = Placeholder(scope_test_, DT_FLOAT);
+  auto y0 = Square(scope_test_, x);
+  auto y1 = Square(scope_test_, y0);
+  auto y2 = Square(scope_test_, y1);
+  // Requesting the gradients for y0 and y2 should return the sum of their
+  // individual gradients.
+  std::vector<Output> grad_outputs;
+  TF_EXPECT_OK(AddSymbolicGradients(scope_test_, {y0, y2}, {x}, &grad_outputs));
+  ClientSession session(scope_test_);
+  std::vector<Tensor> grad_result;
+  TF_EXPECT_OK(session.Run({{x, {3.0f}}}, grad_outputs, &grad_result));
+  EXPECT_EQ(grad_result.size(), 1);
+  EXPECT_EQ(grad_result[0].NumElements(), 1);
+  EXPECT_EQ(grad_result[0].flat<float>()(0), 17502.0f);
+}
+
+TEST_F(GradientsTest, MultiOutputNodeDependentOutputs) {
+  auto x = Placeholder(scope_test_, DT_FLOAT);
+  auto y0 = Square(scope_test_, x);
+  // y1, y2, and y3 all use y0. This means the backwards pass will need to wait
+  // for the gradient for all three.
+  auto y1 = Square(scope_test_, y0);
+  auto y2 = Square(scope_test_, y0);
+  auto y3 = Square(scope_test_, y2);
+  std::vector<Output> grad_outputs;
+  // By requesting y0, y1, and y3 we test that the computation correctly waits
+  // for all the points in backprop where gradients need to be summed from
+  // multiple branches.
+  TF_EXPECT_OK(
+      AddSymbolicGradients(scope_test_, {y0, y1, y3}, {x}, &grad_outputs));
+  ClientSession session(scope_test_);
+  std::vector<Tensor> grad_result;
+  TF_EXPECT_OK(session.Run({{x, {3.0f}}}, grad_outputs, &grad_result));
+  EXPECT_EQ(grad_result.size(), 1);
+  EXPECT_EQ(grad_result[0].NumElements(), 1);
+  EXPECT_EQ(grad_result[0].flat<float>()(0), 17610.0f);
+}
+
+TEST_F(GradientsTest, AddSymbolicGradientsTest) {
+  Scope scope = Scope::NewRootScope();
+  for (int cnt = 0; cnt < 100; ++cnt) {
+    int N = 5 + rand() % 10;
+    // Construct forward graph.
+    OutputList inputs;
+    for (int i = 0; i < N; ++i) {
+      auto a = Const(scope, i, {1});
+      inputs.push_back(a);
+    }
+
+    auto pack = Stack(scope, inputs);
+    TF_ASSERT_OK(scope.status());
+
+    // Construct grad inputs.
+    OutputList output_grads;
+    Tensor ts(DT_INT32, {N, 1});
+    auto v = ts.matrix<int32_t>();
+    for (int i = 0; i < N; ++i) {
+      v(i, 0) = i;
+    }
+    auto dy = Const(scope, ts);
+    output_grads.push_back(dy);
+    // Call AddSymbolicGradients.
+    std::vector<Output> grad_outputs;
+    TF_ASSERT_OK(AddSymbolicGradients(scope, {pack.output}, inputs,
+                                      output_grads, &grad_outputs));
+    ClientSession session((scope));
+    std::vector<Tensor> in_grad;
+    TF_ASSERT_OK(session.Run(grad_outputs, &in_grad));
+    for (int i = 0; i < N; ++i) {
+      test::ExpectTensorEqual<int>(in_grad[i], test::AsTensor<int>({i}, {1}));
+    }
+  }
+}
+
+// StopGradientSingleOutputMultiEdgeTest tests combinations of valid and
+// 'NoGradient' (induced by StopGradient op) returned along multiple edges from
+// a single nodes output.
+class StopGradientSingleOutputMultiEdgeTest : public ::testing::Test {
+ protected:
+  StopGradientSingleOutputMultiEdgeTest() : scope_(Scope::NewRootScope()) {}
+
+  void CheckGrad(const std::vector<bool>& stop_outputs,
+                 const Tensor& expected_grad) {
+    CHECK_EQ(3, stop_outputs.size());
+
+    auto x = Const(scope_, {{1, 0}, {0, 1}});
+    auto y = Const(scope_, {{1, 0}, {0, 1}});
+    auto z = MatMul(scope_, x, y);
+
+    // Create three output going edges from 'z'.
+    // Add StopGradients according to 'stop_outputs'.
+    auto out0 = stop_outputs[0]
+                    ? StopGradient(scope_, (Identity(scope_, z))).output
+                    : Identity(scope_, z).output;
+    auto out1 = stop_outputs[1]
+                    ? StopGradient(scope_, (Identity(scope_, z))).output
+                    : Identity(scope_, z).output;
+    auto out2 = stop_outputs[2]
+                    ? StopGradient(scope_, (Identity(scope_, z))).output
+                    : Identity(scope_, z).output;
+
+    auto g0 = Const(scope_, {{1, 2}, {3, 4}});
+    auto g1 = Const(scope_, {{5, 6}, {7, 8}});
+    auto g2 = Const(scope_, {{9, 10}, {11, 12}});
+
+    // Call AddSymbolicGradients and compare against 'expected_grad'.
+    std::vector<Output> grad_outputs;
+    TF_EXPECT_OK(AddSymbolicGradients(scope_, {out0, out1, out2}, {z},
+                                      {g0, g1, g2}, &grad_outputs));
+
+    if (expected_grad.NumElements() > 0) {
+      Tensor output;
+      test::GetTensor(scope_, grad_outputs[0], &output);
+      test::ExpectTensorEqual<int>(output, expected_grad);
+    } else {
+      EXPECT_EQ(NoGradient(), grad_outputs[0]);
+    }
+  }
+
+  Scope scope_;
+};
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, ValidGradAllEdges) {
+  CheckGrad({false, false, false},
+            test::AsTensor<int>({15, 18, 21, 24}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradFirstEdge) {
+  CheckGrad({true, false, false},
+            test::AsTensor<int>({14, 16, 18, 20}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradSecondEdge) {
+  CheckGrad({false, true, false},
+            test::AsTensor<int>({10, 12, 14, 16}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradThirdEdge) {
+  CheckGrad({false, false, true}, test::AsTensor<int>({6, 8, 10, 12}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradFirstAndSecondEdges) {
+  CheckGrad({true, true, false}, test::AsTensor<int>({9, 10, 11, 12}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradSecondAndThirdEdges) {
+  CheckGrad({false, true, true}, test::AsTensor<int>({1, 2, 3, 4}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradFirstAndThirdEdges) {
+  CheckGrad({true, false, true}, test::AsTensor<int>({5, 6, 7, 8}, {2, 2}));
+}
+
+TEST_F(StopGradientSingleOutputMultiEdgeTest, StopGradAllEdges) {
+  CheckGrad({true, true, true}, Tensor());
+}
+
+// StopGradientMultiOutputTest tests combinations of valid and 'NoGradient'
+// (induced by StopGradient op) returned along a single nodes multiple outputs.
+class StopGradientMultiOutputTest : public ::testing::Test {
+ protected:
+  StopGradientMultiOutputTest() : scope_(Scope::NewRootScope()) {}
+
+  void CheckGrad(const std::vector<bool>& stop_outputs,
+                 const Tensor& expected_grad) {
+    CHECK_EQ(3, stop_outputs.size());
+    auto x = ops::Const(scope_, 1, {3, 2, 4});
+    auto y = Unstack(scope_, x, 3);
+    TF_ASSERT_OK(scope_.status());
+
+    // Add StopGradients according to 'stop_outputs'.
+    auto out0 =
+        stop_outputs[0] ? StopGradient(scope_, y.output[0]) : y.output[0];
+    auto out1 =
+        stop_outputs[1] ? StopGradient(scope_, y.output[1]) : y.output[1];
+    auto out2 =
+        stop_outputs[2] ? StopGradient(scope_, y.output[2]) : y.output[2];
+
+    auto g0 = Const(scope_, {1, 2, 3, 4, 5, 6, 7, 8}, {2, 4});
+    auto g1 = Const(scope_, {9, 10, 11, 12, 13, 14, 15, 16}, {2, 4});
+    auto g2 = Const(scope_, {17, 18, 19, 20, 21, 22, 23, 24}, {2, 4});
+
+    // Call AddSymbolicGradients and compare against 'expected_grad'.
+    std::vector<Output> grad_outputs;
+    TF_EXPECT_OK(AddSymbolicGradients(scope_, {out0, out1, out2}, {x},
+                                      {g0, g1, g2}, &grad_outputs));
+
+    if (expected_grad.NumElements() > 0) {
+      Tensor output;
+      test::GetTensor(scope_, grad_outputs[0], &output);
+      test::ExpectTensorEqual<int>(output, expected_grad);
+    } else {
+      EXPECT_EQ(NoGradient(), grad_outputs[0]);
+    }
+  }
+
+  Scope scope_;
+};
+
+TEST_F(StopGradientMultiOutputTest, ValidGradAllOutputs) {
+  // clang-format off
+  CheckGrad({false, false, false}, test::AsTensor<int>(
+    {1, 2, 3, 4, 5, 6, 7, 8,
+     9, 10, 11, 12, 13, 14, 15, 16,
+     17, 18, 19, 20, 21, 22, 23, 24},
+    {3, 2, 4}));
+  // clang-format on
+}
+
+TEST_F(StopGradientMultiOutputTest, StopGradFirstOutput) {
+  // clang-format off
+  CheckGrad({true, false, false}, test::AsTensor<int>(
+    {0, 0, 0, 0, 0, 0, 0, 0,
+     9, 10, 11, 12, 13, 14, 15, 16,
+     17, 18, 19, 20, 21, 22, 23, 24},
+    {3, 2, 4}));
+  // clang-format on
+}
+
+TEST_F(StopGradientMultiOutputTest, StopGradSecondOutput) {
+  // clang-format off
+  CheckGrad({false, true, false}, test::AsTensor<int>(
+    {1, 2, 3, 4, 5, 6, 7, 8,
+     0, 0, 0, 0, 0, 0, 0, 0,
+     17, 18, 19, 20, 21, 22, 23, 24},
+    {3, 2, 4}));
+  // clang-format on
+}
+
+TEST_F(StopGradientMultiOutputTest, StopGradThirdOutput) {
+  // clang-format off
+  CheckGrad({false, false, true}, test::AsTensor<int>(
+    {1, 2, 3, 4, 5, 6, 7, 8,
+     9, 10, 11, 12, 13, 14, 15, 16,
+     0, 0, 0, 0, 0, 0, 0, 0},
+    {3, 2, 4}));
+  // clang-format on
+}
+
+TEST_F(StopGradientMultiOutputTest, StopGradFirstAndThirdOutputs) {
+  // clang-format off
+  CheckGrad({true, false, true}, test::AsTensor<int>(
+    {0, 0, 0, 0, 0, 0, 0, 0,
+     9, 10, 11, 12, 13, 14, 15, 16,
+     0, 0, 0, 0, 0, 0, 0, 0},
+    {3, 2, 4}));
+  // clang-format on
+}
+
+TEST_F(StopGradientMultiOutputTest, StopAllOutputs) {
+  CheckGrad({true, true, true}, Tensor());
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/framework/ops.cc

@@ -0,0 +1,112 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/core/lib/hash/hash.h"
+
+namespace tensorflow {
+
+Operation::Operation(Node* n) : inputs_(GetInputs(n)), node_(n) {}
+
+Output Operation::input(int32_t i) const {
+  CHECK_NOTNULL(node_);
+  CHECK_GE(i, 0);
+  CHECK_LT(i, node_->num_inputs());
+  // Handle the case where the input was unknown at the time this
+  // Operation was constructed.
+  if (inputs_[i].first == nullptr && inputs_[i].second == -1) {
+    for (const Edge* e : node_->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      if (e->dst_input() == i) {
+        return Output(e->src(), e->src_output());
+      }
+    }
+  }
+  return Output(inputs_[i].first, inputs_[i].second);
+}
+
+Output Operation::output(int32_t i) const {
+  CHECK_NOTNULL(node_);
+  CHECK_GE(i, 0);
+  CHECK_LT(i, node_->num_outputs());
+  return Output(node_, i);
+}
+
+uint64_t Operation::hash(int32_t index) const {
+  return ::tensorflow::Hash64(reinterpret_cast<const char*>(&node_),
+                              sizeof(Node*), index);
+}
+
+Operation::Inputs Operation::GetInputs(Node* node) {
+  Operation::Inputs inputs;
+  if (node != nullptr) {
+    inputs.resize(node->num_inputs(), {nullptr, -1});
+    for (const Edge* e : node->in_edges()) {
+      if (e->IsControlEdge()) continue;
+      inputs[e->dst_input()] = std::make_pair(e->src(), e->src_output());
+    }
+  }
+  return inputs;
+}
+
+Input::Initializer::Initializer(
+    const std::initializer_list<Input::Initializer>& v) {
+  if (v.size() < 1) {
+    // Empty initializer list defaults to float tensor with shape (0,)
+    tensor = Tensor(DT_FLOAT, TensorShape{0});
+    return;
+  }
+  auto const& first = *v.begin();
+  // Check to make sure that the constituent Initializers are all the same
+  // type and same shape.
+  for (auto const& e : v) {
+    if (e.tensor.dtype() != first.tensor.dtype()) {
+      status = absl::InvalidArgumentError(
+          "Initializer list components should all have the same type");
+      return;
+    }
+    if (!TensorShape{e.tensor.shape()}.IsSameSize(
+            TensorShape{first.tensor.shape()})) {
+      status = absl::InvalidArgumentError(
+          "Initializer list components should all have the same shape");
+      return;
+    }
+  }
+
+  // Form the new shape.
+  TensorShape shape{static_cast<int64_t>(v.size())};
+  shape.AppendShape(TensorShape{first.tensor.shape()});
+
+  Tensor t(first.tensor.dtype(), shape);
+
+  // Collate the constituent Tensors.
+  size_t offset = 0;
+  for (auto const& e : v) {
+    Tensor elem = e.tensor;
+    if (first.tensor.dtype() == DT_STRING) {
+      for (int i = 0; i < elem.NumElements(); ++i) {
+        t.flat<tstring>()(offset + i) = elem.flat<tstring>()(i);
+      }
+      offset += elem.NumElements();
+    } else {
+      std::copy_n(elem.tensor_data().data(), elem.TotalBytes(),
+                  const_cast<char*>(t.tensor_data().data()) + offset);
+      offset += elem.TotalBytes();
+    }
+  }
+  tensor = t;
+}
+
+}  // namespace tensorflow

tensorflow/cc/framework/ops.h

@@ -0,0 +1,306 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_OPS_H_
+#define TENSORFLOW_CC_FRAMEWORK_OPS_H_
+
+#include <string>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/lib/hash/hash.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+
+namespace tensorflow {
+
+/// @defgroup core Core Tensorflow API
+
+class Output;
+
+/// @addtogroup core
+/// @{
+
+/// Represents a node in the computation graph.
+class Operation {
+ public:
+  Operation() : node_(nullptr) {}
+  explicit Operation(Node* n);
+
+  int32_t num_inputs() const { return node_->num_inputs(); }
+  DataType input_type(int32_t o) const { return node_->input_type(o); }
+  Output input(int32_t i) const;
+
+  int32_t num_outputs() const { return node_->num_outputs(); }
+  DataType output_type(int32_t o) const { return node_->output_type(o); }
+  Output output(int32_t i) const;
+
+  Node* node() const { return node_; }
+
+  uint64_t hash(int32_t index) const;
+
+  bool operator==(const Operation& other) const { return node_ == other.node_; }
+
+ private:
+  typedef std::vector<std::pair<Node*, int32_t>> Inputs;
+  static Inputs GetInputs(Node* node);
+
+  Inputs inputs_;
+  Node* node_;
+};
+
+/// Represents a tensor value produced by an Operation.
+class Output {
+ public:
+  Output() = default;
+  explicit Output(Node* n) : op_(n) {}
+  Output(Node* n, int32_t index) : op_(n), index_(index) {}
+  Output(const Operation& op, int32_t index) : op_(op), index_(index) {}
+
+  Operation op() const { return op_; }
+  Node* node() const { return op().node(); }
+  int32_t index() const { return index_; }
+  DataType type() const { return op_.output_type(index_); }
+  std::string name() const {
+    return absl::StrCat(node()->name(), ":", index());
+  }
+  bool operator==(const Output& other) const {
+    return op_ == other.op_ && index_ == other.index_;
+  }
+
+  uint64_t hash() const { return op_.hash(index_); }
+
+ private:
+  Operation op_ = Operation(nullptr);
+  int32_t index_ = 0;
+};
+
+/// Hash class that can be used for e.g. storing Outputs in an unordered_map
+struct OutputHash {
+  std::size_t operator()(const Output& output) const {
+    return Hash64Combine(std::hash<Node*>()(output.node()),
+                         std::hash<int32_t>()(output.index()));
+  }
+};
+
+/// Represents a tensor value that can be used as an operand to an Operation.
+class Input {
+ public:
+  /// Initializer enables constructing an Input object from various kinds of C++
+  /// constants such as simple primitive constants and nested initializer lists
+  /// representing a multi-dimensional array. Initializer constructors are all
+  /// templates, so the aforementioned kinds of C++ constants can be used to
+  /// construct an Initializer. Initializer stores the value it got constructed
+  /// with in a Tensor object.
+  struct Initializer {
+    /// Construct from a scalar value of an arithmetic type or a type that can
+    /// be converted to a string (eg. a string literal).
+    template <typename T, typename = typename std::enable_if<
+                              std::is_arithmetic<T>::value ||
+                              std::is_convertible<T, std::string>::value>::type>
+    Initializer(const T& v) {  // NOLINT(runtime/explicit)
+      typedef typename RealType<T>::type RealT;
+      Tensor t(DataTypeToEnum<RealT>::v(), TensorShape());
+      t.flat<RealT>()(0) = RealT(v);
+      tensor = t;
+    }
+
+    Initializer(const Tensor& t) : tensor(t) {}  // NOLINT(runtime/explicit)
+
+    /// Construct from a scalar value and an explicit shape
+    template <typename T, typename = typename std::enable_if<
+                              std::is_arithmetic<T>::value ||
+                              std::is_convertible<T, std::string>::value>::type>
+    Initializer(const T& v, const TensorShape& shape) {
+      typedef typename RealType<T>::type RealT;
+      Tensor t(DataTypeToEnum<RealT>::v(), shape);
+      for (int64_t i = 0; i < t.NumElements(); ++i) {
+        t.flat<RealT>()(i) = RealT(v);
+      }
+      tensor = t;
+    }
+
+    /// Construct from a initializer list of scalars (a one-dimensional tensor).
+    template <typename T, typename = typename std::enable_if<
+                              std::is_arithmetic<T>::value ||
+                              std::is_convertible<T, std::string>::value>::type>
+    Initializer(
+        const std::initializer_list<T>& v) {  // NOLINT(runtime/explicit)
+      typedef typename RealType<T>::type RealT;
+      Tensor t(DataTypeToEnum<RealT>::v(),
+               TensorShape{static_cast<int>(v.size())});
+      std::copy_n(v.begin(), v.size(), t.flat<RealT>().data());
+      tensor = t;
+    }
+
+    /// Construct from a initializer list of scalars and an explicit shape.
+    template <typename T, typename = typename std::enable_if<
+                              std::is_arithmetic<T>::value ||
+                              std::is_convertible<T, std::string>::value>::type>
+    Initializer(const std::initializer_list<T>& v, const TensorShape& shape) {
+      typedef typename RealType<T>::type RealT;
+      Tensor t(DataTypeToEnum<RealT>::v(), shape);
+      if (t.NumElements() != static_cast<int64_t>(v.size())) {
+        status = absl::InvalidArgumentError(absl::StrCat(
+            "Cannot construct a tensor with ", t.NumElements(),
+            " from an initializer list with ", v.size(), " elements"));
+        return;
+      }
+      std::copy_n(v.begin(), v.size(), t.flat<RealT>().data());
+      tensor = t;
+    }
+
+    /// Construct a multi-dimensional tensor from a nested initializer
+    /// list. Note that C++ syntax allows nesting of arbitrarily typed
+    /// initializer lists, so such invalid initializers cannot be disallowed at
+    /// compile time. This function performs checks to make sure that the nested
+    /// initializer list is indeed a valid multi-dimensional tensor.
+    Initializer(const std::initializer_list<Initializer>& v);
+
+    // START_SKIP_DOXYGEN
+    template <typename T, bool = std::is_convertible<T, std::string>::value>
+    struct RealType {
+      typedef tstring type;
+    };
+
+    template <typename T>
+    struct RealType<T, false> {
+      typedef T type;
+    };
+    // END_SKIP_DOXYGEN
+
+    TensorProto AsTensorProto() {
+      TensorProto tensor_proto;
+      if (tensor.NumElements() > 1) {
+        tensor.AsProtoTensorContent(&tensor_proto);
+      } else {
+        tensor.AsProtoField(&tensor_proto);
+      }
+      return tensor_proto;
+    }
+
+    absl::Status status;
+    Tensor tensor;
+  };
+
+  /// All of Input's constructors are implicit. Input can be implicitly
+  /// constructed from the following objects :
+  /// * Output: This is so that the output of an Operation can be directly used
+  ///   as the input to a op wrapper, which takes Inputs.
+  /// * A scalar, or a multi-dimensional tensor specified as a recursive
+  ///   initializer list. This enables directly passing constants as
+  ///   inputs to op wrappers.
+  /// * A Tensor object.
+  Input(const Output& o) : output_(o) {}  // NOLINT(runtime/explicit)
+
+  template <typename T, typename = typename std::enable_if<
+                            std::is_arithmetic<T>::value ||
+                            std::is_convertible<T, std::string>::value>::type>
+  Input(const T& v)  // NOLINT(runtime/explicit)
+      : Input(Initializer(v)) {}
+
+  Input(const Initializer& init)  // NOLINT(runtime/explicit)
+      : status_(init.status),
+        tensor_(init.tensor) {}
+
+  Input(const Tensor& t)  // NOLINT(runtime/explicit)
+      : status_(absl::OkStatus()), tensor_(t) {}
+
+  Input(const std::initializer_list<Initializer>&
+            init) {  // NOLINT(runtime/explicit)
+    for (const auto& i : init) {
+      if (!i.status.ok()) {
+        status_ = i.status;
+        return;
+      }
+    }
+    tensor_ = Initializer(init).tensor;
+  }
+
+  /// Constructor specifying a node name, index and datatype. This should only
+  /// be used for specifying a backward edge, needed by control flow.
+  Input(const std::string& name, int32_t i, DataType dt)
+      : node_name_(name), index_(i), data_type_(dt) {}
+
+  Node* node() const { return output_.node(); }
+  std::string node_name() const { return node_name_; }
+  int32_t index() const {
+    return node_name_.empty() ? output_.index() : index_;
+  }
+  DataType data_type() const { return data_type_; }
+  absl::Status status() const { return status_; }
+  const Tensor& tensor() const { return tensor_; }
+
+ private:
+  absl::Status status_;
+  Output output_ = Output(Operation(nullptr), 0);
+  Tensor tensor_;
+  const std::string node_name_ = "";
+  int32_t index_ = 0;
+  DataType data_type_ = DT_INVALID;
+};
+
+/// A type for representing the output of ops that produce more than one output,
+/// or a list of tensors.
+typedef std::vector<Output> OutputList;
+
+/// A type for representing the input to ops that require a list of tensors.
+class InputList {
+ public:
+  /// Implicitly convert a list of outputs to a list of inputs. This is useful
+  /// to write code such as ops::Concat(ops::Split(x, 4)).
+  InputList(const OutputList& out) {  // NOLINT(runtime/explicit)
+    for (auto const& x : out) {
+      inputs_.push_back(x);
+    }
+  }
+
+  InputList(
+      const std::initializer_list<Input>& inputs)  // NOLINT(runtime/explicit)
+      : inputs_(inputs.begin(), inputs.end()) {}
+
+  InputList(const absl::Span<const Input>& inputs)  // NOLINT(runtime/explicit)
+      : inputs_(inputs.begin(), inputs.end()) {}
+
+  InputList(
+      const std::initializer_list<Output>& out) {  // NOLINT(runtime/explicit)
+    for (auto const& x : out) {
+      inputs_.push_back(x);
+    }
+  }
+
+  typename std::vector<Input>::iterator begin() { return inputs_.begin(); }
+  typename std::vector<Input>::iterator end() { return inputs_.end(); }
+  typename std::vector<Input>::const_iterator begin() const {
+    return inputs_.begin();
+  }
+  typename std::vector<Input>::const_iterator end() const {
+    return inputs_.end();
+  }
+
+ private:
+  std::vector<Input> inputs_;
+};
+
+/// @}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_OPS_H_

tensorflow/cc/framework/scope.cc

@@ -0,0 +1,558 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/scope_internal.h"
+#include "tensorflow/core/common_runtime/shape_refiner.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+
+namespace tensorflow {
+
+Scope::Scope(Impl* impl) : impl_(impl) {}
+
+Scope::Scope(const Scope& other) : impl_(new Impl(*other.impl())) {}
+
+Scope::~Scope() {}
+
+Scope& Scope::operator=(const Scope& other) {
+  // We can't copy Impls because of the const members, use copy ctor instead
+  impl_.reset(new Impl(*other.impl_));
+  return *this;
+}
+
+namespace {
+const char kScopeSeparator[] = "/";
+const char kSuffixSeparator[] = "_";
+}  // namespace
+
+Scope::Impl::Impl(Graph* graph, absl::Status* status, NameMap* name_map,
+                  ShapeRefiner* refiner, bool disable_shape_inference)
+    : graph_(graph),
+      status_(status),
+      name_map_(name_map),
+      refiner_(refiner),
+      scope_used_(nullptr),
+      colocation_constraints_(),
+      disable_shape_inference_(disable_shape_inference) {}
+
+Scope::Impl::Impl(const std::shared_ptr<Graph>& graph,
+                  const std::shared_ptr<absl::Status>& status,
+                  const std::shared_ptr<NameMap>& name_map,
+                  const std::shared_ptr<ShapeRefiner>& refiner)
+    : graph_(graph),
+      status_(status),
+      name_map_(name_map),
+      refiner_(refiner),
+      scope_used_(nullptr),
+      colocation_constraints_(),
+      disable_shape_inference_(refiner_ == nullptr) {}
+
+Scope Scope::NewRootScope() {
+  Graph* graph = new Graph(OpRegistry::Global());
+  ShapeRefiner* refiner =
+      new ShapeRefiner(graph->versions(), graph->op_registry());
+  return Scope(new Impl(graph, new absl::Status, new Impl::NameMap, refiner,
+                        /* disable_shape_inference */ false));
+}
+
+Scope Scope::DisabledShapeInferenceScope() {
+  Graph* graph = new Graph(OpRegistry::Global());
+  ShapeRefiner* refiner =
+      new ShapeRefiner(graph->versions(), graph->op_registry());
+  return Scope(new Impl(graph, new absl::Status, new Impl::NameMap, refiner,
+                        /* disable_shape_inference */ true));
+}
+
+Scope::Impl::Impl(const Scope& other, Tags::ScopeName, const std::string& name,
+                  bool copy_names)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(copy_names ? other.impl()->name_map_
+                           : std::shared_ptr<NameMap>(new NameMap)),
+      refiner_(other.impl()->refiner_),
+      scope_used_(nullptr),
+      control_deps_(other.impl()->control_deps_),
+      name_(name),
+      op_name_(""),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::OpName, const std::string& name,
+                  const std::string& op_name)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(name),
+      op_name_(op_name),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::ControlDeps,
+                  std::vector<Operation> control_deps, bool clear_control_deps)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(
+          clear_control_deps
+              ? std::vector<Operation>()
+              : (control_deps.insert(control_deps.begin(),
+                                     other.impl()->control_deps_.begin(),
+                                     other.impl()->control_deps_.end()),
+                 control_deps)),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::Device, const std::string& device)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(device),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::SingleUseScope,
+                  const std::string& op_name)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(new bool(false)),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(op_name),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::ExitOnError)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(true),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::KernelLabel,
+                  const std::string& kernel_label)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(kernel_label),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::Colocate,
+                  const Operation& colocate_with_op, bool clear_colocations)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(
+          clear_colocations
+              ? std::unordered_set<std::string>()
+              : other.impl()->GetColocationConstraints(colocate_with_op)),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::AssignedDevice,
+                  const std::string& assigned_device)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(assigned_device),
+      xla_cluster_(other.impl()->xla_cluster_),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+Scope::Impl::Impl(const Scope& other, Tags::XlaCluster,
+                  const std::string& xla_cluster)
+    : graph_(other.impl()->graph_),
+      status_(other.impl()->status_),
+      name_map_(other.impl()->name_map_),
+      refiner_(other.impl()->refiner_),
+      scope_used_(other.impl()->scope_used_),
+      control_deps_(other.impl()->control_deps_),
+      name_(other.impl()->name_),
+      op_name_(other.impl()->op_name_),
+      exit_on_error_(other.impl()->exit_on_error_),
+      kernel_label_(other.impl()->kernel_label_),
+      device_(other.impl()->device_),
+      assigned_device_(other.impl()->assigned_device_),
+      xla_cluster_(xla_cluster),
+      colocation_constraints_(other.impl()->colocation_constraints_),
+      disable_shape_inference_(other.impl()->disable_shape_inference_) {}
+
+std::unordered_set<std::string> Scope::Impl::GetColocationConstraints(
+    const Operation& colocate_with_op) const {
+  std::unordered_set<std::string> current_constraints(colocation_constraints_);
+  const AttrSlice attrs = colocate_with_op.node()->attrs();
+  std::vector<std::string> node_constraints;
+  if (TryGetNodeAttr(attrs, kColocationAttrName, &node_constraints)) {
+    for (const std::string& entry : node_constraints) {
+      absl::string_view s(entry);
+      if (absl::ConsumePrefix(&s, kColocationGroupPrefix)) {
+        current_constraints.emplace(s);
+      }
+    }
+  } else {
+    current_constraints.insert(colocate_with_op.node()->name());
+  }
+  return current_constraints;
+}
+
+bool Scope::ok() const { return impl()->status_->ok(); }
+
+Graph* Scope::graph() const { return impl()->graph_.get(); }
+
+std::shared_ptr<Graph> Scope::graph_as_shared_ptr() const {
+  return impl()->graph_;
+}
+
+absl::Status Scope::status() const { return *impl()->status_; }
+
+const std::vector<Operation>& Scope::control_deps() const {
+  return impl()->control_deps_;
+}
+
+void Scope::UpdateStatus(const absl::Status& s) const {
+  impl()->status_->Update(s);
+  if (impl()->exit_on_error_ && !ok()) {
+    LOG(FATAL) << *impl()->status_;
+  }
+}
+
+absl::Status Scope::ToGraphDef(GraphDef* gdef, bool include_debug_info) const {
+  if (!ok()) {
+    return *impl()->status_;
+  }
+  graph()->ToGraphDef(gdef, /*include_flib_def=*/true, include_debug_info);
+  return absl::OkStatus();
+}
+
+absl::Status Scope::ToGraph(Graph* g, GraphConstructorOptions opts) const {
+  if (ok()) {
+    GraphDef graph_def;
+    graph()->ToGraphDef(&graph_def);
+    UpdateStatus(ConvertGraphDefToGraph(opts, std::move(graph_def), g));
+  }
+  return *impl()->status_;
+}
+
+void Scope::UpdateBuilder(NodeBuilder* builder) const {
+  std::vector<Node*> control_inputs;
+  for (const auto& op : impl()->control_deps_) {
+    control_inputs.push_back(op.node());
+  }
+  builder->ControlInputs(control_inputs);
+
+  if (!impl()->kernel_label_.empty()) {
+    builder->Attr("_kernel", impl()->kernel_label_);
+  }
+
+  if (!impl()->colocation_constraints_.empty()) {
+    std::vector<std::string> constraints(
+        impl()->colocation_constraints_.begin(),
+        impl()->colocation_constraints_.end());
+    // Sort the set.
+    std::sort(constraints.begin(), constraints.end());
+    // Add loc:@ prefix
+    std::transform(constraints.begin(), constraints.end(), constraints.begin(),
+                   [](const std::string& s) {
+                     return absl::StrCat(kColocationGroupPrefix, s);
+                   });
+    builder->Attr(kColocationAttrName, constraints);
+  }
+  if (!impl()->device_.empty()) {
+    builder->Device(impl()->device_);
+  }
+  if (!impl()->assigned_device_.empty()) {
+    builder->AssignedDevice(impl()->assigned_device_);
+  }
+  if (!impl()->xla_cluster_.empty()) {
+    builder->XlaCluster(impl()->xla_cluster_);
+  }
+}
+
+std::string Scope::Impl::GetUniqueName(const std::string& prefix,
+                                       bool check_single_use) const {
+  if (check_single_use && single_use_scope()) {
+    if (*scope_used_) {
+      *status_ = absl::AlreadyExistsError(
+          absl::StrCat(prefix, " already exists in the current scope"));
+      return "";
+    }
+    *scope_used_ = true;
+    return prefix;
+  }
+  auto entry = name_map_->find(prefix);
+  if (entry == name_map_->end()) {
+    name_map_->insert({prefix, 0});
+    return prefix;
+  }
+  std::string unique_name;
+  do {
+    unique_name = absl::StrCat(prefix, kSuffixSeparator, ++entry->second);
+  } while (name_map_->find(unique_name) != name_map_->end());
+  name_map_->insert({unique_name, 0});
+  return unique_name;
+}
+
+std::string Scope::Impl::GetNameForOp(const std::string& default_name) const {
+  const std::string unique_name =
+      GetUniqueName(default_name, true /* check_single_use */);
+  const std::string sep =
+      name_.empty() || unique_name.empty() ? "" : kScopeSeparator;
+  return absl::StrCat(name_, sep, unique_name);
+}
+
+std::string Scope::GetUniqueNameForOp(const std::string& default_name) const {
+  if (impl()->single_use_scope()) {
+    if (impl()->op_name_.empty() || *impl()->scope_used_) {
+      *impl()->status_ =
+          absl::InvalidArgumentError("Cannot get a unique name in this scope");
+      return "";
+    }
+    *impl()->scope_used_ = true;
+    return impl()->op_name_;
+  }
+  return impl()->op_name_.empty() ? impl()->GetNameForOp(default_name)
+                                  : impl()->GetNameForOp(impl()->op_name_);
+}
+
+Scope Scope::NewSubScope(const std::string& child_scope_name) const {
+  if (child_scope_name.empty()) {
+    return Scope(new Impl(*this, Impl::Tags::ScopeName(), impl()->name_,
+                          true /* copy_names */));
+  }
+  const std::string unique_name =
+      impl()->GetUniqueName(child_scope_name, false /* check_single_use */);
+  const std::string sep =
+      impl()->name_.empty() || unique_name.empty() ? "" : kScopeSeparator;
+  return Scope(new Impl(*this, Impl::Tags::ScopeName(),
+                        absl::StrCat(impl()->name_, sep, unique_name),
+                        false /* copy_names */));
+}
+
+Scope Scope::WithOpNameImpl(const std::string& op_name) const {
+  if (impl()->single_use_scope()) {
+    UpdateStatus(absl::InvalidArgumentError(
+        absl::StrCat("Cannot set op name ", op_name, " on this scope")));
+    return *this;
+  }
+  return Scope(new Impl(*this, Impl::Tags::OpName(), impl()->name_, op_name));
+}
+
+Scope Scope::WithControlDependencies(
+    const absl::Span<const Operation> control_deps) const {
+  return Scope(
+      new Impl(*this, Impl::Tags::ControlDeps(),
+               std::vector<Operation>(control_deps.begin(), control_deps.end()),
+               /* clear_control_deps */ false));
+}
+
+Scope Scope::WithControlDependencies(const Output& control_dep) const {
+  return Scope(new Impl(*this, Impl::Tags::ControlDeps(),
+                        std::vector<Operation>(1, control_dep.op()),
+                        /* clear_control_deps */ false));
+}
+
+Scope Scope::WithNoControlDependencies() const {
+  return Scope(new Impl(*this, Impl::Tags::ControlDeps(),
+                        std::vector<Operation>(),
+                        /* clear_control_deps */ true));
+}
+
+Scope Scope::WithDevice(const std::string& device) const {
+  return Scope(new Impl(*this, Impl::Tags::Device(), device));
+}
+
+Scope Scope::WithAssignedDevice(const std::string& assigned_device) const {
+  return Scope(new Impl(*this, Impl::Tags::AssignedDevice(), assigned_device));
+}
+
+Scope Scope::WithXlaCluster(const std::string& xla_cluster) const {
+  return Scope(new Impl(*this, Impl::Tags::XlaCluster(), xla_cluster));
+}
+
+Scope Scope::ColocateWith(const Operation& op) const {
+  return Scope(new Impl(*this, Impl::Tags::Colocate(), op,
+                        /* clear_colocations */ false));
+}
+
+Scope Scope::ClearColocation() const {
+  return Scope(new Impl(*this, Impl::Tags::Colocate(), Operation(),
+                        /* clear_colocations */ true));
+}
+
+Scope Scope::ExitOnError() const {
+  return Scope(new Impl(*this, Impl::Tags::ExitOnError()));
+}
+
+Scope Scope::WithKernelLabel(const std::string& kernel_label) const {
+  return Scope(new Impl(*this, Impl::Tags::KernelLabel(), kernel_label));
+}
+
+CompositeOpScopes Scope::GetCompositeOpScopes(
+    const std::string& composite_op_name) const {
+  if (impl()->op_name_.empty() && composite_op_name.empty()) {
+    UpdateStatus(absl::InvalidArgumentError(
+        "Cannot create composite op scopes with empty name"));
+    return {*this, *this};
+  }
+  if (!impl()->single_use_scope()) {
+    Scope child = NewSubScope(impl()->op_name_.empty() ? composite_op_name
+                                                       : impl()->op_name_);
+    const std::string child_op_sep =
+        impl()->name_.empty() ? "" : kSuffixSeparator;
+    const std::string child_name =
+        absl::StrCat(impl()->name_, child_op_sep, child.impl()->name_);
+    return {child,
+            Scope(new Impl(child, Impl::Tags::SingleUseScope(), child_name))};
+  } else {
+    return {Scope(new Impl(*this, Impl::Tags::ScopeName(), impl()->op_name_,
+                           true /* copy_names */)),
+            *this};
+  }
+}
+
+absl::Status Scope::DoShapeInference(Node* node) const {
+  if (impl_->disable_shape_inference_) return absl::OkStatus();
+  return impl_->refiner_->AddNode(node);
+}
+
+class InternalScope {
+ public:
+  // NewScope doesn't take ownership of the inputs.
+  static Scope NewScope(Graph* graph, absl::Status* status,
+                        ShapeRefiner* refiner) {
+    Scope::Impl::NameMap* name_map = new Scope::Impl::NameMap;
+    for (const Node* node : graph->nodes()) {
+      const std::string& name = node->name();
+      (*name_map)[name] = 0;
+      // Add all name prefixes ('/' separated).
+      size_t idx = -1;
+      while ((idx = name.find(kScopeSeparator, idx + 1)) != std::string::npos) {
+        (*name_map)[name.substr(0, idx)] = 0;
+      }
+    }
+    // We provide null destructors for these shared ptrs (except for name_map)
+    // since the caller owns them and doesn't want the scope to destroy them.
+    return Scope(new Scope::Impl(
+        std::shared_ptr<Graph>(graph, [](Graph*) {}),
+        std::shared_ptr<absl::Status>(status, [](absl::Status*) {}),
+        std::shared_ptr<Scope::Impl::NameMap>(name_map),
+        std::shared_ptr<ShapeRefiner>(refiner, [](ShapeRefiner*) {})));
+  }
+};
+
+Scope NewInternalScope(Graph* graph, absl::Status* status,
+                       ShapeRefiner* refiner) {
+  return InternalScope::NewScope(graph, status, refiner);
+}
+
+absl::Status CreateOutputWithScope(std::string op_name,
+                                   absl::Span<const ::tensorflow::Input> inputs,
+                                   const Scope& scope, Output* output) {
+  TF_RETURN_IF_ERROR(scope.status());
+  const auto unique_name = scope.GetUniqueNameForOp(op_name);
+  auto builder = ::tensorflow::NodeBuilder(unique_name, op_name);
+  for (const auto& input : inputs) {
+    TF_RETURN_IF_ERROR(scope.status());
+    builder = builder.Input(input.node());
+  }
+  ::tensorflow::Node* ret;
+  scope.UpdateBuilder(&builder);
+  TF_RETURN_IF_ERROR(scope.status());
+  scope.UpdateStatus(builder.Finalize(scope.graph(), &ret));
+  TF_RETURN_IF_ERROR(scope.status());
+  *output = Output(ret, 0);
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/cc/framework/scope.h

@@ -0,0 +1,271 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_SCOPE_H_
+#define TENSORFLOW_CC_FRAMEWORK_SCOPE_H_
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+
+namespace tensorflow {
+
+class Graph;
+class GraphDef;
+class NodeBuilder;
+struct CompositeOpScopes;
+
+/// @addtogroup core
+/// @{
+
+/// A `Scope` object represents a set of related TensorFlow ops that have the
+/// same properties such as a common name prefix.
+///
+/// A Scope object is a container for TensorFlow Op properties. Op constructors
+/// get a Scope object as a mandatory first argument and the constructed op
+/// acquires the properties in the object.
+///
+/// A simple example:
+///
+///     using namespace ops;
+///     Scope root = Scope::NewRootScope();
+///     auto c1 = Const(root, { {1, 1} });
+///     auto m = MatMul(root, c1, { {41}, {1} });
+///     GraphDef gdef;
+///     Status s = root.ToGraphDef(&gdef);
+///     if (!s.ok()) { ... }
+///
+/// Scope hierarchy:
+///
+/// The Scope class provides various With<> functions that create a new scope.
+/// The new scope typically has one property changed while other properties are
+/// inherited from the parent scope.
+/// NewSubScope(name) method appends `name` to the prefix of names for ops
+/// created within the scope, and WithOpName() changes the suffix which
+/// otherwise defaults to the type of the op.
+///
+/// Name examples:
+///
+///     Scope root = Scope::NewRootScope();
+///     Scope linear = root.NewSubScope("linear");
+///     // W will be named "linear/W"
+///     auto W = Variable(linear.WithOpName("W"),
+///                       {2, 2}, DT_FLOAT);
+///     // b will be named "linear/b_3"
+///     int idx = 3;
+///     auto b = Variable(linear.WithOpName("b_", idx),
+///                       {2}, DT_FLOAT);
+///     auto x = Const(linear, {...});  // name: "linear/Const"
+///     auto m = MatMul(linear, x, W);  // name: "linear/MatMul"
+///     auto r = BiasAdd(linear, m, b); // name: "linear/BiasAdd"
+///
+/// Scope lifetime:
+///
+/// A new scope is created by calling Scope::NewRootScope. This creates some
+/// resources that are shared by all the child scopes that inherit from this
+/// scope, directly or transitively. For instance, a new scope creates a new
+/// Graph object to which operations are added when the new scope or its
+/// children are used by an Op constructor. The new scope also has a Status
+/// object which will be used to indicate errors by Op-constructor functions
+/// called on any child scope. The Op-constructor functions have to check the
+/// scope's status by calling the ok() method before proceeding to construct the
+/// op.
+///
+/// Thread safety:
+///
+/// A `Scope` object is NOT thread-safe. Threads cannot concurrently call
+/// op-constructor functions on the same `Scope` object.
+class Scope {
+ public:
+  Scope(const Scope& other);
+  ~Scope();
+  Scope& operator=(const Scope& other);
+
+  // The following functions are for users making graphs. They return brand new
+  // scopes, or scopes derived from an existing scope object.
+
+  /// Return a new scope.
+  /// This creates a new graph and all operations constructed in this graph
+  /// should use the returned object as the "root" scope.
+  static Scope NewRootScope();
+
+  /// Return a new scope. Ops created with this scope will have
+  /// `name/child_scope_name` as the prefix. The actual name will be unique
+  /// in the current scope. All other properties are inherited from the current
+  /// scope. If `child_scope_name` is empty, the `/` is elided.
+  Scope NewSubScope(const std::string& child_scope_name) const;
+
+  /// Return a new scope. All ops created within the returned scope will have
+  /// names of the form `name/StrCat(fragments...)[_suffix]`
+  template <typename... Ty>
+  Scope WithOpName(Ty... fragments) const {
+    return WithOpNameImpl(absl::StrCat(fragments...));
+  }
+
+  /// Return a new scope. All ops created within the returned scope will have as
+  /// control dependencies the union of operations in the control_deps vector
+  /// and the control dependencies of the current scope.
+  Scope WithControlDependencies(absl::Span<const Operation> control_deps) const;
+  /// Same as above, but convenient to add control dependency on the operation
+  /// producing the control_dep output.
+  Scope WithControlDependencies(const Output& control_dep) const;
+
+  /// Return a new scope. All ops created within the returned scope will have no
+  /// control dependencies on other operations.
+  Scope WithNoControlDependencies() const;
+
+  /// Return a new scope. All ops created within the returned scope will have
+  /// the device field set to 'device'.
+  Scope WithDevice(const std::string& device) const;
+
+  /// Returns a new scope.  All ops created within the returned scope will have
+  /// their assigned device set to `assigned_device`.
+  Scope WithAssignedDevice(const std::string& assigned_device) const;
+
+  /// Returns a new scope.  All ops created within the returned scope will have
+  /// their _XlaCluster attribute set to `xla_cluster`.
+  Scope WithXlaCluster(const std::string& xla_cluster) const;
+
+  /// Return a new scope. All ops created within the returned scope will be
+  /// co-located on the device where op is placed.
+  /// NOTE: This function is intended to be use internal libraries only for
+  /// controlling placement of ops on to devices. Public use is not encouraged
+  /// because the implementation of device placement is subject to change.
+  Scope ColocateWith(const Operation& op) const;
+  /// Convenience function for above.
+  Scope ColocateWith(const Output& out) const { return ColocateWith(out.op()); }
+  /// Clear all colocation constraints.
+  Scope ClearColocation() const;
+
+  /// Return a new scope. The op-constructor functions taking the returned scope
+  /// as the scope argument will exit as soon as an error is detected, instead
+  /// of setting the status on the scope.
+  Scope ExitOnError() const;
+
+  /// Return a new scope. All ops created with the new scope will have
+  /// kernel_label as the value for their '_kernel' attribute;
+  Scope WithKernelLabel(const std::string& kernel_label) const;
+
+  // The following functions are for scope object consumers.
+
+  /// Return a unique name, using default_name if an op name has not been
+  /// specified.
+  std::string GetUniqueNameForOp(const std::string& default_name) const;
+
+  /// Update the status on this scope.
+  /// Note: The status object is shared between all children of this scope.
+  /// If the resulting status is not OkStatus() and exit_on_error_ is set on
+  /// this scope, this function exits by calling LOG(FATAL).
+  void UpdateStatus(const absl::Status& s) const;
+
+  // START_SKIP_DOXYGEN
+
+  /// Update the builder with properties accumulated in this scope. Does not set
+  /// status().
+  // TODO(skyewm): NodeBuilder is not part of public API
+  void UpdateBuilder(NodeBuilder* builder) const;
+  // END_SKIP_DOXYGEN
+
+  CompositeOpScopes GetCompositeOpScopes(
+      const std::string& composite_op_name) const;
+
+  bool ok() const;
+
+  // TODO(skyewm): Graph is not part of public API
+  Graph* graph() const;
+
+  // TODO(skyewm): Graph is not part of public API
+  std::shared_ptr<Graph> graph_as_shared_ptr() const;
+
+  absl::Status status() const;
+
+  /// If status() is ok, convert the Graph object stored in this scope
+  /// to a GraphDef proto and return an ok Status. Otherwise, return the error
+  /// status as is without performing GraphDef conversion. If
+  /// `include_debug_info` is true, populate the `debug_info` field of the
+  /// GraphDef from stack traces in this Graph.
+  absl::Status ToGraphDef(GraphDef* gdef,
+                          bool include_debug_info = false) const;
+
+  // START_SKIP_DOXYGEN
+
+  /// If status() is OkStatus(), construct a Graph object using `opts` as the
+  /// GraphConstructorOptions, and return Status::OK if graph construction was
+  /// successful. Otherwise, return the error status.
+  // TODO(josh11b, keveman): Make this faster; right now it converts
+  // Graph->GraphDef->Graph.  This cleans up the graph (e.g. adds
+  // edges from the source and to the sink node, resolves back edges
+  // by name), and makes sure the resulting graph is valid.
+  absl::Status ToGraph(
+      Graph* g, GraphConstructorOptions opts = GraphConstructorOptions{}) const;
+
+  // Calls AddNode() using this scope's ShapeRefiner. This exists in the public
+  // API to prevent custom op wrappers from needing access to shape_refiner.h or
+  // scope_internal.h.
+  // TODO(skyewm): remove this from public API
+  absl::Status DoShapeInference(Node* node) const;
+
+  // Creates a new root scope that causes all DoShapeInference() calls to return
+  // OkStatus() (on the returned scope and any subscopes). Used for testing.
+  // TODO(skyewm): fix tests that still require this and eventually remove, or
+  // at least remove from public API
+  static Scope DisabledShapeInferenceScope();
+  // END_SKIP_DOXYGEN
+
+  const std::vector<Operation>& control_deps() const;
+
+  // START_SKIP_DOXYGEN
+  class Impl;
+  Impl* impl() { return impl_.get(); }
+  const Impl* impl() const { return impl_.get(); }
+  // END_SKIP_DOXYGEN
+
+ private:
+  Scope WithOpNameImpl(const std::string& op_name) const;
+
+  friend class InternalScope;
+  std::unique_ptr<Impl> impl_;
+  explicit Scope(Impl*);
+};
+
+/// A helper struct to hold the scopes that would be used by a function
+/// constructing a composite op.
+struct CompositeOpScopes {
+  /// Scope to be used for creating the local ops (primitive or other composite
+  /// ops).
+  Scope child;
+  /// Scope to be used for creating the last op.
+  Scope last;
+};
+
+// Creates a node of the given operation, with the given inputs, and assigns the
+// result to output. This does not support the ability to add additional
+// attributes.
+absl::Status CreateOutputWithScope(std::string op_name,
+                                   absl::Span<const ::tensorflow::Input> inputs,
+                                   const Scope& scope, Output* output);
+/// @}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_SCOPE_H_

tensorflow/cc/framework/scope_internal.h

@@ -0,0 +1,136 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_
+#define TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+
+class ShapeRefiner;
+
+// NewInternalScope returns a new scope which doesn't take ownership of
+// graph, status, name_map, and refiner.
+// This is intended to enable the C API (which are used by other language
+// bindings) to create a Scope and access C++ functionality (i.e. gradients).
+//
+// Shape inference is disabled if `refiner` is nullptr.
+Scope NewInternalScope(Graph* graph, absl::Status* status,
+                       ShapeRefiner* refiner);
+
+class Scope::Impl {
+ public:
+  // A NameMap is used to keep track of suffixes for names used in a scope. A
+  // name that has not been used so far in a scope will get no suffix. Later
+  // uses of the same name will get suffixes _1, _2, _3, etc. Multiple scopes
+  // can share the same NameMap. For instance, a new scope created using
+  // WithControlDependencies() would share the same NameMap with the parent.
+  typedef std::unordered_map<std::string, int> NameMap;
+
+  Impl(const std::shared_ptr<Graph>& graph,
+       const std::shared_ptr<absl::Status>& status,
+       const std::shared_ptr<NameMap>& name_map,
+       const std::shared_ptr<ShapeRefiner>& refiner);
+
+  const std::string& name() const { return name_; }
+  const std::vector<Operation>& control_deps() const { return control_deps_; }
+
+ private:
+  friend class Scope;
+
+  // Tag types to choose the constructor to dispatch.
+  struct Tags {
+    enum class ScopeName;
+    enum class OpName;
+    enum class ControlDeps;
+    enum class Device;
+    enum class SingleUseScope;
+    enum class ExitOnError;
+    enum class KernelLabel;
+    enum class Colocate;
+    enum class AssignedDevice;
+    enum class XlaCluster;
+  };
+
+  Impl(Graph* graph, absl::Status* status, NameMap* name_map,
+       ShapeRefiner* refiner, bool disable_shape_inference);
+  Impl(const Scope& other, Tags::ScopeName, const std::string& name,
+       bool copy_names);
+  Impl(const Scope& other, Tags::OpName, const std::string& name,
+       const std::string& op_name);
+  Impl(const Scope& other, Tags::ControlDeps,
+       std::vector<Operation> control_deps, bool clear_control_deps);
+  Impl(const Scope& other, Tags::Device, const std::string& device);
+  Impl(const Scope& other, Tags::SingleUseScope, const std::string& op_name);
+  Impl(const Scope& other, Tags::ExitOnError);
+  Impl(const Scope& other, Tags::KernelLabel, const std::string& kernel_label);
+  Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op,
+       bool clear_colocations);
+  Impl(const Scope& other, Tags::AssignedDevice,
+       const std::string& assigned_device);
+  Impl(const Scope& other, Tags::XlaCluster, const std::string& xla_cluster);
+
+  std::unordered_set<std::string> GetColocationConstraints(
+      const Operation& colocate_with_op) const;
+
+  // Helper functions to get a unique names.
+  std::string GetUniqueName(const std::string& prefix,
+                            bool check_single_use) const;
+  std::string GetNameForOp(const std::string& default_name) const;
+
+  bool single_use_scope() const { return scope_used_ != nullptr; }
+
+  // The graph, status, and name maps are shared by all child scopes
+  // created from a single 'root' scope. A root scope is created by calling the
+  // Scope::NewRootScope function, which creates a new graph, a new status and
+  // the name maps.
+  std::shared_ptr<Graph> graph_ = nullptr;
+  std::shared_ptr<absl::Status> status_ = nullptr;
+  std::shared_ptr<NameMap> name_map_ = nullptr;
+  std::shared_ptr<ShapeRefiner> refiner_ = nullptr;
+
+  // If scope_used_ is not nullptr, op_name_ should be empty and
+  // GetUniqueNameForOp can only be called once on this scope. More calls to
+  // GetUniqueNameForOp will cause an error status to be set on this scope.
+  std::shared_ptr<bool> scope_used_ = nullptr;
+
+  const std::vector<Operation> control_deps_;
+
+  // The fully-qualified name of this scope (i.e. includes any parent scope
+  // names).
+  const std::string name_ = "";
+  const std::string op_name_ = "";
+  const bool exit_on_error_ = false;
+  const std::string kernel_label_ = "";
+  const std::string device_ = "";
+  const std::string assigned_device_ = "";
+  const std::string xla_cluster_ = "";
+  const std::unordered_set<std::string> colocation_constraints_;
+
+  // If true, Scope::DoShapeInference() always returns Status:OK().
+  // TODO(skyewm): remove this when possible
+  const bool disable_shape_inference_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_

tensorflow/cc/framework/scope_test.cc

@@ -0,0 +1,162 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/scope.h"
+
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+TEST(ScopeTest, BasicNames) {
+  Scope root = Scope::NewRootScope();
+  EXPECT_EQ(root.GetUniqueNameForOp("add"), "add");
+  EXPECT_EQ(root.GetUniqueNameForOp("add"), "add_1");
+  EXPECT_EQ(root.GetUniqueNameForOp("add"), "add_2");
+  EXPECT_EQ(root.GetUniqueNameForOp("mul"), "mul");
+}
+
+TEST(ScopeTest, OpAndScopeNameCollision) {
+  Scope root = Scope::NewRootScope();
+  EXPECT_EQ(root.GetUniqueNameForOp("foo"), "foo");
+  EXPECT_EQ(root.GetUniqueNameForOp("foo"), "foo_1");
+  EXPECT_EQ(root.GetUniqueNameForOp("foo_1"), "foo_1_1");
+  EXPECT_EQ(root.GetUniqueNameForOp("foo_2"), "foo_2");
+  EXPECT_EQ(root.GetUniqueNameForOp("foo"), "foo_3");
+  EXPECT_EQ(root.GetUniqueNameForOp("foo_2"), "foo_2_1");
+}
+
+TEST(ScopeTest, HierarchicalNames) {
+  Scope root = Scope::NewRootScope();
+  Scope child = root.NewSubScope("child");
+  EXPECT_EQ(child.GetUniqueNameForOp("add"), "child/add");
+  EXPECT_EQ(child.GetUniqueNameForOp("add"), "child/add_1");
+  EXPECT_EQ(child.GetUniqueNameForOp("mul"), "child/mul");
+
+  Scope child_1 = root.NewSubScope("child");
+  EXPECT_EQ(child_1.GetUniqueNameForOp("add"), "child_1/add");
+  EXPECT_EQ(child_1.GetUniqueNameForOp("add"), "child_1/add_1");
+  EXPECT_EQ(child_1.GetUniqueNameForOp("mul"), "child_1/mul");
+
+  Scope c_c = root.NewSubScope("c").NewSubScope("c");
+  EXPECT_EQ(c_c.GetUniqueNameForOp("add"), "c/c/add");
+
+  Scope c_1 = root.NewSubScope("c");
+  Scope c_1_c = c_1.NewSubScope("c");
+  EXPECT_EQ(c_1_c.GetUniqueNameForOp("add"), "c_1/c/add");
+
+  Scope c_1_c_1 = c_1.NewSubScope("c");
+  EXPECT_EQ(c_1_c_1.GetUniqueNameForOp("add"), "c_1/c_1/add");
+
+  EXPECT_EQ(root.NewSubScope("").NewSubScope("").GetUniqueNameForOp("d"), "d");
+  EXPECT_EQ(root.NewSubScope("").GetUniqueNameForOp("d"), "d_1");
+  EXPECT_EQ(root.GetUniqueNameForOp("d"), "d_2");
+}
+
+TEST(ScopeTest, ScopeAndOpNames) {
+  Scope root = Scope::NewRootScope();
+  Scope child = root.NewSubScope("child");
+
+  EXPECT_EQ(child.GetUniqueNameForOp("add"), "child/add");
+  EXPECT_EQ(root.GetUniqueNameForOp("child"), "child_1");
+
+  EXPECT_EQ(root.NewSubScope("child").GetUniqueNameForOp("p"), "child_2/p");
+}
+
+namespace {
+
+std::string LastOp(const Scope& scope) {
+  return scope.GetUniqueNameForOp("Last");
+}
+
+std::vector<std::string> AnotherCompositeOp(const Scope& scope) {
+  auto cop_scopes = scope.GetCompositeOpScopes("another_cop");
+  const std::string c1 = cop_scopes.child.GetUniqueNameForOp("c1");
+  const std::string c2 = cop_scopes.child.GetUniqueNameForOp("mul");
+  return {c1, c2, LastOp(cop_scopes.last)};
+}
+
+std::vector<std::string> LinearOp(const Scope& scope) {
+  auto cop_scopes = scope.GetCompositeOpScopes("linear");
+  Scope linear = cop_scopes.child;
+  const std::string mul_op_name = linear.GetUniqueNameForOp("mul");
+  const std::string bias_add_op_name = linear.GetUniqueNameForOp("bias_add");
+  auto cop_names = AnotherCompositeOp(cop_scopes.last);
+  return {mul_op_name, bias_add_op_name, cop_names[0], cop_names[1],
+          cop_names[2]};
+}
+
+}  // namespace
+
+TEST(ScopeTest, CompositeOp) {
+  Scope root = Scope::NewRootScope();
+  const auto names1 = LinearOp(root);
+
+  EXPECT_EQ(names1[0], "linear/mul");
+  EXPECT_EQ(names1[1], "linear/bias_add");
+  EXPECT_EQ(names1[2], "linear/c1");
+  EXPECT_EQ(names1[3], "linear/mul_1");
+  EXPECT_EQ(names1[4], "linear");
+
+  EXPECT_EQ(root.GetUniqueNameForOp("linear"), "linear_1");
+
+  const auto names2 = LinearOp(root);
+
+  EXPECT_EQ(names2[0], "linear_2/mul");
+  EXPECT_EQ(names2[1], "linear_2/bias_add");
+  EXPECT_EQ(names2[2], "linear_2/c1");
+  EXPECT_EQ(names2[3], "linear_2/mul_1");
+  EXPECT_EQ(names2[4], "linear_2");
+
+  const auto names3 = LinearOp(root.WithOpName("c"));
+
+  EXPECT_EQ(names3[0], "c/mul");
+  EXPECT_EQ(names3[1], "c/bias_add");
+  EXPECT_EQ(names3[2], "c/c1");
+  EXPECT_EQ(names3[3], "c/mul_1");
+  EXPECT_EQ(names3[4], "c");
+}
+
+TEST(ScopeTest, SingleUseScope) {
+  Scope root = Scope::NewRootScope();
+  auto cop_scopes = root.GetCompositeOpScopes("cop");
+  // cop_scopes.last is a single use scope
+  EXPECT_EQ(cop_scopes.last.GetUniqueNameForOp("foo"), "cop");
+  cop_scopes.last.GetUniqueNameForOp("foo");
+  // Error status should be set on cop_scopes.last
+  EXPECT_FALSE(cop_scopes.last.ok());
+}
+
+TEST(ScopeTest, ControlDeps) {
+  Scope root = Scope::NewRootScope();
+  auto c1 = Operation();
+  auto c2 = Operation();
+  Scope c = root.WithControlDependencies({c1, c2});
+  EXPECT_EQ(c.control_deps().size(), 2);
+  Scope c_c = c.WithControlDependencies({Operation()});
+  EXPECT_EQ(c_c.control_deps().size(), 3);
+}
+
+TEST(ScopeTest, CreateOutput) {
+  Scope root = Scope::NewRootScope();
+  Output a = ops::Placeholder(root.WithOpName("a"), DT_FLOAT);
+  Output add;
+  ASSERT_TRUE(
+      CreateOutputWithScope("Add", {a, a}, root.WithOpName("add"), &add).ok());
+  EXPECT_EQ(add.node()->name(), "add");
+  EXPECT_EQ(add.node()->type_string(), "Add");
+}
+
+}  // namespace tensorflow

tensorflow/cc/framework/test_op.cc

@@ -0,0 +1,57 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+
+namespace tensorflow {
+
+REGISTER_OP("ThrowAway1")
+    .Input("ret: int32")
+    .Input("unique_name: float")
+    .Input("for: int32")
+    .Attr("scope: int")
+    .Attr("builder: int = 1")
+    .Attr("while: int")
+    .SetShapeFn(shape_inference::UnknownShape)
+    .Doc(R"doc(
+Op to test keywords and reserved words in input and attr names.
+
+ret: Return value.
+for: Keyword as name for input.
+while: Keyword as name for attr.
+)doc");
+
+REGISTER_OP("ThrowAway2")
+    .Attr("scope: int = 2")
+    .Attr("throw_away2: int = 2")
+    .Attr("attrs: int = 4")
+    .Attr("node: int = 4")
+    .SetShapeFn(shape_inference::UnknownShape);
+
+REGISTER_OP("ThrowAway3")
+    .Output("node: int32")
+    .SetShapeFn(shape_inference::UnknownShape);
+
+REGISTER_OP("ThrowAway4")
+    .Input("node: int32")
+    .SetShapeFn(shape_inference::UnknownShape);
+
+REGISTER_OP("ThrowAway5")
+    .Output("foo: int32")
+    .Attr("node: int = 4")
+    .SetShapeFn(shape_inference::UnknownShape);
+
+}  // namespace tensorflow

tensorflow/cc/framework/testutil.cc

@@ -0,0 +1,54 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/testutil.h"
+
+#include <utility>
+
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/graph/default_device.h"
+
+namespace tensorflow {
+namespace test {
+
+void GetTensors(const Scope& scope, OutputList tensors,
+                std::vector<Tensor>* out) {
+  ClientSession session(scope);
+  TF_CHECK_OK(session.Run(tensors, out));
+}
+
+void GetTensor(const Scope& scope, Output tensor, Tensor* out) {
+  std::vector<Tensor> outputs;
+  GetTensors(scope, {std::move(tensor)}, &outputs);
+  *out = outputs[0];
+}
+
+void GetTensors(const Scope& scope, const std::vector<Output>& assign_vars,
+                const OutputList& tensors, std::vector<Tensor>* out) {
+  ClientSession session(scope);
+  TF_CHECK_OK(session.Run(assign_vars, nullptr));
+  TF_CHECK_OK(session.Run(tensors, out));
+}
+
+void GetTensor(const Scope& scope, const std::vector<Output>& assign_vars,
+               Output tensor, Tensor* out) {
+  std::vector<Tensor> outputs;
+  GetTensors(scope, assign_vars, {std::move(tensor)}, &outputs);
+  *out = outputs[0];
+}
+
+}  // end namespace test
+}  // end namespace tensorflow

tensorflow/cc/framework/testutil.h

@@ -0,0 +1,49 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_
+#define TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_
+
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+namespace test {
+
+/// Computes the outputs listed in 'tensors', returns the tensors in 'out'.
+void GetTensors(const Scope& scope, OutputList tensors,
+                std::vector<Tensor>* out);
+
+// Computes the outputs listed in 'tensors', returns the tensors in 'out'.
+// assign_vars are extra outputs that should be run
+// e.g. to assign values to variables.
+void GetTensors(const Scope& scope, const std::vector<Output>& assign_vars,
+                const OutputList& tensors, std::vector<Tensor>* out);
+
+/// Computes the output 'tensor', returning the resulting tensor in 'out'.
+void GetTensor(const Scope& scope, Output tensor, Tensor* out);
+
+// Computes the output 'tensor', returning the resulting tensor in 'out'.
+// assign_vars are extra outputs that should be run
+// e.g. to assign values to variables.
+void GetTensor(const Scope& scope, const std::vector<Output>& assign_vars,
+               Output tensor, Tensor* out);
+
+}  // namespace test
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_TESTUTIL_H_

tensorflow/cc/framework/while_gradients.cc

@@ -0,0 +1,201 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/framework/while_gradients.h"
+
+#include <string>
+
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/framework/scope_internal.h"
+#include "tensorflow/cc/ops/control_flow_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/cc/ops/while_loop.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::BodyGraphBuilderFn;
+using ops::BuildWhileLoop;
+using ops::CondGraphBuilderFn;
+
+Output ToOutput(OutputTensor output_tensor) {
+  return Output(const_cast<Node*>(output_tensor.node), output_tensor.index);
+}
+
+std::vector<Output> ToOutputVector(
+    const std::vector<OutputTensor>& output_tensors) {
+  const int n = output_tensors.size();
+  std::vector<Output> result;
+  result.reserve(n);
+  for (int i = 0; i < n; ++i) result.push_back(ToOutput(output_tensors[i]));
+  return result;
+}
+
+// The backprop loop counter and main backprop loop run in their own execution
+// frame (conceptually, the main forward loop and forward loop counter run
+// together in a frame, then the backprop loop counter and backprop loop run
+// together in a different frame). This returns the frame name to use for the
+// backprop while loops.
+// TODO(skyewm): make sure this is unique among existing frame names
+std::string BackPropFrameName(const std::string& forward_frame_name) {
+  return absl::StrCat(forward_frame_name, "_backprop");
+}
+
+// Creates a loop that counts the number of iterations performed by the
+// while loop associated with `while_ctx`. The returned output yields the
+// iteration count.
+absl::Status AddForwardLoopCounter(WhileContext* while_ctx, const Scope& scope,
+                                   Output* count) {
+  // Create while loop:
+  //   i = 0
+  //   while forward loop predicate is true:
+  //     ++i
+
+  Output zero = ops::Const(scope, 0, {});
+
+  // Condition function that returns condition output from original while loop.
+  CondGraphBuilderFn cond_fn = [while_ctx](const Scope& scope,
+                                           const std::vector<Output>& inputs,
+                                           Output* output) {
+    *output = ToOutput(while_ctx->cond_output());
+    return absl::OkStatus();
+  };
+
+  // Body function that adds one to input.
+  BodyGraphBuilderFn body_fn = [](const Scope& scope,
+                                  const std::vector<Output>& inputs,
+                                  std::vector<Output>* outputs) {
+    DCHECK_EQ(inputs.size(), 1);
+    outputs->emplace_back(ops::Add(scope, inputs[0], 1));
+    return scope.status();
+  };
+
+  // Note that this loop runs in the same execution frame as the forward loop.
+  std::vector<Output> outputs;
+  TF_RETURN_IF_ERROR(BuildWhileLoop(scope, {zero}, cond_fn, body_fn,
+                                    while_ctx->frame_name(), &outputs,
+                                    /* create_while_ctx */ false));
+  *count = outputs[0];
+  return absl::OkStatus();
+}
+
+// Creates a loop that executes `loop_count` times. The returned output is the
+// boolean predicate indicating if the loop is still executing. This is used to
+// drive the gradient computation for the while loop associated with
+// `while_ctx`.
+absl::Status AddBackPropLoopCounter(WhileContext* while_ctx,
+                                    const Output& loop_count,
+                                    const Scope& scope,
+                                    Output* backprop_execution_pred) {
+  // Create while loop:
+  //   n = loop_count
+  //   while n > 0:
+  //     --n
+
+  // Condition function that returns input > 0.
+  CondGraphBuilderFn cond_fn = [](const Scope& scope,
+                                  const std::vector<Output>& inputs,
+                                  Output* output) {
+    DCHECK_EQ(inputs.size(), 1);
+    *output = ops::Greater(scope, inputs[0], 0);
+    return scope.status();
+  };
+
+  // Body function that subtracts one from input.
+  BodyGraphBuilderFn body_fn = [](const Scope& scope,
+                                  const std::vector<Output>& inputs,
+                                  std::vector<Output>* outputs) {
+    DCHECK_EQ(inputs.size(), 1);
+    outputs->emplace_back(ops::Subtract(scope, inputs[0], 1));
+    return scope.status();
+  };
+
+  std::string frame_name = BackPropFrameName(while_ctx->frame_name());
+  std::vector<Output> outputs;
+  TF_RETURN_IF_ERROR(BuildWhileLoop(
+      scope, {loop_count}, cond_fn, body_fn, frame_name, &outputs,
+      /* create_while_ctx */ false, backprop_execution_pred));
+  return absl::OkStatus();
+}
+
+// Creates the main backprop loop that computes the gradient of the loop
+// associated with `while_ctx`. `grad_inputs` are the partial derivatives
+// w.r.t. the loop outputs, i.e. the exit nodes. `backprop_execution_pred` is
+// the predicate to use for the backprop loop (see AddBackPropLoopCounter()).
+// The partial derivatives w.r.t. the loop inputs, i.e. the input loop vars, are
+// returned in `grad_outputs`.
+absl::Status AddWhileGradientLoop(WhileContext* while_ctx,
+                                  const std::vector<Output>& grad_inputs,
+                                  const Output& backprop_execution_pred,
+                                  const Scope& parent_scope,
+                                  std::vector<Output>* grad_outputs) {
+  DCHECK_EQ(grad_inputs.size(), while_ctx->body_outputs().size());
+  DCHECK_EQ(while_ctx->body_inputs().size(), while_ctx->body_outputs().size());
+
+  Scope scope = parent_scope.NewSubScope("while");
+
+  // Create while loop:
+  //   while backprop_execution_pred:
+  //     forward loop body gradient
+
+  // Condition function that returns 'backprop_execution_pred'.
+  CondGraphBuilderFn cond_fn = [backprop_execution_pred](
+                                   const Scope& scope,
+                                   const std::vector<Output>& inputs,
+                                   Output* output) {
+    *output = backprop_execution_pred;
+    return absl::OkStatus();
+  };
+
+  // Body function that builds while body gradient subgraph.
+  BodyGraphBuilderFn body_fn = [while_ctx](const Scope& scope,
+                                           const std::vector<Output>& inputs,
+                                           std::vector<Output>* outputs) {
+    std::vector<Output> body_outputs =
+        ToOutputVector(while_ctx->body_outputs());
+    std::vector<Output> body_inputs = ToOutputVector(while_ctx->body_inputs());
+    return AddSymbolicGradients(scope, body_outputs, body_inputs, inputs,
+                                outputs);
+  };
+
+  std::string frame_name = BackPropFrameName(while_ctx->frame_name());
+  TF_RETURN_IF_ERROR(BuildWhileLoop(scope, grad_inputs, cond_fn, body_fn,
+                                    frame_name, grad_outputs,
+                                    /* create_while_ctx */ false));
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status AddWhileLoopGradient(WhileContext* while_ctx, const Scope& scope,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs) {
+  Output forward_loop_count;
+  TF_RETURN_IF_ERROR(AddForwardLoopCounter(
+      while_ctx, scope.NewSubScope("ForwardLoopCounter"), &forward_loop_count));
+
+  // TODO(skyewm): can we combine the backprop loop counter and main gradient
+  // loop into a single loop? The original Python code doesn't combine the
+  // loops, but I'm not sure why.
+  Output backprop_counter_cond;
+  TF_RETURN_IF_ERROR(AddBackPropLoopCounter(
+      while_ctx, forward_loop_count, scope.NewSubScope("BackPropLoopCounter"),
+      &backprop_counter_cond));
+
+  return AddWhileGradientLoop(while_ctx, grad_inputs, backprop_counter_cond,
+                              scope, grad_outputs);
+}
+
+}  // namespace tensorflow

tensorflow/cc/framework/while_gradients.h

@@ -0,0 +1,42 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_
+#define TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_
+
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/graph/while_context.h"
+
+// Utility functions for constructing while loop gradients
+
+namespace tensorflow {
+
+// Adds the gradient computation for the while loop associated with
+// `while_ctx`. `grad_inputs` are the partial derivatives w.r.t. the loop
+// outputs, i.e. the exit nodes.  The partial derivatives w.r.t. the loop
+// inputs, i.e. the input loop vars, are returned in `grad_outputs`.
+// `grad_inputs` and `grad_outputs` are both in loop-variable order, as defined
+// by the original inputs to BuildWhileLoop().
+// TODO(skyewm): maybe comment on NoGradient once it's supported
+absl::Status AddWhileLoopGradient(WhileContext* while_ctx, const Scope& scope,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_FRAMEWORK_WHILE_GRADIENTS_H_

tensorflow/cc/framework/while_gradients_test.cc

@@ -0,0 +1,233 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/cc/ops/while_loop.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+class WhileGradientsTest : public ::testing::Test {
+ protected:
+  WhileGradientsTest() : scope_(Scope::NewRootScope()) {}
+
+  void Init(int num_inputs, DataType dtype = DT_INT32) {
+    for (int i = 0; i < num_inputs; ++i) {
+      inputs_.push_back(ops::Placeholder(scope_, dtype));
+    }
+  }
+
+  void CreateLoop(const ops::CondGraphBuilderFn& cond,
+                  const ops::BodyGraphBuilderFn& body,
+                  const std::vector<Output>* inputs = nullptr) {
+    if (inputs == nullptr) inputs = &inputs_;
+    TF_ASSERT_OK(ops::BuildWhileLoop(scope_, *inputs, cond, body, "test_loop",
+                                     &outputs_));
+  }
+
+  void CreateBackprop() {
+    TF_ASSERT_OK(
+        AddSymbolicGradients(scope_, outputs_, inputs_, &grad_outputs_));
+    ASSERT_EQ(grad_outputs_.size(), inputs_.size());
+  }
+
+  template <typename T>
+  void Run(const std::vector<Input::Initializer>& input_values,
+           const std::vector<T>& expected_grad_values) {
+    Run<T>(ClientSession(scope_), input_values, expected_grad_values);
+  }
+
+  template <typename T>
+  void Run(const ClientSession& session,
+           const std::vector<Input::Initializer>& input_values,
+           const std::vector<T>& expected_grad_values,
+           const RunOptions& run_options = RunOptions(),
+           RunMetadata* run_metadata = nullptr) {
+    DCHECK_EQ(input_values.size(), inputs_.size());
+    ClientSession::FeedType feeds;
+    for (int i = 0; i < inputs_.size(); ++i) {
+      feeds.emplace(inputs_[i], input_values[i]);
+    }
+
+    std::vector<Operation> run_outputs;
+    std::vector<Tensor> out_tensors;
+    TF_ASSERT_OK(session.Run(run_options, feeds, grad_outputs_, run_outputs,
+                             &out_tensors, run_metadata));
+    ASSERT_EQ(out_tensors.size(), grad_outputs_.size());
+
+    DCHECK_EQ(expected_grad_values.size(), out_tensors.size());
+    for (int i = 0; i < out_tensors.size(); ++i) {
+      test::ExpectTensorEqual<T>(
+          out_tensors[i], test::AsTensor<T>({expected_grad_values[i]}, {}));
+    }
+  }
+
+  Scope scope_;
+  std::vector<Output> inputs_;
+  std::vector<Output> outputs_;
+  std::vector<Output> grad_outputs_;
+};
+
+TEST_F(WhileGradientsTest, Basic) {
+  // Create loop: while (i < 10) i += 1
+  Init(1);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = ops::Less(s, inputs[0], 10);
+        return s.status();
+      },
+      [](const Scope& s, const std::vector<Output>& inputs,
+         std::vector<Output>* outputs) {
+        // Use AddN, rather than Add, because the gradient function doesn't
+        // depend on the input shapes, and thus we do not need to store
+        // intermediate values in a stack.
+        outputs->push_back(ops::AddN(s, {inputs[0], 1}));
+        return s.status();
+      });
+  CreateBackprop();
+
+  Run<int>({1}, {1});
+  Run<int>({11}, {1});
+}
+
+TEST_F(WhileGradientsTest, MultipleLoopVars) {
+  // Create loop: while (i < 10) i += j; j += 1; k = k
+  Init(3);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = ops::Less(s, inputs[0], 10);
+        return s.status();
+      },
+      [](const Scope& s, const std::vector<Output>& inputs,
+         std::vector<Output>* outputs) {
+        outputs->push_back(ops::AddN(s, {inputs[0], inputs[1]}));
+        outputs->push_back(ops::AddN(s, {inputs[1], 1}));
+        outputs->push_back(inputs[2]);
+        return s.status();
+      });
+  CreateBackprop();
+
+  // The following execution traces illustrate why we expect dF/dj to be 5:
+  //
+  //  i  j  k
+  // ---------
+  //  0  1  2 <-- initial values
+  //  1  2  2
+  //  3  3  2
+  //  6  4  2
+  // 10  5  2 <-- while output values
+  // outputs sum = 17
+  //
+  //  i  j  k
+  // ---------
+  //  0  2  2 <-- initial values (add 1 to j)
+  //  2  3  2
+  //  5  4  2
+  //  9  5  2
+  // 14  6  2 <-- while output values
+  // outputs sum = 22
+  //
+  // Calculate the "slope" between j=1 and j=2:
+  // 22 - 17 = 5 => dF/dj = 5
+  Run<int>({0, 1, 2}, {1, 5, 1});
+
+  Run<int>({1, 1, 0}, {1, 5, 1});
+  Run<int>({0, 0, 0}, {1, 6, 1});
+}
+
+TEST_F(WhileGradientsTest, Chaining) {
+  Init(2, DT_DOUBLE);
+
+  // Multiply each input by 2 before passing to while loop to make sure chaining
+  // works properly
+  std::vector<Output> loop_inputs = {ops::Multiply(scope_, inputs_[0], 2.0),
+                                     ops::Multiply(scope_, inputs_[1], 2.0)};
+
+  // Create loop: while (i > 0 && j > 0) i -= 1
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = ops::LogicalAnd(s, ops::Greater(s, inputs[0], 0.0),
+                                  ops::Greater(s, inputs[1], 0.0));
+        return s.status();
+      },
+      [](const Scope& s, const std::vector<Output>& inputs,
+         std::vector<Output>* outputs) {
+        outputs->push_back(ops::AddN(s, {inputs[0], -1.0}));
+        outputs->push_back(inputs[1]);
+        return s.status();
+      },
+      &loop_inputs);
+
+  // Take negative of first output to make sure chaining works properly
+  outputs_[0] = ops::Neg(scope_, outputs_[0]);
+
+  CreateBackprop();
+
+  Run<double>({1.0, 1.0}, {-2.0, 2.0});
+  Run<double>({0.0, 0.0}, {-2.0, 2.0});
+}
+
+TEST_F(WhileGradientsTest, MultipleDevices) {
+  // Make sure loop is created on cpu0
+  scope_ = scope_.WithDevice("/cpu:0");
+
+  // Create loop: while (i < 10) i += j
+  Init(2);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = ops::Less(s, inputs[0], 10);
+        return s.status();
+      },
+      [](const Scope& s, const std::vector<Output>& inputs,
+         std::vector<Output>* outputs) {
+        // Place body on cpu1
+        Scope cpu1_scope = s.WithDevice("/cpu:1");
+        outputs->push_back(ops::AddN(cpu1_scope, {inputs[0], inputs[1]}));
+        outputs->push_back(inputs[1]);
+        return cpu1_scope.status();
+      });
+
+  // Build gradient graph on cpu1
+  Scope cpu1_scope = scope_.WithDevice("/cpu:1");
+  TF_ASSERT_OK(
+      AddSymbolicGradients(cpu1_scope, outputs_, inputs_, &grad_outputs_));
+  ASSERT_EQ(grad_outputs_.size(), inputs_.size());
+
+  // Run with two CPU devices and output partition graphs
+  SessionOptions session_options;
+  (*session_options.config.mutable_device_count())["CPU"] = 2;
+  RunOptions run_options;
+  run_options.set_output_partition_graphs(true);
+  RunMetadata run_metadata;
+  Run<int>(ClientSession(scope_, session_options), {0, 1}, {1, 11}, run_options,
+           &run_metadata);
+
+  // Check that at least one node ran on each device
+  ASSERT_EQ(run_metadata.partition_graphs().size(), 2);
+  for (const GraphDef& partition_graph : run_metadata.partition_graphs()) {
+    EXPECT_GE(partition_graph.node().size(), 1);
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow