chore: import upstream snapshot with attribution

tensorflow/cc/client/client_session.cc

@@ -0,0 +1,188 @@
+/* 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/client/client_session.h"
+
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+
+class ClientSession::Impl {
+ private:
+  friend class ClientSession;
+
+  Impl(Session* session, std::shared_ptr<Graph> graph)
+      : session_(session), graph_(std::move(graph)) {}
+
+  static SessionOptions MakeDefaultSessionOptions(const std::string& target);
+  absl::Status MaybeExtendGraph() const;
+
+  std::unique_ptr<Session> session_;
+  std::shared_ptr<Graph> graph_;
+
+  mutable mutex mu_;
+  mutable int last_num_graph_nodes_ TF_GUARDED_BY(mu_) = 0;
+};
+
+ClientSession::ClientSession(const Scope& scope, const std::string& target)
+    : ClientSession(scope, Impl::MakeDefaultSessionOptions(target)) {}
+
+ClientSession::ClientSession(const Scope& scope) : ClientSession(scope, "") {}
+
+ClientSession::ClientSession(const Scope& scope,
+                             const SessionOptions& session_options) {
+  Session* new_session;
+  absl::Status status = NewSession(session_options, &new_session);
+  TF_CHECK_OK(status) << status;
+  impl_.reset(new Impl(new_session, scope.graph_as_shared_ptr()));
+  CHECK_NOTNULL(impl()->session_.get());
+}
+
+// Define destructor here so we can forward declare `Impl` in client_session.h.
+// If we define a dtor in the header file or use the default dtor,
+// unique_ptr<Impl> needs the complete type.
+ClientSession::~ClientSession() {}
+
+SessionOptions ClientSession::Impl::MakeDefaultSessionOptions(
+    const std::string& target) {
+  SessionOptions options;
+  options.env = Env::Default();
+  options.target = target;
+  return options;
+}
+
+absl::Status ClientSession::Run(const std::vector<Output>& fetch_outputs,
+                                std::vector<Tensor>* outputs) const {
+  return Run(FeedType{}, fetch_outputs, {}, outputs);
+}
+
+absl::Status ClientSession::Run(const FeedType& inputs,
+                                const std::vector<Output>& fetch_outputs,
+                                std::vector<Tensor>* outputs) const {
+  return Run(inputs, fetch_outputs, {}, outputs);
+}
+
+absl::Status ClientSession::Run(const FeedType& inputs,
+                                const std::vector<Output>& fetch_outputs,
+                                const std::vector<Operation>& run_outputs,
+                                std::vector<Tensor>* outputs) const {
+  return Run(RunOptions(), inputs, fetch_outputs, run_outputs, outputs,
+             nullptr);
+}
+
+absl::Status ClientSession::Impl::MaybeExtendGraph() const {
+  mutex_lock l(mu_);
+  int num_nodes = graph_->num_node_ids();
+  if (num_nodes > last_num_graph_nodes_) {
+    GraphDef graph_def;
+    graph_->ToGraphDefSubRange(&graph_def, last_num_graph_nodes_);
+    last_num_graph_nodes_ = num_nodes;
+    return session_->Extend(graph_def);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ClientSession::Run(const RunOptions& run_options,
+                                const FeedType& inputs,
+                                const std::vector<Output>& fetch_outputs,
+                                const std::vector<Operation>& run_outputs,
+                                std::vector<Tensor>* outputs,
+                                RunMetadata* run_metadata) const {
+  std::vector<std::pair<std::string, Tensor>> feeds;
+  feeds.reserve(inputs.size());
+  for (auto const& feed : inputs) {
+    TF_RETURN_IF_ERROR(feed.second.status);
+    feeds.emplace_back(std::piecewise_construct,
+                       std::forward_as_tuple(feed.first.name()),
+                       std::forward_as_tuple(feed.second.tensor));
+  }
+
+  std::vector<std::string> output_tensor_names;
+  output_tensor_names.reserve(fetch_outputs.size());
+  for (auto const& output : fetch_outputs) {
+    output_tensor_names.push_back(output.name());
+  }
+  std::vector<std::string> target_node_names;
+  target_node_names.reserve(run_outputs.size());
+  for (auto const& output : run_outputs) {
+    target_node_names.push_back(output.node()->name());
+  }
+  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());
+  return impl()->session_->Run(run_options, feeds, output_tensor_names,
+                               target_node_names, outputs, run_metadata);
+}
+
+absl::Status ClientSession::Run(
+    const RunOptions& run_options, const FeedType& inputs,
+    const std::vector<Output>& fetch_outputs,
+    const std::vector<Operation>& run_outputs, std::vector<Tensor>* outputs,
+    RunMetadata* run_metadata,
+    const thread::ThreadPoolOptions& threadpool_options) const {
+  std::vector<std::pair<std::string, Tensor>> feeds;
+  for (auto const& feed : inputs) {
+    TF_RETURN_IF_ERROR(feed.second.status);
+    feeds.emplace_back(feed.first.name(), feed.second.tensor);
+  }
+  std::vector<std::string> output_tensor_names;
+  output_tensor_names.reserve(fetch_outputs.size());
+  for (auto const& output : fetch_outputs) {
+    output_tensor_names.push_back(output.name());
+  }
+  std::vector<std::string> target_node_names;
+  target_node_names.reserve(run_outputs.size());
+  for (auto const& output : run_outputs) {
+    target_node_names.push_back(output.node()->name());
+  }
+  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());
+  return impl()->session_->Run(run_options, feeds, output_tensor_names,
+                               target_node_names, outputs, run_metadata,
+                               threadpool_options);
+}
+
+absl::Status ClientSession::MakeCallable(
+    const CallableOptions& callable_options, CallableHandle* out_handle) {
+  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());
+  return impl()->session_->MakeCallable(callable_options, out_handle);
+}
+
+absl::Status ClientSession::RunCallable(CallableHandle handle,
+                                        const std::vector<Tensor>& feed_tensors,
+                                        std::vector<Tensor>* fetch_tensors,
+                                        RunMetadata* run_metadata) {
+  return impl()->session_->RunCallable(handle, feed_tensors, fetch_tensors,
+                                       run_metadata);
+}
+
+absl::Status ClientSession::RunCallable(
+    CallableHandle handle, const std::vector<Tensor>& feed_tensors,
+    std::vector<Tensor>* fetch_tensors, RunMetadata* run_metadata,
+    const thread::ThreadPoolOptions& options) {
+  return impl()->session_->RunCallable(handle, feed_tensors, fetch_tensors,
+                                       run_metadata, options);
+}
+
+absl::Status ClientSession::ReleaseCallable(CallableHandle handle) {
+  return impl()->session_->ReleaseCallable(handle);
+}
+
+}  // end namespace tensorflow

tensorflow/cc/client/client_session.h

@@ -0,0 +1,164 @@
+/* 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_CLIENT_CLIENT_SESSION_H_
+#define TENSORFLOW_CC_CLIENT_CLIENT_SESSION_H_
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tsl {
+namespace thread {
+struct ThreadPoolOptions;
+}
+}  // namespace tsl
+
+namespace tensorflow {
+
+namespace thread {
+using tsl::thread::ThreadPoolOptions;
+}
+
+/// @addtogroup core
+/// @{
+
+/// A `ClientSession` object lets the caller drive the evaluation of the
+/// TensorFlow graph constructed with the C++ API.
+///
+/// Example:
+///
+///     Scope root = Scope::NewRootScope();
+///     auto a = Placeholder(root, DT_INT32);
+///     auto c = Add(root, a, {41});
+///
+///     ClientSession session(root);
+///     std::vector<Tensor> outputs;
+///
+///     Status s = session.Run({ {a, {1}} }, {c}, &outputs);
+///     if (!s.ok()) { ... }
+class ClientSession {
+ public:
+  /// A data type to represent feeds to a Run call.
+  ///
+  /// This is a map of `Output` objects returned by op-constructors to the value
+  /// to feed them with. See `Input::Initializer` for details on what can be
+  /// used as feed values.
+  typedef std::unordered_map<Output, Input::Initializer, OutputHash> FeedType;
+
+  /// Create a new session to evaluate the graph contained in `scope` by
+  /// connecting to the TensorFlow runtime specified by `target`.
+  ClientSession(const Scope& scope, const std::string& target);
+
+  /// Same as above, but use the empty string ("") as the target specification.
+  explicit ClientSession(const Scope& scope);
+
+  /// Create a new session, configuring it with `session_options`.
+  ClientSession(const Scope& scope, const SessionOptions& session_options);
+
+  ~ClientSession();
+
+  /// Evaluate the tensors in `fetch_outputs`. The values are returned as
+  /// `Tensor` objects in `outputs`. The number and order of `outputs` will
+  /// match `fetch_outputs`.
+  absl::Status Run(const std::vector<Output>& fetch_outputs,
+                   std::vector<Tensor>* outputs) const;
+
+  /// Same as above, but use the mapping in `inputs` as feeds.
+  absl::Status Run(const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   std::vector<Tensor>* outputs) const;
+
+  /// Same as above. Additionally runs the operations ins `run_outputs`.
+  absl::Status Run(const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   const std::vector<Operation>& run_outputs,
+                   std::vector<Tensor>* outputs) const;
+
+  /// Use `run_options` to turn on performance profiling. `run_metadata`, if not
+  /// null, is filled in with the profiling results.
+  absl::Status Run(const RunOptions& run_options, const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   const std::vector<Operation>& run_outputs,
+                   std::vector<Tensor>* outputs,
+                   RunMetadata* run_metadata) const;
+
+  /// Same as above. Additionally allows user to provide custom threadpool
+  /// implementation via ThreadPoolOptions.
+  absl::Status Run(const RunOptions& run_options, const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   const std::vector<Operation>& run_outputs,
+                   std::vector<Tensor>* outputs, RunMetadata* run_metadata,
+                   const thread::ThreadPoolOptions& threadpool_options) const;
+
+  /// \brief A handle to a subgraph, created with
+  /// `ClientSession::MakeCallable()`.
+  typedef int64_t CallableHandle;
+
+  /// \brief Creates a `handle` for invoking the subgraph defined by
+  /// `callable_options`.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status MakeCallable(const CallableOptions& callable_options,
+                            CallableHandle* out_handle);
+
+  /// \brief Invokes the subgraph named by `handle` with the given options and
+  /// input tensors.
+  ///
+  /// The order of tensors in `feed_tensors` must match the order of names in
+  /// `CallableOptions::feed()` and the order of tensors in `fetch_tensors` will
+  /// match the order of names in `CallableOptions::fetch()` when this subgraph
+  /// was created.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status RunCallable(CallableHandle handle,
+                           const std::vector<Tensor>& feed_tensors,
+                           std::vector<Tensor>* fetch_tensors,
+                           RunMetadata* run_metadata);
+
+  /// \brief Invokes the subgraph named by `handle` with the given options and
+  /// input tensors.
+  ///
+  /// The order of tensors in `feed_tensors` must match the order of names in
+  /// `CallableOptions::feed()` and the order of tensors in `fetch_tensors` will
+  /// match the order of names in `CallableOptions::fetch()` when this subgraph
+  /// was created.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status RunCallable(CallableHandle handle,
+                           const std::vector<Tensor>& feed_tensors,
+                           std::vector<Tensor>* fetch_tensors,
+                           RunMetadata* run_metadata,
+                           const thread::ThreadPoolOptions& options);
+
+  /// \brief Releases resources associated with the given `handle` in this
+  /// session.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status ReleaseCallable(CallableHandle handle);
+
+ private:
+  class Impl;
+  std::unique_ptr<Impl> impl_;
+  Impl* impl() { return impl_.get(); }
+  const Impl* impl() const { return impl_.get(); }
+};
+
+/// @}
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_CC_CLIENT_CLIENT_SESSION_H_

tensorflow/cc/client/client_session_test.cc

@@ -0,0 +1,270 @@
+/* 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.
+==============================================================================*/
+
+#define EIGEN_USE_THREADS
+
+#include "tensorflow/cc/client/client_session.h"
+
+#include <utility>
+#include <vector>
+
+#include "absl/synchronization/barrier.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/core/threadpool_options.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/work_sharder.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Add;
+using ops::BatchMatMul;
+using ops::Const;
+using ops::Mul;
+using ops::Placeholder;
+using ops::Sub;
+
+tensorflow::SessionOptions GetSessionOptions() {
+  tensorflow::SessionOptions options;
+  // Disable optimizations for static graph to allow calls to Session::Extend.
+  options.config.mutable_experimental()->set_disable_optimize_for_static_graph(
+      true);
+  return options;
+}
+
+class CustomThreadPoolImpl : public thread::ThreadPoolInterface {
+ public:
+  explicit CustomThreadPoolImpl(int numThreads) {
+    underlying_threadpool_.reset(new thread::ThreadPool(
+        tensorflow::Env::Default(), "custom_threadpool", numThreads));
+    num_schedule_called_ = 0;
+  }
+
+  void Schedule(std::function<void()> fn) override {
+    num_schedule_called_ += 1;
+    underlying_threadpool_->Schedule(std::move(fn));
+  }
+
+  void ScheduleWithHint(std::function<void()> fn, int start, int end) override {
+    num_schedule_called_ += 1;
+    underlying_threadpool_->ScheduleWithHint(std::move(fn), start, end);
+  }
+
+  void Cancel() override {}
+
+  int NumThreads() const override {
+    return underlying_threadpool_->NumThreads();
+  }
+
+  int CurrentThreadId() const override {
+    return underlying_threadpool_->CurrentThreadId();
+  }
+
+  int GetNumScheduleCalled() { return num_schedule_called_; }
+
+ private:
+  int num_schedule_called_;
+  std::unique_ptr<tensorflow::thread::ThreadPool> underlying_threadpool_;
+};
+
+TEST(ClientSessionTest, Basic) {
+  Scope root = Scope::NewRootScope();
+  auto c = Const(root, {{1, 1}});
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  TF_EXPECT_OK(session.Run({c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({1, 1}, {1, 2}));
+}
+
+TEST(ClientSessionTest, Feed) {
+  Scope root = Scope::NewRootScope();
+  auto a = Placeholder(root, DT_INT32);
+  auto b = Placeholder(root, DT_INT32);
+  auto c = Add(root, a, b);
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  TF_EXPECT_OK(session.Run({{a, 1}, {b, 41}}, {c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({42}, {}));
+}
+
+TEST(ClientSessionTest, Extend) {
+  Scope root = Scope::NewRootScope();
+  auto a = Placeholder(root, DT_INT32, Placeholder::Shape({2}));
+  auto c = Add(root, a, {2, 2});
+  ClientSession session(root, GetSessionOptions());
+  std::vector<Tensor> outputs;
+
+  TF_EXPECT_OK(session.Run({{a, {1, 1}}}, {c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({3, 3}, {2}));
+
+  auto d = Add(root, c, {39, 39});
+  outputs.clear();
+  TF_EXPECT_OK(session.Run({{a, {-10, 1}}}, {d}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({31, 42}, {2}));
+}
+
+TEST(ClientSessionTest, MultiThreadedWithDefaultThreadpool) {
+  Scope root = Scope::NewRootScope();
+  auto a = Add(root, {1, 2}, {3, 4});
+  auto b = Mul(root, {1, 2}, {3, 4});
+  ClientSession session(root, GetSessionOptions());
+  {
+    thread::ThreadPool thread_pool(Env::Default(), "pool", 2);
+    thread_pool.Schedule([&session, a]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run({a}, &outputs));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({4, 6}, {2}));
+    });
+    thread_pool.Schedule([&session, b]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run({b}, &outputs));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({3, 8}, {2}));
+    });
+  }
+  auto c = Sub(root, b, a);
+  std::vector<Tensor> outputs;
+  TF_EXPECT_OK(session.Run({c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({-1, 2}, {2}));
+}
+
+TEST(ClientSessionTest, MultiThreadedWithCustomThreadpool) {
+  Scope root = Scope::NewRootScope();
+  int num_threads = 3;
+  auto a = Add(root, {1, 2}, {3, 4});
+  auto b = Mul(root, {1, 2}, {3, 4});
+  ClientSession session(root, GetSessionOptions());
+
+  auto inter_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(inter_op_threadpool->GetNumScheduleCalled(), 0);
+
+  auto intra_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(intra_op_threadpool->GetNumScheduleCalled(), 0);
+
+  tensorflow::thread::ThreadPoolOptions threadPoolOptions;
+  threadPoolOptions.inter_op_threadpool = inter_op_threadpool.get();
+  threadPoolOptions.intra_op_threadpool = intra_op_threadpool.get();
+
+  {
+    thread::ThreadPool thread_pool(Env::Default(), "pool", 2);
+    thread_pool.Schedule([&session, a]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run(RunOptions(), ClientSession::FeedType{}, {a}, {},
+                               &outputs, nullptr, thread::ThreadPoolOptions()));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({4, 6}, {2}));
+    });
+    thread_pool.Schedule([&session, b]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run(RunOptions(), ClientSession::FeedType{}, {b}, {},
+                               &outputs, nullptr, thread::ThreadPoolOptions()));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({3, 8}, {2}));
+    });
+  }
+  auto c = Sub(root, b, a);
+  std::vector<Tensor> outputs;
+  TF_EXPECT_OK(session.Run(RunOptions(), ClientSession::FeedType{}, {c}, {},
+                           &outputs, nullptr, thread::ThreadPoolOptions()));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({-1, 2}, {2}));
+}
+
+TEST(ClientSessionTest, CallableWithDefaultThreadPool) {
+  Scope root = Scope::NewRootScope();
+  auto a = Placeholder(root, DT_INT32);
+  auto b = Placeholder(root, DT_INT32);
+  auto c = Add(root, a, b);
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  CallableOptions options;
+  options.add_feed(a.node()->name());
+  options.add_feed(b.node()->name());
+  options.add_fetch(c.node()->name());
+  ClientSession::CallableHandle callable;
+  TF_CHECK_OK(session.MakeCallable(options, &callable));
+  TF_EXPECT_OK(session.RunCallable(
+      callable, {test::AsTensor<int>({1}, {}), test::AsTensor<int>({41}, {})},
+      &outputs, nullptr));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({42}, {}));
+  TF_EXPECT_OK(session.ReleaseCallable(callable));
+}
+
+TEST(ClientSessionTest, CallableWithCustomThreadPool) {
+  Scope root = Scope::NewRootScope();
+  int num_threads = 3;
+
+  TensorShape data_shape({1, 1});
+  auto a = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
+  auto b = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
+  auto c = BatchMatMul(root, a, b);
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  auto inter_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(inter_op_threadpool->GetNumScheduleCalled(), 0);
+
+  auto intra_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(intra_op_threadpool->GetNumScheduleCalled(), 0);
+
+  tensorflow::thread::ThreadPoolOptions threadPoolOptions;
+  threadPoolOptions.inter_op_threadpool = inter_op_threadpool.get();
+  threadPoolOptions.intra_op_threadpool = intra_op_threadpool.get();
+
+  CallableOptions options;
+  options.add_feed(a.node()->name());
+  options.add_feed(b.node()->name());
+  options.add_fetch(c.node()->name());
+  ClientSession::CallableHandle callable;
+  TF_CHECK_OK(session.MakeCallable(options, &callable));
+
+  // This is needed to have BatchMatMul computation be scheduled in the
+  // intra_op_threadpool.
+  absl::Barrier barrier(num_threads + 1);
+  for (int i = 0; i < num_threads; i++) {
+    intra_op_threadpool->Schedule([&barrier, num_threads]() {
+      tensorflow::SetPerThreadMaxParallelism(num_threads - 1);
+      barrier.Block();
+    });
+  }
+  barrier.Block();
+
+  TF_EXPECT_OK(session.RunCallable(
+      callable,
+      {test::AsTensor<int>({2}, {1, 1}), test::AsTensor<int>({10}, {1, 1})},
+      &outputs, nullptr, threadPoolOptions));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({20}, {1, 1}));
+  TF_EXPECT_OK(session.ReleaseCallable(callable));
+  ASSERT_GT(inter_op_threadpool->GetNumScheduleCalled(), 0);
+  ASSERT_GT(intra_op_threadpool->GetNumScheduleCalled(), 0);
+
+  // Free intra_op_threadpool and wait for its threads to exit before freeing
+  // other objects (e.g. barrier). This is needed to avoid data race.
+  intra_op_threadpool.reset();
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/experimental/base/public/BUILD

@@ -0,0 +1,81 @@
+# Experimental C++ APIs for TensorFlow.
+# New TF C++ APIs under the tensorflow::cc namespace aim to guarantee ABI stability.
+# Users are expected to compile against public c++ headers, and link against
+# libtensorflow (https://www.tensorflow.org/install/lang_c).
+# We aim to achieve ABI stability in new C++ APIs by only using types
+# on the API surface that:
+# 1. Have a header-only implementation
+# 2. Are std:: types
+# 3. Wrap an opaque C type
+
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    # This is intentionally public
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "runtime",
+    hdrs = [
+        "runtime.h",
+    ],
+    deps = [
+        ":status",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+    ],
+)
+
+cc_library(
+    name = "runtime_builder",
+    hdrs = [
+        "runtime_builder.h",
+    ],
+    deps = [
+        ":runtime",
+        ":status",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+    ],
+)
+
+cc_library(
+    name = "status",
+    hdrs = [
+        "status.h",
+    ],
+    deps = [
+        "//tensorflow/c:tf_status",
+    ],
+)
+
+cc_library(
+    name = "tensor",
+    hdrs = [
+        "tensor.h",
+    ],
+    deps = [
+        ":status",
+        "//tensorflow/c:tf_datatype",
+        "//tensorflow/c:tf_tensor",
+    ],
+)
+
+cc_library(
+    name = "tensorhandle",
+    hdrs = [
+        "tensorhandle.h",
+    ],
+    deps = [
+        ":runtime",
+        ":status",
+        ":tensor",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+    ],
+)

tensorflow/cc/experimental/base/public/runtime.h

@@ -0,0 +1,71 @@
+/* Copyright 2020 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_EXPERIMENTAL_BASE_PUBLIC_RUNTIME_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_RUNTIME_H_
+
+#include <memory>
+
+#include "tensorflow/c/eager/c_api_experimental.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// Runtime represents an opaque instance of a Tensorflow runtime, with its own
+// resources, threadpools, etc. Clients are expected to construct a Runtime
+// object through tensorflow::cc::RuntimeBuilder::Build, after setting any
+// relevant configuration options. Many Tensorflow functions take a reference to
+// the runtime as an argument (eg: tensorflow::cc::SavedModelAPI::Load), and
+// may have different implementations depending on the runtime. For many of
+// these Runtime-attached objects (such as tensorflow::cc::TensorHandle), the
+// Runtime must outlive these objects.
+class Runtime {
+ public:
+  // Runtime is movable, but not copyable.
+  Runtime(Runtime&&) = default;
+  Runtime& operator=(Runtime&&) = default;
+
+ private:
+  friend class RuntimeBuilder;
+  friend class SavedModelAPI;
+  friend class TensorHandle;
+
+  // Wraps a TFE_Context. Takes ownership of ctx.
+  explicit Runtime(TFE_Context* ctx) : ctx_(ctx) {}
+
+  // Deletes the currently wrapped TFE_Context, swaps it with ctx,
+  // and takes ownership of ctx.
+  void Reset(TFE_Context* ctx) { ctx_.reset(ctx); }
+
+  // Returns the TFE_Context that this object wraps. This object
+  // retains ownership of the pointer.
+  TFE_Context* GetTFEContext() const { return ctx_.get(); }
+
+  // Runtime is not copyable
+  Runtime(const Runtime&) = delete;
+  Runtime& operator=(const Runtime&) = delete;
+
+  struct TFEContextDeleter {
+    void operator()(TFE_Context* p) const { TFE_DeleteContext(p); }
+  };
+  std::unique_ptr<TFE_Context, TFEContextDeleter> ctx_;
+};
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_RUNTIME_H_

tensorflow/cc/experimental/base/public/runtime_builder.h

@@ -0,0 +1,86 @@
+/* Copyright 2020 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_EXPERIMENTAL_BASE_PUBLIC_RUNTIME_BUILDER_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_RUNTIME_BUILDER_H_
+
+#include <memory>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/cc/experimental/base/public/runtime.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// RuntimeBuilder is a builder used to construct a tensorflow::cc::Runtime.
+// Use this to set configuration options, like threadpool size, etc.
+class RuntimeBuilder {
+ public:
+  RuntimeBuilder() : options_(TFE_NewContextOptions()) {}
+
+  // If `use_tfrt` is true, we will use the new Tensorflow Runtime
+  // (https://blog.tensorflow.org/2020/04/tfrt-new-tensorflow-runtime.html) as
+  // our runtime implementation.
+  RuntimeBuilder& SetUseTFRT(bool use_tfrt);
+
+  // Build a Tensorflow Runtime.
+  //
+  // Params:
+  //  status - Set to OK on success and an appropriate error on failure.
+  // Returns:
+  //  If status is not OK, returns nullptr. Otherwise, returns a
+  //  unique_ptr<tensorflow::cc::Runtime>.
+  std::unique_ptr<Runtime> Build(Status* status);
+
+  // RuntimeBuilder is movable, but not copyable.
+  RuntimeBuilder(RuntimeBuilder&&) = default;
+  RuntimeBuilder& operator=(RuntimeBuilder&&) = default;
+
+ private:
+  // RuntimeBuilder is not copyable
+  RuntimeBuilder(const RuntimeBuilder&) = delete;
+  RuntimeBuilder& operator=(const RuntimeBuilder&) = delete;
+
+  struct TFEContextOptionsDeleter {
+    void operator()(TFE_ContextOptions* p) const {
+      TFE_DeleteContextOptions(p);
+    }
+  };
+  std::unique_ptr<TFE_ContextOptions, TFEContextOptionsDeleter> options_;
+};
+
+inline RuntimeBuilder& RuntimeBuilder::SetUseTFRT(bool use_tfrt) {
+  TFE_ContextOptionsSetTfrt(options_.get(), use_tfrt);
+  return *this;
+}
+
+inline std::unique_ptr<Runtime> RuntimeBuilder::Build(Status* status) {
+  TFE_Context* result = TFE_NewContext(options_.get(), status->GetTFStatus());
+  if (!status->ok()) {
+    return nullptr;
+  }
+  // We can't use std::make_unique here because of its interaction with a
+  // private constructor: https://abseil.io/tips/134
+  return std::unique_ptr<Runtime>(new Runtime(result));
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_RUNTIME_BUILDER_H_

tensorflow/cc/experimental/base/public/status.h

@@ -0,0 +1,96 @@
+/* Copyright 2020 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_EXPERIMENTAL_BASE_PUBLIC_STATUS_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_STATUS_H_
+
+#include <memory>
+#include <string>
+
+#include "tensorflow/c/tf_status.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// Status is a wrapper around an error code and an optional error message.
+// The set of error codes are defined here:
+// https://github.com/tensorflow/tensorflow/blob/08931c1e3e9eb2e26230502d678408e66730826c/tensorflow/c/tf_status.h#L39-L60
+// Many Tensorflow APIs return a Status, or take a Status as an out parameter.
+// Clients should check for status.ok() after calling these APIs, and either
+// handle or propagate the error appropriately.
+// TODO(bmzhao): Add a detailed code example before moving out of experimental.
+class Status {
+ public:
+  // Create a success status
+  Status() : status_(TF_NewStatus()) {}
+
+  // Return the status code
+  TF_Code code() const;
+
+  // Returns the error message in Status.
+  std::string message() const;
+
+  // Returns the error message in Status.
+  bool ok() const;
+
+  // Record <code, msg> in Status. Any previous information is lost.
+  // A common use is to clear a status: SetStatus(TF_OK, "");
+  void SetStatus(TF_Code code, const std::string& msg);
+
+  // Status is movable, but not copyable.
+  Status(Status&&) = default;
+  Status& operator=(Status&&) = default;
+
+ private:
+  friend class RuntimeBuilder;
+  friend class Runtime;
+  friend class SavedModelAPI;
+  friend class TensorHandle;
+
+  // Wraps a TF_Status*, and takes ownership of it.
+  explicit Status(TF_Status* status) : status_(status) {}
+
+  // Status is not copyable
+  Status(const Status&) = delete;
+  Status& operator=(const Status&) = delete;
+
+  // Returns the TF_Status that this object wraps. This object
+  // retains ownership of the pointer.
+  TF_Status* GetTFStatus() const { return status_.get(); }
+
+  struct TFStatusDeleter {
+    void operator()(TF_Status* p) const { TF_DeleteStatus(p); }
+  };
+  std::unique_ptr<TF_Status, TFStatusDeleter> status_;
+};
+
+inline TF_Code Status::code() const { return TF_GetCode(status_.get()); }
+
+inline std::string Status::message() const {
+  return std::string(TF_Message(status_.get()));
+}
+
+inline bool Status::ok() const { return code() == TF_OK; }
+
+inline void Status::SetStatus(TF_Code code, const std::string& msg) {
+  TF_SetStatus(status_.get(), code, msg.c_str());
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_STATUS_H_

tensorflow/cc/experimental/base/public/tensor.h

@@ -0,0 +1,175 @@
+/* Copyright 2020 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_EXPERIMENTAL_BASE_PUBLIC_TENSOR_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_TENSOR_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <functional>
+#include <memory>
+#include <vector>
+
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// Tensor represents an n-dimensional array of values.
+class Tensor {
+ public:
+  using DeleterCallback = std::function<void(void*, size_t)>;
+
+  // Constructs a Tensor from user provided buffer.
+  //
+  // Params:
+  //  dtype - The dtype of the tensor's data.
+  //  shape - A shape vector, where each element corresponds to the size of
+  //          the tensor's corresponding dimension.
+  //  data - Pointer to a buffer of memory to construct a Tensor out of.
+  //  len - The length (in bytes) of `data`
+  //  deleter - A std::function to be called when the Tensor no longer needs the
+  //            memory in `data`. This can be used to free `data`, or
+  //            perhaps decrement a refcount associated with `data`, etc.
+  //  status - Set to OK on success and an error on failure.
+  // Returns:
+  // If an error occurred, status->ok() will be false, and the returned
+  // Tensor must not be used.
+  // TODO(bmzhao): Add Runtime as an argument to this function so we can swap to
+  // a TFRT backed tensor.
+  // TODO(bmzhao): Add benchmarks on overhead for this function; we can
+  // consider using int64_t* + length rather than vector.
+  static Tensor FromBuffer(TF_DataType dtype, const std::vector<int64_t>& shape,
+                           void* data, size_t len, DeleterCallback deleter,
+                           Status* status);
+
+  // TODO(bmzhao): In the case we construct a tensor from non-owned memory,
+  // we should offer a way to deep copy the tensor into a new tensor, which
+  // owns the underlying memory. This could be a .deepcopy()/clone() method.
+
+  // TODO(bmzhao): In the future, we want to relax the non-copyability
+  // constraint. To do so, we can add a C API function that acts like
+  // CopyFrom:
+  // https://github.com/tensorflow/tensorflow/blob/08931c1e3e9eb2e26230502d678408e66730826c/tensorflow/core/framework/tensor.h#L301-L311
+
+  // Tensor is movable, but not copyable
+  Tensor(Tensor&&) = default;
+  Tensor& operator=(Tensor&&) = default;
+
+  // Returns the number of dimensions in the tensor. Can be -1, which represents
+  // unknown rank.
+  int dims() const;
+
+  // Returns the number of elements in dimension `d`.
+  // REQUIRES: `0 <= d < dims()`
+  int64_t dim_size(int d) const;
+
+  // Returns a pointer to the underlying data buffer.
+  void* data() const;
+
+  // Returns the data type of the tensor.
+  TF_DataType dtype() const;
+
+  // Returns the number of elements in the tensor. For a tensor with a partially
+  // defined shape, -1 means not fully defined.
+  int64_t num_elements() const;
+
+  // Returns the size of the underlying data in bytes.
+  size_t num_bytes() const;
+
+ private:
+  friend class TensorHandle;
+  friend class Runtime;
+
+  // Wraps a TF_Tensor. Takes ownership of handle.
+  explicit Tensor(TF_Tensor* tensor) : tensor_(tensor) {}
+
+  // Tensor is not copyable
+  Tensor(const Tensor&) = delete;
+  Tensor& operator=(const Tensor&) = delete;
+
+  // Returns the underlying TF_Tensor that this object wraps.
+  // This object retains ownership of the pointer.
+  TF_Tensor* GetTFTensor() const { return tensor_.get(); }
+
+  struct DeleterStruct {
+    std::function<void(void*, size_t)> deleter;
+  };
+
+  static void DeleterFunction(void* memory, size_t len, void* deleter_struct) {
+    DeleterStruct* deleter = reinterpret_cast<DeleterStruct*>(deleter_struct);
+    deleter->deleter(memory, len);
+    delete deleter;
+  }
+
+  struct TFTensorDeleter {
+    void operator()(TF_Tensor* p) const { TF_DeleteTensor(p); }
+  };
+  std::unique_ptr<TF_Tensor, TFTensorDeleter> tensor_;
+};
+
+inline void* Tensor::data() const { return TF_TensorData(tensor_.get()); }
+
+inline int Tensor::dims() const { return TF_NumDims(tensor_.get()); }
+
+inline int64_t Tensor::dim_size(int d) const {
+  return TF_Dim(tensor_.get(), d);
+}
+
+inline TF_DataType Tensor::dtype() const {
+  return TF_TensorType(tensor_.get());
+}
+
+inline int64_t Tensor::num_elements() const {
+  return TF_TensorElementCount(tensor_.get());
+}
+
+inline size_t Tensor::num_bytes() const {
+  return TF_TensorByteSize(tensor_.get());
+}
+
+inline Tensor Tensor::FromBuffer(TF_DataType dtype,
+                                 const std::vector<int64_t>& shape, void* data,
+                                 size_t len, DeleterCallback deleter,
+                                 Status* status) {
+  // Credit to apassos@ for this technique:
+  // Despite the fact that our API takes a std::function deleter, we are able
+  // to maintain ABI stability because:
+  // 1. Only a function pointer is sent across the C API (&DeleterFunction)
+  // 2. DeleterFunction is defined in the same build artifact that constructed
+  //    the std::function (so there isn't confusion about std::function ABI).
+  // Note that 2. is satisfied by the fact that this is a header-only API, where
+  // the function implementations are inline.
+
+  DeleterStruct* deleter_struct = new DeleterStruct{deleter};
+  TF_Tensor* tensor = TF_NewTensor(dtype, shape.data(), shape.size(), data, len,
+                                   &DeleterFunction, deleter_struct);
+  if (tensor == nullptr) {
+    status->SetStatus(TF_INVALID_ARGUMENT,
+                      "Failed to create tensor for input buffer");
+    return Tensor(nullptr);
+  }
+  return Tensor(tensor);
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_TENSOR_H_

tensorflow/cc/experimental/base/public/tensorhandle.h

@@ -0,0 +1,98 @@
+/* Copyright 2020 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_EXPERIMENTAL_BASE_PUBLIC_TENSORHANDLE_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_TENSORHANDLE_H_
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/cc/experimental/base/public/runtime.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/cc/experimental/base/public/tensor.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// An opaque representation of a tensor computed/managed by the Tensorflow
+// runtime (tensorflow:cc::Runtime). Unlike a tensor, a Tensorhandle may refer
+// to tensors placed in memory of different devices or remote address spaces.
+// Note that tensorflow::cc::Runtime MUST outlive all TensorHandles created
+// from it.
+class TensorHandle {
+ public:
+  // Unwraps a Tensor from the given TensorHandle. If an error occurred,
+  // status->ok() will be false, and the returned Tensor must not be used.
+  Tensor Resolve(Status* status);
+
+  // Constructs a TensorHandle from a Tensor. If an error occurred,
+  // status->ok() will be false, and the returned TensorHandle must not be used.
+  static TensorHandle FromTensor(const Tensor& tensor, const Runtime& runtime,
+                                 Status* status);
+
+  // TensorHandle is movable, and not copyable
+  TensorHandle(TensorHandle&&) = default;
+  TensorHandle& operator=(TensorHandle&&) = default;
+
+ private:
+  // Wraps a TFE_TensorHandle. Takes ownership of handle.
+  explicit TensorHandle(TFE_TensorHandle* handle) : handle_(handle) {}
+
+  // TensorHandle is not copyable
+  TensorHandle(const TensorHandle&) = delete;
+  TensorHandle& operator=(const TensorHandle&) = delete;
+
+  // Returns the underlying TFE_TensorHandle that this object wraps.
+  // This object retains ownership of the pointer.
+  TFE_TensorHandle* GetTFETensorHandle() const { return handle_.get(); }
+
+  // Deletes the currently wrapped TFE_TensorHandle, and swaps it with handle,
+  // and takes ownership of handle.
+  void Reset(TFE_TensorHandle* handle) { handle_.reset(handle); }
+
+  struct TFETensorHandleDeleter {
+    void operator()(TFE_TensorHandle* p) const { TFE_DeleteTensorHandle(p); }
+  };
+  std::unique_ptr<TFE_TensorHandle, TFETensorHandleDeleter> handle_;
+};
+
+inline Tensor TensorHandle::Resolve(Status* status) {
+  TF_Tensor* tensor =
+      TFE_TensorHandleResolve(handle_.get(), status->GetTFStatus());
+  if (!status->ok()) {
+    return Tensor(nullptr);
+  }
+  return Tensor(tensor);
+}
+
+inline TensorHandle TensorHandle::FromTensor(const Tensor& tensor,
+                                             const Runtime& runtime,
+                                             Status* status) {
+  TFE_TensorHandle* tensor_handle = TFE_NewTensorHandleFromTensor(
+      runtime.GetTFEContext(), tensor.GetTFTensor(), status->GetTFStatus());
+  if (!status->ok()) {
+    return TensorHandle(nullptr);
+  }
+  return TensorHandle(tensor_handle);
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_BASE_PUBLIC_TENSORHANDLE_H_

tensorflow/cc/experimental/base/tests/BUILD

@@ -0,0 +1,54 @@
+# Tests for the C++ header-only base types.
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "tensor_types_test_util",
+    testonly = True,
+    hdrs = ["tensor_types_test_util.h"],
+    deps = [
+        "//tensorflow/c:tf_datatype",
+    ],
+)
+
+tf_cc_test(
+    name = "tensor_test",
+    srcs = [
+        "tensor_test.cc",
+    ],
+    deps = [
+        ":tensor_types_test_util",
+        "//tensorflow/c:tf_datatype",
+        "//tensorflow/cc/experimental/base/public:status",
+        "//tensorflow/cc/experimental/base/public:tensor",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tf_cc_test(
+    name = "tensorhandle_test",
+    srcs = [
+        "tensorhandle_test.cc",
+    ],
+    deps = [
+        ":tensor_types_test_util",
+        "//tensorflow/c:tf_datatype",
+        "//tensorflow/cc/experimental/base/public:runtime",
+        "//tensorflow/cc/experimental/base/public:runtime_builder",
+        "//tensorflow/cc/experimental/base/public:status",
+        "//tensorflow/cc/experimental/base/public:tensor",
+        "//tensorflow/cc/experimental/base/public:tensorhandle",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/cc/experimental/base/tests/tensor_test.cc

@@ -0,0 +1,167 @@
+/* Copyright 2020 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/experimental/base/public/tensor.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/cc/experimental/base/tests/tensor_types_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace {
+
+using tensorflow::experimental::cc::Status;
+using tensorflow::experimental::cc::Tensor;
+
+using SimpleTypes = ::testing::Types<
+    tensorflow::FloatType, tensorflow::DoubleType, tensorflow::Int32Type,
+    tensorflow::UINT8Type, tensorflow::INT8Type, tensorflow::INT64Type,
+    tensorflow::UINT16Type, tensorflow::UINT32Type, tensorflow::UINT64Type>;
+
+template <typename T>
+class ConstructScalarTensorTest : public ::testing::Test {};
+TYPED_TEST_SUITE(ConstructScalarTensorTest, SimpleTypes);
+
+// This test constructs a scalar tensor for each of the types in "SimpleTypes",
+// and verifies the expected dimensions, dtype, value, number of bytes, and
+// number of elements.
+TYPED_TEST(ConstructScalarTensorTest, ValidTensorAttributesAfterConstruction) {
+  Status status;
+  TF_DataType dtype = TypeParam::kDType;
+  typename TypeParam::type value = 42;
+  Tensor tensor = Tensor::FromBuffer(/*dtype=*/dtype, /*shape=*/{},
+                                     /*data=*/&value,
+                                     /*len=*/sizeof(value),
+                                     /*deleter=*/[](void*, size_t) {}, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  EXPECT_EQ(tensor.dims(), 0);
+  EXPECT_EQ(tensor.dtype(), dtype);
+  EXPECT_EQ(*reinterpret_cast<typename TypeParam::type*>(tensor.data()), 42);
+  EXPECT_EQ(tensor.num_bytes(), sizeof(typename TypeParam::type));
+  EXPECT_EQ(tensor.num_elements(), 1);
+}
+
+template <typename T>
+class Construct1DTensorTest : public ::testing::Test {};
+TYPED_TEST_SUITE(Construct1DTensorTest, SimpleTypes);
+
+// This test constructs a 1D tensor for each of the types in "SimpleTypes",
+// and verifies the expected dimensions, dtype, value, number of bytes, and
+// number of elements.
+TYPED_TEST(Construct1DTensorTest, ValidTensorAttributesAfterConstruction) {
+  Status status;
+  TF_DataType dtype = TypeParam::kDType;
+  // This is our 1D tensor of varying dtype.
+  std::vector<typename TypeParam::type> value = {42, 100, 0, 1, 4, 29};
+  // Shape is Rank 1 vector.
+  std::vector<int64_t> shape;
+  shape.push_back(value.size());
+
+  Tensor tensor = Tensor::FromBuffer(
+      /*dtype=*/dtype, /*shape=*/shape,
+      /*data=*/value.data(),
+      /*len=*/value.size() * sizeof(typename TypeParam::type),
+      /*deleter=*/[](void*, size_t) {}, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  EXPECT_EQ(tensor.dims(), 1);
+  EXPECT_EQ(tensor.dtype(), dtype);
+  absl::Span<const typename TypeParam::type> tensor_view(
+      reinterpret_cast<typename TypeParam::type*>(tensor.data()), value.size());
+  EXPECT_EQ(tensor_view[0], 42);
+  EXPECT_EQ(tensor_view[1], 100);
+  EXPECT_EQ(tensor_view[2], 0);
+  EXPECT_EQ(tensor_view[3], 1);
+  EXPECT_EQ(tensor_view[4], 4);
+  EXPECT_EQ(tensor_view[5], 29);
+
+  EXPECT_EQ(tensor.num_bytes(),
+            value.size() * sizeof(typename TypeParam::type));
+  EXPECT_EQ(tensor.num_elements(), value.size());
+}
+
+template <typename T>
+class Construct2DTensorTest : public ::testing::Test {};
+TYPED_TEST_SUITE(Construct2DTensorTest, SimpleTypes);
+
+// This test constructs a 2D tensor for each of the types in "SimpleTypes",
+// and verifies the expected dimensions, dtype, value, number of bytes, and
+// number of elements.
+TYPED_TEST(Construct2DTensorTest, ValidTensorAttributesAfterConstruction) {
+  Status status;
+  TF_DataType dtype = TypeParam::kDType;
+  // This is our 1D tensor of varying dtype.
+  std::vector<typename TypeParam::type> value = {42, 100, 0, 1, 4, 29};
+  // Shape is Rank 2 vector with shape 2 x 3.
+  std::vector<int64_t> shape({2, 3});
+
+  Tensor tensor = Tensor::FromBuffer(
+      /*dtype=*/dtype, /*shape=*/shape,
+      /*data=*/value.data(),
+      /*len=*/value.size() * sizeof(typename TypeParam::type),
+      /*deleter=*/[](void*, size_t) {}, &status);
+
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  EXPECT_EQ(tensor.dims(), 2);
+  EXPECT_EQ(tensor.dtype(), dtype);
+  absl::Span<const typename TypeParam::type> tensor_view(
+      reinterpret_cast<typename TypeParam::type*>(tensor.data()), value.size());
+  EXPECT_EQ(tensor_view[0], 42);
+  EXPECT_EQ(tensor_view[1], 100);
+  EXPECT_EQ(tensor_view[2], 0);
+  EXPECT_EQ(tensor_view[3], 1);
+  EXPECT_EQ(tensor_view[4], 4);
+  EXPECT_EQ(tensor_view[5], 29);
+
+  EXPECT_EQ(tensor.num_bytes(),
+            value.size() * sizeof(typename TypeParam::type));
+  EXPECT_EQ(tensor.num_elements(), value.size());
+}
+
+TEST(CPPTensorAPI, ConstructTensorFromBuffer) {
+  bool done = false;
+  Status status;
+  std::vector<int32_t> data_vector({12, 14, 20, 18, 39, 42, 100});
+  {
+    // data_vector is a rank 1 tensor.
+    std::vector<int64_t> shape;
+    shape.push_back(data_vector.size());
+
+    Tensor::DeleterCallback callback = [&done](void* data, size_t len) {
+      done = true;
+    };
+
+    Tensor tensor =
+        Tensor::FromBuffer(/*dtype=*/TF_INT32, /*shape=*/shape,
+                           /*data=*/data_vector.data(),
+                           /*len=*/data_vector.size() * sizeof(int32_t),
+                           /*deleter=*/callback, &status);
+    ASSERT_TRUE(status.ok()) << status.message();
+  }
+  // At this point, tensor has been destroyed, and the deleter callback should
+  // have run.
+  EXPECT_TRUE(done);
+}
+
+}  // namespace

tensorflow/cc/experimental/base/tests/tensor_types_test_util.h

@@ -0,0 +1,76 @@
+/* Copyright 2020 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_EXPERIMENTAL_BASE_TESTS_TENSOR_TYPES_TEST_UTIL_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_BASE_TESTS_TENSOR_TYPES_TEST_UTIL_H_
+
+#include <stdint.h>
+
+#include "tensorflow/c/tf_datatype.h"
+
+namespace tensorflow {
+
+// Each of the following struct types have two members: a kDType that
+// corresponds to a TF_Datatype enum value, and a typedef "type"
+// of its corresponding C++ type. These types allow us to write Dtype-agnostic
+// tests via GoogleTest's TypedTests:
+// https://github.com/google/googletest/blob/e589a337170554c48bc658cc857cf15080c9eacc/googletest/docs/advanced.md#typed-tests
+struct FloatType {
+  using type = float;
+  static constexpr TF_DataType kDType = TF_FLOAT;
+};
+
+struct DoubleType {
+  using type = double;
+  static constexpr TF_DataType kDType = TF_DOUBLE;
+};
+
+struct Int32Type {
+  using type = int32_t;
+  static constexpr TF_DataType kDType = TF_INT32;
+};
+
+struct UINT8Type {
+  using type = uint8_t;
+  static constexpr TF_DataType kDType = TF_UINT8;
+};
+
+struct INT8Type {
+  using type = int8_t;
+  static constexpr TF_DataType kDType = TF_INT8;
+};
+
+struct INT64Type {
+  using type = int64_t;
+  static constexpr TF_DataType kDType = TF_INT64;
+};
+
+struct UINT16Type {
+  using type = uint16_t;
+  static constexpr TF_DataType kDType = TF_UINT16;
+};
+
+struct UINT32Type {
+  using type = uint32_t;
+  static constexpr TF_DataType kDType = TF_UINT32;
+};
+
+struct UINT64Type {
+  using type = uint64_t;
+  static constexpr TF_DataType kDType = TF_UINT64;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_BASE_TESTS_TENSOR_TYPES_TEST_UTIL_H_

tensorflow/cc/experimental/base/tests/tensorhandle_test.cc

@@ -0,0 +1,187 @@
+/* Copyright 2020 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/experimental/base/public/tensorhandle.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/cc/experimental/base/public/runtime.h"
+#include "tensorflow/cc/experimental/base/public/runtime_builder.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/cc/experimental/base/public/tensor.h"
+#include "tensorflow/cc/experimental/base/tests/tensor_types_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+using tensorflow::experimental::cc::Runtime;
+using tensorflow::experimental::cc::RuntimeBuilder;
+using tensorflow::experimental::cc::Status;
+using tensorflow::experimental::cc::Tensor;
+using tensorflow::experimental::cc::TensorHandle;
+
+using SimpleTypes = ::testing::Types<
+    tensorflow::FloatType, tensorflow::DoubleType, tensorflow::Int32Type,
+    tensorflow::UINT8Type, tensorflow::INT8Type, tensorflow::INT64Type,
+    tensorflow::UINT16Type, tensorflow::UINT32Type, tensorflow::UINT64Type>;
+
+template <typename T>
+class ConstructScalarTensorHandleTest : public ::testing::Test {};
+TYPED_TEST_SUITE(ConstructScalarTensorHandleTest, SimpleTypes);
+
+// This test constructs a scalar tensor for each of the types in "SimpleTypes",
+// then wraps it in a TensorHandle. We then unwrap it back into a Tensor, and
+// verify the expected dims, dtype, value, num bytes, and num elements.
+TYPED_TEST(ConstructScalarTensorHandleTest,
+           ValidTensorAttributesAfterConstruction) {
+  Status status;
+  RuntimeBuilder runtime_builder;
+  std::unique_ptr<Runtime> runtime = runtime_builder.Build(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  TF_DataType dtype = TypeParam::kDType;
+  typename TypeParam::type value = 42;
+  Tensor original_tensor =
+      Tensor::FromBuffer(/*dtype=*/dtype, /*shape=*/{},
+                         /*data=*/&value,
+                         /*len=*/sizeof(value),
+                         /*deleter=*/[](void*, size_t) {}, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  TensorHandle handle =
+      TensorHandle::FromTensor(original_tensor, *runtime, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  Tensor tensor = handle.Resolve(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  EXPECT_EQ(tensor.dims(), 0);
+  EXPECT_EQ(tensor.dtype(), dtype);
+  EXPECT_EQ(*reinterpret_cast<typename TypeParam::type*>(tensor.data()), 42);
+  EXPECT_EQ(tensor.num_bytes(), sizeof(typename TypeParam::type));
+  EXPECT_EQ(tensor.num_elements(), 1);
+}
+
+template <typename T>
+class Construct1DTensorHandleTest : public ::testing::Test {};
+TYPED_TEST_SUITE(Construct1DTensorHandleTest, SimpleTypes);
+
+// This test constructs a 1D tensor for each of the types in "SimpleTypes",
+// and verifies the expected dimensions, dtype, value, number of bytes, and
+// number of elements.
+TYPED_TEST(Construct1DTensorHandleTest,
+           ValidTensorAttributesAfterConstruction) {
+  Status status;
+  RuntimeBuilder runtime_builder;
+  std::unique_ptr<Runtime> runtime = runtime_builder.Build(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  TF_DataType dtype = TypeParam::kDType;
+  // This is our 1D tensor of varying dtype.
+  std::vector<typename TypeParam::type> value = {42, 100, 0, 1, 4, 29};
+  // Shape is Rank 1 vector.
+  std::vector<int64_t> shape;
+  shape.push_back(value.size());
+
+  Tensor original_tensor = Tensor::FromBuffer(
+      /*dtype=*/dtype, /*shape=*/shape,
+      /*data=*/value.data(),
+      /*len=*/value.size() * sizeof(typename TypeParam::type),
+      /*deleter=*/[](void*, size_t) {}, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  TensorHandle handle =
+      TensorHandle::FromTensor(original_tensor, *runtime, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  Tensor tensor = handle.Resolve(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  EXPECT_EQ(tensor.dims(), 1);
+  EXPECT_EQ(tensor.dtype(), dtype);
+  absl::Span<const typename TypeParam::type> tensor_view(
+      reinterpret_cast<typename TypeParam::type*>(tensor.data()), value.size());
+  EXPECT_EQ(tensor_view[0], 42);
+  EXPECT_EQ(tensor_view[1], 100);
+  EXPECT_EQ(tensor_view[2], 0);
+  EXPECT_EQ(tensor_view[3], 1);
+  EXPECT_EQ(tensor_view[4], 4);
+  EXPECT_EQ(tensor_view[5], 29);
+
+  EXPECT_EQ(tensor.num_bytes(),
+            value.size() * sizeof(typename TypeParam::type));
+  EXPECT_EQ(tensor.num_elements(), value.size());
+}
+
+template <typename T>
+class Construct2DTensorHandleTest : public ::testing::Test {};
+TYPED_TEST_SUITE(Construct2DTensorHandleTest, SimpleTypes);
+
+// This test constructs a 2D tensor for each of the types in "SimpleTypes",
+// and verifies the expected dimensions, dtype, value, number of bytes, and
+// number of elements.
+TYPED_TEST(Construct2DTensorHandleTest,
+           ValidTensorAttributesAfterConstruction) {
+  Status status;
+  RuntimeBuilder runtime_builder;
+  std::unique_ptr<Runtime> runtime = runtime_builder.Build(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  TF_DataType dtype = TypeParam::kDType;
+  // This is our 1D tensor of varying dtype.
+  std::vector<typename TypeParam::type> value = {42, 100, 0, 1, 4, 29};
+  // Shape is Rank 2 vector with shape 2 x 3.
+  std::vector<int64_t> shape({2, 3});
+
+  Tensor original_tensor = Tensor::FromBuffer(
+      /*dtype=*/dtype, /*shape=*/shape,
+      /*data=*/value.data(),
+      /*len=*/value.size() * sizeof(typename TypeParam::type),
+      /*deleter=*/[](void*, size_t) {}, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  TensorHandle handle =
+      TensorHandle::FromTensor(original_tensor, *runtime, &status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  Tensor tensor = handle.Resolve(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  EXPECT_EQ(tensor.dims(), 2);
+  EXPECT_EQ(tensor.dtype(), dtype);
+  absl::Span<const typename TypeParam::type> tensor_view(
+      reinterpret_cast<typename TypeParam::type*>(tensor.data()), value.size());
+  EXPECT_EQ(tensor_view[0], 42);
+  EXPECT_EQ(tensor_view[1], 100);
+  EXPECT_EQ(tensor_view[2], 0);
+  EXPECT_EQ(tensor_view[3], 1);
+  EXPECT_EQ(tensor_view[4], 4);
+  EXPECT_EQ(tensor_view[5], 29);
+
+  EXPECT_EQ(tensor.num_bytes(),
+            value.size() * sizeof(typename TypeParam::type));
+  EXPECT_EQ(tensor.num_elements(), value.size());
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/experimental/libexport/BUILD

@@ -0,0 +1,67 @@
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+)
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility =
+        [
+            "//tensorflow/python/saved_model:__subpackages__",
+        ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "save",
+    srcs = ["save.cc"],
+    hdrs = ["save.h"],
+    deps = [
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:status",
+    ],
+)
+
+tf_cc_test(
+    name = "save_test",
+    size = "small",
+    srcs = ["save_test.cc"],
+    deps = [
+        ":save",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:path",
+    ],
+)
+
+cc_library(
+    name = "load",
+    srcs = ["load.cc"],
+    hdrs = ["load.h"],
+    deps = [
+        "//tensorflow/cc/saved_model:constants",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:env",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/util/tensor_bundle",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
+tf_cc_test(
+    name = "load_test",
+    size = "small",
+    srcs = ["load_test.cc"],
+    deps = [
+        ":load",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:path",
+    ],
+)

tensorflow/cc/experimental/libexport/load.cc

@@ -0,0 +1,126 @@
+/* Copyright 2021 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/experimental/libexport/load.h"
+
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
+
+namespace tensorflow {
+namespace libexport {
+
+using protobuf::RepeatedPtrField;
+
+absl::StatusOr<TFPackage> TFPackage::Load(const std::string& path) {
+  // Load the proto
+  TFPackage tf_package;
+  const std::string saved_model_pb_path =
+      io::JoinPath(path, kSavedModelFilenamePb);
+  const std::string saved_model_pbtxt_path =
+      io::JoinPath(path, kSavedModelFilenamePbTxt);
+  if (Env::Default()->FileExists(saved_model_pb_path).ok()) {
+    TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), saved_model_pb_path,
+                                       &tf_package.saved_model_proto_));
+  } else if (Env::Default()->FileExists(saved_model_pbtxt_path).ok()) {
+    TF_RETURN_IF_ERROR(ReadTextProto(Env::Default(), saved_model_pbtxt_path,
+                                     &tf_package.saved_model_proto_));
+  } else {
+    return absl::Status(
+        absl::StatusCode::kNotFound,
+        "Could not find SavedModel .pb or .pbtxt at supplied export "
+        "directory path: " +
+            path);
+  }
+
+  // Load the trackable object graph for restoring checkpoint values
+  const std::string variables_dir =
+      tensorflow::io::JoinPath(path, tensorflow::kSavedModelVariablesDirectory);
+  // TODO(b/228181641): revisit non-explicit-checkpoint-loading behavior when
+  // MLAs come along
+  if (Env::Default()->FileExists(variables_dir).ok()) {
+    tf_package.has_checkpoint_ = true;
+    tf_package.variables_filepath_ = tensorflow::io::JoinPath(
+        variables_dir, tensorflow::kSavedModelVariablesFilename);
+    tf_package.variable_reader_ = std::make_unique<tensorflow::BundleReader>(
+        tensorflow::Env::Default(), tf_package.variables_filepath_);
+    tensorflow::Tensor object_graph_tensor;
+    TF_RETURN_IF_ERROR(tf_package.variable_reader_->Lookup(
+        tensorflow::kObjectGraphProtoKey, &object_graph_tensor));
+    const auto* object_graph_string =
+        reinterpret_cast<const tensorflow::tstring*>(
+            object_graph_tensor.tensor_data().data());
+    // TODO(danielellis): make sure parse was successful
+    tf_package.trackable_object_graph_.ParseFromString(*object_graph_string);
+  } else {
+    tf_package.has_checkpoint_ = false;
+    LOG(INFO)
+        << "No checkpoint found, assuming this is a program-only SavedModel";
+  }
+
+  // Build a map of node names to their corresponding nodes.
+  //
+  // See `GetGraphDefNode` for more details.
+  const auto& nodes =
+      tf_package.saved_model_proto_.meta_graphs(0).graph_def().node();
+  for (const auto& node : nodes) {
+    tf_package.graph_def_nodes_by_name_[node.name()] = &node;
+  }
+  return tf_package;
+}
+
+absl::StatusOr<std::string> TFPackage::GetVariableCheckpointKey(int index) {
+  // TODO(danielellis): make sure valid index
+  const auto& trackable_object = trackable_object_graph_.nodes(index);
+  const TrackableObjectGraph::TrackableObject::SerializedTensor*
+      serialized_tensor = nullptr;
+  for (auto& maybe_serialized_tensor : trackable_object.attributes()) {
+    if (maybe_serialized_tensor.name() == "VARIABLE_VALUE") {
+      serialized_tensor = &maybe_serialized_tensor;
+    }
+  }
+  if (serialized_tensor == nullptr) {
+    return absl::Status(absl::StatusCode::kInternal,
+                        "Failed to find variable value field.");
+  }
+  return serialized_tensor->checkpoint_key();
+}
+
+const SavedObjectGraph& TFPackage::GetObjectGraph() {
+  return saved_model_proto_.mutable_meta_graphs(0)->object_graph_def();
+}
+
+absl::StatusOr<const tensorflow::NodeDef*> TFPackage::GetGraphDefNode(
+    std::string name) {
+  const auto& iter = graph_def_nodes_by_name_.find(name);
+  if (iter == graph_def_nodes_by_name_.end()) {
+    return absl::Status(absl::StatusCode::kInternal,
+                        absl::StrCat("Failed to find node named ", name));
+  }
+  return iter->second;
+}
+
+const RepeatedPtrField<FunctionDef>& TFPackage::GetFunctionDefs() {
+  auto& function_library =
+      saved_model_proto_.mutable_meta_graphs(0)->graph_def().library();
+  return function_library.function();
+}
+
+}  // namespace libexport
+}  // namespace tensorflow

tensorflow/cc/experimental/libexport/load.h

@@ -0,0 +1,108 @@
+/* Copyright 2021 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_EXPERIMENTAL_LIBEXPORT_LOAD_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_LIBEXPORT_LOAD_H_
+
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
+#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
+#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
+
+namespace tensorflow {
+namespace libexport {
+
+// A low-level representation of a SavedModel.
+//
+// This class should only ever be a thin wrapper around disk (or other storage)
+// access for a SavedModel.  Higher level functionality should be layered on top
+// by other functions and classes.
+//
+// In the future, this class can also provide a mechanism for automatic version
+// migration.  This will allow the calling code to always work against the most
+// recent version of SavedModel.
+class TFPackage {
+ public:
+  // Load a SavedModel, parsing the associated protobuf for later access.
+  static absl::StatusOr<TFPackage> Load(const std::string& path);
+
+  // Reads and returns a checkpoint key associated with a variable.
+  //
+  // The variable is identified by the index in the object graph node list.
+  //
+  // RestoreV2 is the operation that will ultimately be responsible for reading
+  // and restoring the variable(s)' values.  Variable values are indexed in the
+  // checkpoint files by "checkpoint keys".  These keys along with dtype and
+  // shape / slice information allow RestoreV2 to look up a variable's value in
+  // the SavedModel and restore it into a tensor.
+  absl::StatusOr<std::string> GetVariableCheckpointKey(int index);
+
+  // Retrieves the object graph from the SavedModel.
+  //
+  // For now, we're returning the object graph directly (i.e. the parsed proto)
+  // rather than adding abstraction on top.  We may later find we would like an
+  // intermediate abstraction layer to make traversal easier, but for now the
+  // extra complexity doesn't seem justified.  Regardless of what we choose,
+  // that logic should live outside this class; this class should continue to
+  // have the clearly-defined, singular responsibility of reading and parsing
+  // the low-level, serialized format.
+  const SavedObjectGraph& GetObjectGraph();
+
+  // Retrieves a specific GraphDef node by name.
+  //
+  // GraphDef nodes are stored as a repeating list of nodes.  At module load
+  // time, a module may have constants that need to be restored.  To restore
+  // these constants, they are looked up in the GraphDef's nodes by their name.
+  // Since we may need to load many constants, we create a hash map of these
+  // names to their corresponding nodes at load time in order to look them up
+  // in constant time.
+  absl::StatusOr<const tensorflow::NodeDef*> GetGraphDefNode(std::string name);
+
+  // Returns a list of function defs in the SavedModel.
+  const protobuf::RepeatedPtrField<FunctionDef>& GetFunctionDefs();
+
+  // Returns a BundleReader for reading variable values.
+  //
+  // This TFPackage retains ownership of the underlying reader.
+  tensorflow::BundleReader* GetVariableReader() {
+    return variable_reader_.get();
+  }
+
+  // Returns whether or not we found a valid checkpoint when loading the
+  // package.
+  bool HasCheckpoint() { return has_checkpoint_; }
+
+  // Returns the path to the variables file.
+  const std::string GetVariablesFilepath() const { return variables_filepath_; }
+
+ private:
+  SavedModel saved_model_proto_;
+  TrackableObjectGraph trackable_object_graph_;
+  std::unique_ptr<tensorflow::BundleReader> variable_reader_;
+  std::string variables_filepath_;
+  bool has_checkpoint_;
+  absl::flat_hash_map<std::string, const NodeDef*> graph_def_nodes_by_name_;
+};
+
+}  // namespace libexport
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_LIBEXPORT_LOAD_H_

tensorflow/cc/experimental/libexport/load_test.cc

@@ -0,0 +1,33 @@
+/* Copyright 2021 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/experimental/libexport/load.h"
+
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace libexport {
+namespace {
+
+TEST(LoadTest, TestDiskSavedModelLoad) {
+  absl::StatusOr<TFPackage> result = TFPackage::Load("test");
+  EXPECT_FALSE(result.status().ok());
+}
+
+}  // namespace
+}  // namespace libexport
+}  // namespace tensorflow

tensorflow/cc/experimental/libexport/save.cc

@@ -0,0 +1,28 @@
+/* Copyright 2021 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/experimental/libexport/save.h"
+
+#include "tensorflow/core/platform/env.h"
+
+namespace tensorflow {
+namespace libexport {
+
+absl::Status Save(const std::string& export_dir) {
+  TF_RETURN_IF_ERROR(Env::Default()->RecursivelyCreateDir(export_dir));
+  return absl::OkStatus();
+}
+
+}  // namespace libexport
+}  // namespace tensorflow

tensorflow/cc/experimental/libexport/save.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 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_EXPERIMENTAL_LIBEXPORT_SAVE_H_
+#define TENSORFLOW_CC_EXPERIMENTAL_LIBEXPORT_SAVE_H_
+
+#include <string>
+
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace libexport {
+
+// Writes a saved model to disk.
+//
+// Writes a saved model to the given `export_dir`.
+TF_EXPORT absl::Status Save(const std::string& export_dir);
+
+}  // namespace libexport
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_EXPERIMENTAL_EXPORT_EXPORT_H_

tensorflow/cc/experimental/libexport/save_test.cc

@@ -0,0 +1,36 @@
+/* Copyright 2021 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/experimental/libexport/save.h"
+
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace libexport {
+namespace {
+
+TEST(SaveTest, TestDirectoryStructure) {
+  const std::string base_dir = tensorflow::io::JoinPath(
+      tensorflow::testing::TmpDir(), "test_directory_structure");
+  TF_ASSERT_OK(Save(base_dir));
+  TF_ASSERT_OK(Env::Default()->IsDirectory(base_dir));
+}
+
+}  // namespace
+}  // namespace libexport
+}  // namespace tensorflow

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

tensorflow/cc/gradients/README.md

@@ -0,0 +1,56 @@
+# C++ gradients
+
+Gradients are currently being ported from
+[python](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/ops)
+to C++ (in this directory).
+
+Contributions are welcome and much appreciated; please follow the instructions
+below.
+
+1.  Create the op gradient function in `foo_grad.cc` corresponding to the
+    `foo_grad.py` file where the op originated (i.e. `array_grad.py` op
+    gradients should be written in `array_grad.cc`).
+
+2.  Write the op gradient with the following naming scheme:
+
+    ```
+    Status OpNameGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+      ...
+      return scope.status();
+    }
+    REGISTER_GRADIENT_OP("OpName", OpNameGrad);
+    ```
+
+3.  Ops gradients are implemented by using the
+    [C++ API](https://www.tensorflow.org/api_docs/cc/).
+
+4.  Tests should be included in `foo_grad_test.cc`. Please see
+    [`array_grad_test.cc`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/gradients/array_grad_test.cc)
+    for many examples. Tests are as simple as, creating a placeholder input for
+    the op's inputs and calling `RunTest` (`RunTest` uses a
+    [gradient checker](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/framework/gradient_checker.cc)
+    to verify that the theoretical gradient matches the numeric gradient). For
+    example:
+
+    ```
+    TEST_F(ArrayGradTest, IdentityGrad) {
+      TensorShape shape({5, 2});
+      auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+      auto y = Identity(scope_, x);
+      RunTest(x, shape, y, shape);
+    }
+    ```
+
+NOTE: There are some ops that require features from the C++ API that are not yet
+implemented.
+
+*   Ops that require PartialTensorShape information cannot yet be implemented.
+
+*   Ops that require SparseTensor or IndexSlices (currently only in python)
+    cannot yet be implemented.
+
+*   Maybe more.
+
+For questions: Please create an issue assigned to suharshs.

tensorflow/cc/gradients/array_grad.cc

@@ -0,0 +1,785 @@
+/* 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 <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+REGISTER_NO_GRADIENT_OP("Const");
+REGISTER_NO_GRADIENT_OP("StopGradient");
+REGISTER_NO_GRADIENT_OP("ConcatOffset");
+REGISTER_NO_GRADIENT_OP("EditDistance");
+REGISTER_NO_GRADIENT_OP("ZerosLike");
+REGISTER_NO_GRADIENT_OP("InvertPermutation");
+REGISTER_NO_GRADIENT_OP("Shape");
+REGISTER_NO_GRADIENT_OP("ShapeN");
+REGISTER_NO_GRADIENT_OP("Rank");
+REGISTER_NO_GRADIENT_OP("Size");
+REGISTER_NO_GRADIENT_OP("BroadcastGradientArgs");
+REGISTER_NO_GRADIENT_OP("OneHot");
+
+absl::Status PackGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  int N;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "N", &N));
+  int axis;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));
+
+  grad_outputs->reserve(N);
+  auto grad_op = Unstack(scope, grad_inputs[0], N, Unstack::Axis(axis));
+  for (const Output& o : grad_op.output) {
+    grad_outputs->emplace_back(o);
+  }
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Pack", PackGrad);
+
+absl::Status UnpackGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  int axis;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));
+  grad_outputs->push_back(Stack(scope, grad_inputs, Stack::Axis(axis)));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Unpack", UnpackGrad);
+
+absl::Status IdentityGrad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Identity", IdentityGrad);
+
+absl::Status RefIdentityGrad(const Scope& scope, const Operation& op,
+                             const std::vector<Output>& grad_inputs,
+                             std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("RefIdentity", RefIdentityGrad);
+
+absl::Status QuantizeAndDequantizeGrad(const Scope& scope, const Operation& op,
+                                       const std::vector<Output>& grad_inputs,
+                                       std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("QuantizeAndDequantize", QuantizeAndDequantizeGrad);
+
+absl::Status QuantizeAndDequantizeV4GradHelper(
+    const Scope& scope, const Operation& op,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
+  Input input = Shape(scope, op.input(0));
+  Input input_min = op.input(1);
+  Input input_max = op.input(2);
+  int64_t axis;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));
+  auto qdq_v4_grad = QuantizeAndDequantizeV4Grad(
+      scope, grad_inputs[0], input, input_min, input_max,
+      QuantizeAndDequantizeV4Grad::Axis(axis));
+  grad_outputs->push_back(qdq_v4_grad.input_backprop);
+  grad_outputs->push_back(qdq_v4_grad.input_min_backprop);
+  grad_outputs->push_back(qdq_v4_grad.input_max_backprop);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("QuantizeAndDequantizeV4",
+                     QuantizeAndDequantizeV4GradHelper);
+
+absl::Status QuantizeAndDequantizeV3Grad(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("QuantizeAndDequantizeV3", QuantizeAndDequantizeV3Grad);
+
+absl::Status SplitGrad(const Scope& scope, const Operation& op,
+                       const std::vector<Output>& grad_inputs,
+                       std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(Concat(scope, grad_inputs, op.input(0)));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Split", SplitGrad);
+
+absl::Status SplitVGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() < 3) {
+    return absl::InvalidArgumentError("SplitV requires 3 arguments");
+  }
+  grad_outputs->push_back(Concat(scope, grad_inputs, op.input(2)));
+  for (int i = 0; i < op.num_inputs() - 1; ++i) {
+    grad_outputs->push_back(NoGradient());
+  }
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SplitV", SplitVGrad);
+
+absl::Status FillGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  // y = fill(fill_shape, x)
+  // No gradient returned for the fill_shape argument.
+  grad_outputs->push_back(NoGradient());
+  // The gradient for x (which must be a scalar) is just the sum of
+  // all the gradients from the shape it fills.
+  // We use ReduceSum to implement this, which needs an argument providing
+  // the indices of all the dimensions of the incoming gradient.
+  // grad(x) = reduce_sum(grad(y), [0..rank(grad(y))])
+  auto all_dims = Range(scope, Const(scope, 0), Rank(scope, grad_inputs[0]),
+                        Const(scope, 1));
+  grad_outputs->push_back(ReduceSum(scope, grad_inputs[0], all_dims));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Fill", FillGrad);
+
+absl::Status DiagGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(DiagPart(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Diag", DiagGrad);
+
+absl::Status DiagPartGrad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Diag(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DiagPart", DiagPartGrad);
+
+absl::Status MatrixDiagGrad(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(MatrixDiagPart(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MatrixDiag", MatrixDiagGrad);
+
+absl::Status MatrixBandPartGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  auto num_lower = op.input(1);
+  auto num_upper = op.input(2);
+  grad_outputs->push_back(
+      MatrixBandPart(scope, grad_inputs[0], num_lower, num_upper));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MatrixBandPart", MatrixBandPartGrad);
+
+absl::Status GatherNdGrad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  auto ref = op.input(0);
+  auto indices = op.input(1);
+  Shape::Attrs shape_attrs;
+  shape_attrs.out_type_ = indices.type();
+  auto ref_shape = Shape(scope, ref, shape_attrs);
+  grad_outputs->push_back(ScatterNd(scope, indices, grad_inputs[0], ref_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("GatherNd", GatherNdGrad);
+
+absl::Status CheckNumericsGrad(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string message;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "message", &message));
+  std::string err_msg = absl::StrCat(
+      "Not a number (NaN) or infinity (Inf) values detected in gradient. ",
+      message);
+  grad_outputs->push_back(CheckNumerics(scope, grad_inputs[0], err_msg));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("CheckNumerics", CheckNumericsGrad);
+
+absl::Status ReshapeGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  auto input_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(Reshape(scope, grad_inputs[0], input_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Reshape", ReshapeGrad);
+
+absl::Status ExpandDimsGrad(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto input_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(Reshape(scope, grad_inputs[0], input_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ExpandDims", ExpandDimsGrad);
+
+absl::Status SqueezeGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  auto input_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(Reshape(scope, grad_inputs[0], input_shape));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Squeeze", SqueezeGrad);
+
+absl::Status TransposeGrad(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto inverted_perm = InvertPermutation(scope, op.input(1));
+  grad_outputs->push_back(Transpose(scope, grad_inputs[0], inverted_perm));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Transpose", TransposeGrad);
+
+absl::Status ReverseSequenceGrad(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  auto seq_lengths = op.input(1);
+  int batch_dim;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "batch_dim", &batch_dim));
+  int seq_dim;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "seq_dim", &seq_dim));
+  grad_outputs->push_back(
+      ReverseSequence(scope, grad_inputs[0], seq_lengths, seq_dim,
+                      ReverseSequence::BatchDim(batch_dim)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ReverseSequence", ReverseSequenceGrad);
+
+absl::Status ReverseGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  auto reverse_dims = op.input(1);
+  grad_outputs->push_back(Reverse(scope, grad_inputs[0], reverse_dims));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ReverseV2", ReverseGrad);
+
+absl::Status ScatterNdGrad(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto indices = op.input(0);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(GatherNd(scope, grad_inputs[0], indices));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ScatterNd", ScatterNdGrad);
+
+absl::Status ScatterNdNonAliasingAddGrad(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  auto indices = op.input(1);
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(GatherNd(scope, grad_inputs[0], indices));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ScatterNdNonAliasingAdd", ScatterNdNonAliasingAddGrad);
+
+template <bool IsPadV2>
+absl::Status PadGrad(const Scope& scope, const Operation& op,
+                     const std::vector<Output>& grad_inputs,
+                     std::vector<Output>* grad_outputs) {
+  auto x = op.input(0);
+  auto a = op.input(1);  // [Rank(x), 2]
+  // Takes a slice of a. The 1st column. [Rank(x), 1].
+  auto size = Stack(scope, {Rank(scope, x), 1});
+  auto pad_before = Slice(scope, a, {0, 0}, size);
+  // Make it a 1-D tensor.
+  auto begin = Reshape(scope, pad_before, {-1});
+  grad_outputs->push_back(Slice(scope, grad_inputs[0], begin, Shape(scope, x)));
+  grad_outputs->push_back(NoGradient());
+  // PadV2 adds a "constant_values" input.
+  if (IsPadV2) {
+    grad_outputs->push_back(NoGradient());
+  }
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Pad", PadGrad<false>);
+REGISTER_GRADIENT_OP("PadV2", PadGrad<true>);
+
+absl::Status SpaceToBatchGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(
+      BatchToSpace(scope, grad_inputs[0], op.input(1), block_size));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SpaceToBatch", SpaceToBatchGrad);
+
+absl::Status SpaceToBatchNDGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(
+      BatchToSpaceND(scope, grad_inputs[0], op.input(1), op.input(2)));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SpaceToBatchND", SpaceToBatchNDGrad);
+
+absl::Status BatchToSpaceGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(
+      SpaceToBatch(scope, grad_inputs[0], op.input(1), block_size));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("BatchToSpace", BatchToSpaceGrad);
+
+absl::Status BatchToSpaceNDGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(
+      SpaceToBatchND(scope, grad_inputs[0], op.input(1), op.input(2)));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("BatchToSpaceND", BatchToSpaceNDGrad);
+
+absl::Status SpaceToDepthGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(DepthToSpace(scope, grad_inputs[0], block_size));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SpaceToDepth", SpaceToDepthGrad);
+
+absl::Status DepthToSpaceGrad(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  int block_size;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "block_size", &block_size));
+  grad_outputs->push_back(SpaceToDepth(scope, grad_inputs[0], block_size));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DepthToSpace", DepthToSpaceGrad);
+
+absl::Status MirrorPadGrad(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  std::string mode;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "mode", &mode));
+  grad_outputs->push_back(tensorflow::ops::internal::MirrorPadGrad(
+      scope, grad_inputs[0], op.input(1), mode));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MirrorPad", MirrorPadGrad);
+
+// TODO(suharshs): b/34770860. This gradient was within 1e-3 but not 1e-4.
+absl::Status MirrorPadGradGrad(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string mode;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "mode", &mode));
+  grad_outputs->push_back(MirrorPad(scope, grad_inputs[0], op.input(1), mode));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MirrorPadGrad", MirrorPadGradGrad);
+
+absl::Status StridedSliceGradHelper(const Scope& scope, const Operation& op,
+                                    const std::vector<Output>& grad_inputs,
+                                    std::vector<Output>* grad_outputs) {
+  Input x = Shape(scope, op.input(0));
+  Input begin = op.input(1);
+  Input end = op.input(2);
+  Input strides = op.input(3);
+  int64_t begin_mask;
+  int64_t end_mask;
+  int64_t ellipsis_mask;
+  int64_t new_axis_mask;
+  int64_t shrink_axis_mask;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "begin_mask", &begin_mask));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "end_mask", &end_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "ellipsis_mask", &ellipsis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "new_axis_mask", &new_axis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "shrink_axis_mask", &shrink_axis_mask));
+  grad_outputs->push_back(
+      StridedSliceGrad(scope, x, begin, end, strides, grad_inputs[0],
+                       StridedSliceGrad::BeginMask(begin_mask)
+                           .EndMask(end_mask)
+                           .EllipsisMask(ellipsis_mask)
+                           .NewAxisMask(new_axis_mask)
+                           .ShrinkAxisMask(shrink_axis_mask)));
+  // No gradients returned for begin, end and strides
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("StridedSlice", StridedSliceGradHelper);
+
+absl::Status SliceGrad(const Scope& scope, const Operation& op,
+                       const std::vector<Output>& grad_inputs,
+                       std::vector<Output>* grad_outputs) {
+  // Propagate the incoming gradient along all the selected values,
+  // and zero everywhere else. Use the Pad operator for this.
+  //
+  // First create an Nx2 padding where N is the number of input
+  // dimensions. The first column is the number of prepended zeros
+  // for each dimension, and the second column is the number of
+  // appended zeros.
+  //
+  // The first column is just the begin vector.
+  // The second column is the shape of the input element-wise
+  // subtracted by begin+size
+
+  // Running example:
+  // input.shape = [3, 5, 3]
+  // begin = [1, 2, 1], size = [1, 3, 2]
+  Input input = op.input(0);
+  Input begin = op.input(1);
+  // input_rank = 3
+  auto input_rank = Rank(scope, input);
+  // slice_size = [1, 3, 2]
+  auto slice_size = Shape(scope, op.output(0));
+  // padding_shape = [3, 1]
+  auto padding_shape = Stack(scope, {input_rank, 1});
+  // before_padding = [[1]
+  //                   [2]
+  //                   [1]]
+  Input before_padding = Reshape(scope, begin, padding_shape);
+  // after_padding_sizes = shape(input) - slice_size - begin
+  //                     = [3, 5, 3] - [1, 3, 2] - [1, 2, 1]
+  //                     = [1, 0, 0]
+  auto after_padding_sizes =
+      Sub(scope, Sub(scope, Shape(scope, input), slice_size), begin);
+  // after_padding = [[1]
+  //                  [0]
+  //                  [0]]
+  Input after_padding = Reshape(scope, after_padding_sizes, padding_shape);
+  // paddings = [[1 1]
+  //             [2 0]
+  //             [1 0]]
+  auto paddings =
+      Concat(scope, {before_padding, after_padding}, Const(scope, 1));
+  grad_outputs->push_back(Pad(scope, grad_inputs[0], paddings));
+  // Nothing propagated for "begin" and "size" inputs
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Slice", SliceGrad);
+
+absl::Status ConcatGradHelper(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs,
+                              int start_value_index, int end_value_index,
+                              int dim_index) {
+  if (end_value_index >= op.num_inputs()) {
+    return absl::InternalError("Invalid input index");
+  }
+  std::vector<Output> inputs;
+  inputs.reserve(end_value_index - start_value_index);
+  for (int i = start_value_index; i < end_value_index; ++i) {
+    inputs.push_back(op.input(i));
+  }
+
+  auto shapes = ShapeN(scope, inputs);
+  const auto unique_name = scope.GetUniqueNameForOp("ConcatOffset");
+  auto builder =
+      ::tensorflow::NodeBuilder(unique_name, "ConcatOffset")
+          .Input(::tensorflow::ops::AsNodeOut(scope, op.input(dim_index)))
+          .Input(::tensorflow::ops::AsNodeOutList(scope, shapes.output));
+  scope.UpdateBuilder(&builder);
+  ::tensorflow::Node* concat_offset_node;
+  scope.UpdateStatus(builder.Finalize(scope.graph(), &concat_offset_node));
+  scope.UpdateStatus(scope.DoShapeInference(concat_offset_node));
+  if (concat_offset_node->num_outputs() != inputs.size()) {
+    return absl::InternalError("ConcatOffset has invalid output count");
+  }
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Concat grad should have 1 input");
+  }
+
+  // For each dx[i], we take a slice of dy. The offset and size of the
+  // slice is given by offset[i] and shape[i].
+  const Output& dy = grad_inputs[0];
+  for (int i = 0; i < inputs.size(); ++i) {
+    grad_outputs->push_back(
+        Slice(scope, dy, Output(concat_offset_node, i), shapes.output[i]));
+  }
+
+  // Insert a NoGradient for the axis.
+  grad_outputs->insert(grad_outputs->begin() + dim_index, NoGradient());
+  return scope.status();
+}
+
+absl::Status ConcatV2Grad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  return ConcatGradHelper(scope, op, grad_inputs, grad_outputs,
+                          /*start_value_index=*/0,
+                          /*end_value_index=*/op.num_inputs() - 1,
+                          /*dim+index=*/op.num_inputs() - 1);
+}
+
+REGISTER_GRADIENT_OP("ConcatV2", ConcatV2Grad);
+
+absl::Status BroadcastToGrad(const Scope& scope, const Operation& op,
+                             const std::vector<Output>& grad_inputs,
+                             std::vector<Output>* grad_outputs) {
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError(
+        "BroadcastTo grad should have 1 grad input");
+  }
+  if (op.num_inputs() != 2) {
+    return absl::InvalidArgumentError("BroadcastTo requires 2 inputs");
+  }
+
+  auto x_shape = Shape(scope, op.input(0));
+  auto args = internal::BroadcastGradientArgs(scope, x_shape, op.input(1));
+  auto sum_gx = Sum(scope, grad_inputs[0], args.r0);
+  grad_outputs->push_back(Reshape(scope, sum_gx, x_shape));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("BroadcastTo", BroadcastToGrad);
+
+absl::Status TileGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 2) {
+    return absl::InvalidArgumentError("Tile requires 2 inputs");
+  }
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Tile grad requires 1 grad input");
+  }
+
+  Shape::Attrs shape_attrs;
+  shape_attrs.out_type_ = op.input_type(1);
+  auto input_shape = Shape(scope, op.input(0), shape_attrs);
+  // We interleave multiples and input_shape to get split_shape,
+  // reshape grad to split_shape, and reduce along all even
+  // dimensions (the tiled dimensions) to get the result
+  // with shape input_shape.  For example
+  //   input_shape = [20, 30, 40]
+  //   multiples = [2, 3, 4]
+  //   split_shape = [2, 20, 3, 30, 4, 40]
+  //   axes = [0, 2, 4]
+  auto stack = Stack(scope, {op.input(1), input_shape.output});
+  auto perm = Range(scope, Sub(scope, Rank(scope, stack), 1), -1, -1);
+  auto split_shape = Reshape(scope, Transpose(scope, stack, perm), {-1});
+  auto axes = Range(scope, Const(scope, 0), Size(scope, split_shape.output), 2);
+  auto input_grad = ReduceSum(
+      scope, Reshape(scope, grad_inputs[0], split_shape.output), axes.output);
+  grad_outputs->push_back(input_grad.output);
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Tile", TileGrad);
+
+// Create a constant of the provided d_type;
+Output ConstHelper(const Scope& scope, int value, DataType d_type) {
+  return Cast(scope, Const(scope, value), d_type);
+}
+
+// Adds the batch offsets to the given indices and returns the results.
+Output GetBatchIndices(const Scope& scope, const Output& params_shape,
+                       const Output& indices, int batch_dims) {
+  Output batch_indices = indices;
+  auto indices_ndims = Rank(scope, indices);
+  auto casted_params_shape = Cast(scope, params_shape, indices.type());
+  Output accum_dim_value = ConstHelper(scope, 1, indices.type());
+  for (int dim = batch_dims; dim > 0; dim--) {
+    Output dim_value = Slice(scope, casted_params_shape, {dim - 1}, {1});
+    accum_dim_value = Multiply(scope, accum_dim_value,
+                               Slice(scope, casted_params_shape, {dim}, {1}));
+    auto start = ConstHelper(scope, 0, indices.type());
+    auto step = ConstHelper(scope, 1, indices.type());
+    Output dim_indices = Range(scope, start, Squeeze(scope, dim_value), step);
+    dim_indices = Multiply(scope, dim_indices, accum_dim_value);
+    auto one = Cast(scope, Const(scope, {1}), indices.type());
+    auto dim_shape = Concat(
+        scope,
+        {Output(Tile(scope, one, Const(scope, {dim - 1}))), dim_value,
+         Output(Tile(scope, one,
+                     ExpandDims(scope, Sub(scope, indices_ndims, dim), 0)))},
+        /*axis=*/0);
+    batch_indices =
+        Add(scope, batch_indices, Reshape(scope, dim_indices, dim_shape));
+  }
+
+  return batch_indices;
+}
+
+Output BatchGatherGrad(const Scope& scope, Output params_shape, Output values,
+                       Output indices, int batch_dims, Output gather_dim_size) {
+  // Axis is the first non-batch dimension.
+  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
+  Output outer_shape, flat_values_shape;
+  if (batch_dims != 0) {
+    auto values_shape = Shape(scope, values);
+    // Add the batch offsets to indices and flatten the batch dimensions.
+    outer_shape = Slice(scope, values_shape, {0}, {batch_dims});
+    auto inner_shape =
+        Slice(scope, Slice(scope, values_shape, {batch_dims}, {-1}), {1}, {-1});
+    auto batch_size = Prod(scope, outer_shape, /*axis=*/0);
+    flat_values_shape = Concat(scope, {{-1}, inner_shape}, /*axis=*/0);
+    gather_dim_size = Multiply(scope, gather_dim_size, batch_size);
+    indices = GetBatchIndices(scope, params_shape, indices, batch_dims);
+    values = Reshape(scope, values, flat_values_shape);
+  }
+
+  indices = Reshape(scope, indices, indices_size);
+  Output params_grad =
+      UnsortedSegmentSum(scope, values, indices, gather_dim_size);
+
+  if (batch_dims != 0) {
+    // Put back the batch dimensions.
+    params_grad = Reshape(scope, params_grad, params_shape);
+  }
+  return params_grad;
+}
+
+absl::Status GatherV2Grad(const Scope& scope, const Operation& op,
+                          const std::vector<Output>& grad_inputs,
+                          std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 3) {
+    return absl::InvalidArgumentError("Gather requires 3 inputs");
+  }
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Gather grad requires 1 grad input");
+  }
+
+  // params can be large, so colocate the shape calculation with it.
+  // params can be very large for sparse model, array_ops.shape raises
+  // exception on the Windows platform when any dimension is larger than
+  // int32. params_shape is not used in optimizer apply_sparse gradients,
+  // so it's fine to convert it back to int32 regardless of truncation.
+  auto params = op.input(0);
+  auto colocate_scope = scope.ColocateWith(params);
+  Shape::Attrs shape_attrs;
+  shape_attrs.out_type_ = DT_INT64;
+  auto params_shape64 = Shape(colocate_scope, params, shape_attrs);
+  Output params_shape = Cast(colocate_scope, params_shape64, DT_INT32);
+
+  auto indices = op.input(1);
+  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
+  auto axis = op.input(2);
+  auto axis_expand = ExpandDims(scope, axis, 0);
+
+  int batch_dims;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "batch_dims", &batch_dims));
+  if (batch_dims < 0) {
+    // TODO(bdodson): Figure out if we can find the param rank here, like the
+    // python implementation does.
+    return absl::InvalidArgumentError(
+        "C++ GatherV2 gradient does not support negative batch_dims.");
+  }
+
+  // Handle axis by transposing the axis dimension to be the first non-batch
+  // dimension, compute the gradient and transpose the result back.
+  auto outer_shape = Slice(scope, params_shape, {0}, axis_expand);
+  auto inner_shape =
+      Slice(scope, Slice(scope, params_shape, axis_expand, {-1}), {1}, {-1});
+  auto values_shape = Concat(scope, {outer_shape, {-1}, inner_shape}, 0);
+  auto values_dims = Size(scope, values_shape);
+  auto axis_dims = Size(scope, outer_shape);
+
+  Output outer_batches_indices = Range(scope, 0, batch_dims, /*delta=*/1);
+  Output batch_axis_indices = Range(scope, batch_dims, axis_dims, /*delta=*/1);
+  Output inner_axes_indices =
+      Range(scope, Add(scope, axis_dims, 1), values_dims, /*delta=*/1);
+  Output axis_dims_expand = ExpandDims(scope, axis_dims, 0);
+
+  auto values = Reshape(scope, grad_inputs[0], values_shape);
+
+  // Move values[axis] up to values[batch_dims]
+  Output transpose_dims = Concat(scope,
+                                 {outer_batches_indices, axis_dims_expand,
+                                  batch_axis_indices, inner_axes_indices},
+                                 0);
+  auto values_transpose = Transpose(scope, values, transpose_dims);
+  Output gather_dim_size =
+      Squeeze(scope, Slice(scope, params_shape, axis_expand, {1}));
+  params_shape = Gather(scope, params_shape, transpose_dims);
+
+  auto params_grad = BatchGatherGrad(scope, params_shape, values_transpose,
+                                     indices, batch_dims, gather_dim_size);
+
+  // Inverts the above transpose by moving dimension batch_dims back to its
+  // original position.
+  Output invert_transpose_dims = Concat(scope,
+                                        {outer_batches_indices,
+                                         Add(scope, batch_axis_indices, 1),
+                                         {batch_dims},
+                                         inner_axes_indices},
+                                        0);
+
+  params_grad = Transpose(scope, params_grad, invert_transpose_dims);
+
+  grad_outputs->push_back(params_grad);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("GatherV2", GatherV2Grad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/array_grad_test.cc

@@ -0,0 +1,555 @@
+/* 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 <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using namespace ops;  // NOLINT(build/namespaces)
+using ops::internal::MirrorPadGrad;
+
+class ArrayGradTest : public ::testing::Test {
+ protected:
+  ArrayGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(ArrayGradTest, StackGrad_Axis0) {
+  TensorShape x_shape({1, 2, 3});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  auto y = Stack(scope_, xs, Stack::Axis(0));
+  TensorShape y_shape({2, 1, 2, 3});
+  RunTest(xs, {x_shape, x_shape}, {y}, {y_shape});
+}
+
+TEST_F(ArrayGradTest, StackGrad_Axis1) {
+  TensorShape x_shape({1, 2, 3});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)));
+  auto y = Stack(scope_, xs, Stack::Axis(1));
+  TensorShape y_shape({1, 2, 2, 3});
+  RunTest(xs, {x_shape, x_shape}, {y}, {y_shape});
+}
+
+TEST_F(ArrayGradTest, UnstackGrad_Axis0) {
+  TensorShape x_shape({4, 2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Unstacking the first dimension results in 4 outputs.
+  std::vector<TensorShape> y_shapes(4, TensorShape({2, 3}));
+  auto y = Unstack(scope_, x, 4, Unstack::Axis(0));
+  RunTest({x}, {x_shape}, y.output, y_shapes);
+}
+
+TEST_F(ArrayGradTest, UnstackGrad_Axis1) {
+  TensorShape x_shape({4, 2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Unstacking the second dimension results in 2 outputs.
+  std::vector<TensorShape> y_shapes(2, TensorShape({4, 3}));
+  auto y = Unstack(scope_, x, 2, Unstack::Axis(1));
+  RunTest({x}, {x_shape}, y.output, y_shapes);
+}
+
+TEST_F(ArrayGradTest, IdentityGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Identity(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, SplitGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, 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});
+  RunTest({x}, {x_shape}, y.output, {y_shape, y_shape});
+}
+
+TEST_F(ArrayGradTest, SplitVGrad) {
+  TensorShape x_shape({2, 6});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = SplitV(scope_, x, {1, 2, 3}, /*axis=*/1, /*num_split=*/3);
+  RunTest({x}, {x_shape}, y.output,
+          {TensorShape({2, 1}), TensorShape({2, 2}), TensorShape({2, 3})});
+}
+
+TEST_F(ArrayGradTest, FillGrad) {
+  TensorShape x_shape({});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({2, 5, 3});
+  auto y = Fill(scope_, {2, 5, 3}, x);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, DiagGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Diag(scope_, x);
+  TensorShape y_shape({5, 2, 5, 2});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, DiagPartGrad) {
+  TensorShape x_shape({5, 2, 5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = DiagPart(scope_, x);
+  TensorShape y_shape({5, 2});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MatrixDiagGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = MatrixDiag(scope_, x);
+  TensorShape y_shape({5, 2, 2});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MatrixBandPartGrad) {
+  TensorShape shape({5, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  const int64_t num_lower = 1;
+  const int64_t num_upper = 2;
+  auto y = MatrixBandPart(scope_, x, num_lower, num_upper);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, GatherNdGrad_SimpleIndexing) {
+  TensorShape x_shape({2, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto indices = Const(scope_, {{0, 0}, {1, 1}});
+  TensorShape y_shape({2});
+  auto y = GatherNd(scope_, x, indices);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherNdGrad_SliceIndexing) {
+  TensorShape shape({2, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto indices = Const(scope_, {{1}, {0}});
+  auto y = GatherNd(scope_, x, indices);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, GatherNdGrad_SliceIndexing_Int64) {
+  TensorShape shape({2, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto indices = Cast(scope_, Const(scope_, {{1}, {0}}), DT_INT64);
+  auto y = GatherNd(scope_, x, indices);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, CheckNumericsGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = CheckNumerics(scope_, x, "CheckNumerics failed");
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, ReshapeGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({2, 5});
+  auto y = Reshape(scope_, x, {2, 5});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ExpandDimsGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({1, 5, 2});
+  auto y = ExpandDims(scope_, x, 0);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SqueezeGrad) {
+  TensorShape x_shape({1, 5, 1, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({5, 2});
+  auto y = Squeeze(scope_, x);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, TransposeGrad) {
+  TensorShape x_shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({2, 5});
+  auto y = Transpose(scope_, x, {1, 0});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ReverseSequenceGrad) {
+  TensorShape shape({5, 2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto seq_lengths = Const(scope_, {1, 2, 3, 4, 5});
+  // batch_dim defaults to 0.
+  auto y = ReverseSequence(scope_, x, seq_lengths, /* seq_dim */ 2);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, ReverseGrad) {
+  TensorShape shape({5, 2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Reverse(scope_, x, {0, 2});
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(ArrayGradTest, ScatterNdGrad_SimpleIndexing) {
+  TensorShape updates_shape({4});
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{4}, {3}, {1}, {7}});
+  TensorShape y_shape({8});
+  auto y = ScatterNd(scope_, indices, updates, {8});
+  RunTest(updates, updates_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ScatterNdGrad_SliceIndexing) {
+  TensorShape updates_shape({2, 4, 4});
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{0}, {2}});
+  TensorShape y_shape({4, 4, 4});
+  auto y = ScatterNd(scope_, indices, updates, {4, 4, 4});
+  RunTest(updates, updates_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, ScatterNdNonAliasingAddGrad_SimpleIndexing) {
+  TensorShape updates_shape({4});
+  TensorShape input_shape({8});
+  auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape));
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{4}, {3}, {1}, {7}});
+  auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates);
+  RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape});
+}
+
+TEST_F(ArrayGradTest, ScatterNdNonAliasingAddGrad_SliceIndexing) {
+  TensorShape updates_shape({2, 4, 4});
+  TensorShape input_shape({4, 4, 4});
+  auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape));
+  auto updates =
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
+  auto indices = Const(scope_, {{0}, {2}});
+  auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates);
+  RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape});
+}
+
+TEST_F(ArrayGradTest, PadGrad) {
+  TensorShape x_shape({2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({4, 7});
+  auto y = Pad(scope_, x, paddings);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SpaceToBatchGrad) {
+  TensorShape x_shape({1, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {1, 1}});
+  TensorShape y_shape({4, 2, 2, 1});
+  auto y = SpaceToBatch(scope_, x, paddings, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SpaceToBatchNdGrad) {
+  TensorShape x_shape({2, 2, 4, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto block_shape = Const(scope_, {2, 2});
+  auto paddings = Const(scope_, {{0, 0}, {2, 0}});
+  TensorShape y_shape({8, 1, 3, 1});
+  auto y = SpaceToBatchND(scope_, x, block_shape, paddings);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, BatchToSpaceGrad) {
+  TensorShape x_shape({4, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {1, 1}});
+  TensorShape y_shape({1, 2, 2, 1});
+  auto y = BatchToSpace(scope_, x, paddings, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, BatchToSpaceNdGrad) {
+  TensorShape x_shape({8, 1, 3, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto block_shape = Const(scope_, {2, 2});
+  auto paddings = Const(scope_, {{0, 0}, {2, 0}});
+  TensorShape y_shape({2, 2, 4, 1});
+  auto y = BatchToSpaceND(scope_, x, block_shape, paddings);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, SpaceToDepthGrad) {
+  TensorShape x_shape({1, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({1, 1, 1, 4});
+  auto y = SpaceToDepth(scope_, x, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, DepthToSpaceGrad) {
+  TensorShape x_shape({1, 1, 1, 4});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({1, 2, 2, 1});
+  auto y = DepthToSpace(scope_, x, /* block_size */ 2);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGrad_Reflect) {
+  TensorShape x_shape({2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({4, 7});
+  auto y = MirrorPad(scope_, x, paddings, "REFLECT");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGrad_Symmetric) {
+  TensorShape x_shape({2, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({4, 7});
+  auto y = MirrorPad(scope_, x, paddings, "SYMMETRIC");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGradGrad_Reflect) {
+  TensorShape x_shape({4, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({2, 3});
+  auto y = MirrorPadGrad(scope_, x, paddings, "REFLECT");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, MirrorPadGradGrad_Symmetric) {
+  TensorShape x_shape({4, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto paddings = Const(scope_, {{1, 1}, {2, 2}});
+  TensorShape y_shape({2, 3});
+  auto y = MirrorPadGrad(scope_, x, paddings, "SYMMETRIC");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, StridedSliceGrad) {
+  TensorShape x_shape({6, 4, 4});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+
+  // y = x[2:6:2, 1:3, 1:3]
+  auto y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1});
+  // y.shape = [2, 2, 2];
+  RunTest(x, x_shape, y, {2, 2, 2});
+
+  // y = x[2:6:2, 1:3, 1:3]
+  // begin_mask = 1<<1 (ignore begin_index = 1)
+  // end_mask = 1<<2 (ignore end_index = 2)
+  y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1},
+                   StridedSlice::BeginMask(1 << 1).EndMask(1 << 2));
+  // y.shape = [2, 3, 3];
+  RunTest(x, x_shape, y, {2, 3, 3});
+
+  // y = [tf.newaxis, 2:6:2, 1:3, 1:3]
+  y = StridedSlice(scope_, x, {0, 2, 1, 1}, {0, 6, 3, 3}, {1, 2, 1, 1},
+                   StridedSlice::NewAxisMask(1 << 0));
+  // y.shape = [1, 2, 2, 2];
+  RunTest(x, x_shape, y, {1, 2, 2, 2});
+}
+
+TEST_F(ArrayGradTest, SliceGrad) {
+  TensorShape x_shape({3, 5, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Slice(scope_, x, {1, 2, 1}, {1, 3, 2});
+  RunTest(x, x_shape, y, {1, 3, 2});
+}
+
+TEST_F(ArrayGradTest, ConcatV2Grad) {
+  TensorShape shape({3, 2, 5});
+  std::vector<Output> xs;
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
+  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
+  auto axis = Const(scope_, 0);
+  auto y = Concat(scope_, xs, axis);
+  TensorShape result_shape({9, 2, 5});
+  RunTest(xs, {shape, shape, shape}, {y}, {result_shape});
+}
+
+TEST_F(ArrayGradTest, BroadcastToGrad) {
+  TensorShape x_shape({2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  TensorShape y_shape({3, 2, 5});
+  auto y = BroadcastTo(scope_, x, Const(scope_, {3, 2, 5}));
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, TileGrad) {
+  TensorShape x_shape({2, 5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Tile(scope_, x, Const(scope_, {3, 2}));
+  TensorShape y_shape({6, 10});
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_Simple) {
+  TensorShape shape({100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5}, /*axis=*/0);
+  TensorShape y_shape({4});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_MoreParamDims) {
+  TensorShape shape({100, 2, 3, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5}, /*axis=*/0);
+  TensorShape y_shape({4, 2, 3, 2});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_MoreIndexDims) {
+  TensorShape shape({100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {{2, 0}, {2, 5}}, /*axis=*/0);
+  TensorShape y_shape({2, 2});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_DifferentAxis) {
+  TensorShape shape({2, 10, 10, 2, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5, 5}, /*axis=*/1);
+  TensorShape y_shape({2, 5, 10, 2, 7});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_DifferentAxis2) {
+  TensorShape shape({2, 3, 100, 2, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5, 5}, /*axis=*/2);
+  TensorShape y_shape({2, 3, 5, 2, 7});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_LastAxis) {
+  TensorShape shape({2, 3, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {2, 0, 2, 5, 5}, /*axis=*/2);
+  TensorShape y_shape({2, 3, 5});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_LastAxis2) {
+  TensorShape shape({2, 3, 7, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = GatherV2(scope_, x, {9, 8, 7, 6}, /*axis=*/3);
+  TensorShape y_shape({2, 3, 7, 4});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_BatchDim) {
+  TensorShape shape({2, 100, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 1;
+  auto y =
+      GatherV2(scope_, x, {{2, 0, 2, 5}, {1, 1, 7, 10}}, /*axis=*/1, attrs);
+  TensorShape y_shape({2, 4, 3});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_BatchDim2) {
+  TensorShape shape({2, 19});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 1;
+  auto y = GatherV2(scope_, x, {{0}, {0}}, /*axis=*/1, attrs);
+  TensorShape y_shape({2, 1});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_BatchDimWithAxis) {
+  TensorShape shape({2, 1, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 1;
+  auto y = GatherV2(scope_, x, {{0}, {0}}, /*axis=*/2, attrs);
+  TensorShape y_shape({2, 1, 1});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_TwoBatchDims) {
+  TensorShape shape({2, 2, 100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 2;
+  auto y = GatherV2(scope_, x, {{{2, 0}, {2, 5}}, {{1, 1}, {7, 10}}},
+                    /*axis=*/2, attrs);
+  TensorShape y_shape({2, 2, 2});
+  RunTest(x, shape, y, y_shape);
+}
+
+TEST_F(ArrayGradTest, GatherV2Grad_TwoBatchDimsWithAxis) {
+  TensorShape shape({2, 2, 3, 100});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  GatherV2::Attrs attrs;
+  attrs.batch_dims_ = 2;
+  auto y = GatherV2(scope_, x, {{{2, 0}, {2, 5}}, {{1, 1}, {7, 10}}},
+                    /*axis=*/3, attrs);
+  TensorShape y_shape({2, 2, 2, 3});
+  RunTest(x, shape, y, y_shape);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/data_flow_grad.cc

@@ -0,0 +1,159 @@
+/* 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 <cstdint>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/data_flow_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+REGISTER_NO_GRADIENT_OP("Queue");
+REGISTER_NO_GRADIENT_OP("QueueEnqueue");
+REGISTER_NO_GRADIENT_OP("QueueEnqueueMany");
+REGISTER_NO_GRADIENT_OP("QueueDequeue");
+REGISTER_NO_GRADIENT_OP("QueueDequeueMany");
+REGISTER_NO_GRADIENT_OP("QueueDequeueUpTo");
+REGISTER_NO_GRADIENT_OP("QueueClose");
+REGISTER_NO_GRADIENT_OP("QueueSize");
+REGISTER_NO_GRADIENT_OP("Stack");
+REGISTER_NO_GRADIENT_OP("StackPush");
+REGISTER_NO_GRADIENT_OP("StackPop");
+REGISTER_NO_GRADIENT_OP("StackClose");
+REGISTER_NO_GRADIENT_OP("GetSessionHandle");
+REGISTER_NO_GRADIENT_OP("GetSessionHandleV2");
+REGISTER_NO_GRADIENT_OP("GetSessionTensor");
+REGISTER_NO_GRADIENT_OP("DeleteSessionTensor");
+
+absl::Status DynamicPartitionGrad(const Scope& scope, const Operation& op,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs) {
+  // DynamicPartition only moves input values into various positions
+  // in the output, so the gradient operation only has to map incoming
+  // gradients into their input source locations.
+  // running example:
+  // data = [10, 20, 30, 40, 50]
+  // partitions = [0, 0, 1, 1, 0]
+  // num_partitions = 2
+  // dynamic_partition(data, partitions, num_partitions) = {
+  //   [10, 20, 50],
+  //   [30, 40]
+  // }
+  // grads = {
+  //   [g1, g2, g3],
+  //   [g4, g5]
+  // }
+  // The desired propagation of the gradients back to the data inputs is:
+  // [g1, g2, g4, g5, g3]
+  auto data = op.input(0);
+  auto partitions = op.input(1);
+  int32_t num_partitions;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "num_partitions", &num_partitions));
+
+  // Note: the shape of the partitions is a prefix of the data shape.
+  // shape(partitions) = [5]
+  auto partitions_shape = Shape(scope, partitions);
+  // We now create a partitions-shaped tensor with integers from
+  // [0..size(partitions)) This will be dynamic_partitioned with the
+  // input parameters, providing the destination index for a given
+  // source item.
+  // partitions_size = prod([5]) = 5
+  // reshape(range(partitions_size), [5]) = [0, 1, 2, 3, 4]
+  auto zero = Const(scope, 0);
+  auto one = Const(scope, 1);
+  auto original_indices = Reshape(
+      scope, Range(scope, zero, Prod(scope, partitions_shape, zero), one),
+      partitions_shape);
+  // dynamic_partition(
+  //   [0, 1, 2, 3, 4],
+  //   [0, 0, 1, 1, 0], 2)
+  //  = { [0, 1, 4],
+  //      [2, 3] }
+  auto partitioned_indices =
+      DynamicPartition(scope, original_indices, partitions, num_partitions);
+
+  // Invert these indices with dynamic_stitch to map the incoming
+  // gradients to their source inputs.
+  // dynamic_stitch(
+  //   { [0, 1, 4], [2, 3] },
+  //   { [g1, g2, g3], [g4, g5] })
+  // = [g1, g2, g4, g5, g3]
+  auto reconstructed =
+      DynamicStitch(scope, partitioned_indices.outputs, grad_inputs);
+  // reshape back into a data-shaped tensor to propagate gradients for the data
+  // input.
+  grad_outputs->push_back(Reshape(scope, reconstructed, Shape(scope, data)));
+  // Stop propagation along the partitions input
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DynamicPartition", DynamicPartitionGrad);
+
+absl::Status DynamicStitchGrad(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  // Running example:
+  // indices = {2, [1, 0]}
+  // data = {[d_1, d_2], [[d_3, d_4], [d_5, d_6]]}
+  // out = [[d_5, d_6], [d_3, d_4], [d_1, d_2]]
+  // grad = [[g_1, g_2], [g_3, g_4], [g_5, g_6]]
+
+  // indices and data are two equal-sized lists passed
+  // into DynamicStitch.
+  // num_values = 2
+  int32_t num_values = op.num_inputs() / 2;
+
+  // Stop propagation along the indices list
+  for (int32_t i = 0; i < num_values; i++) {
+    grad_outputs->push_back(NoGradient());
+  }
+
+  // DynamicStitch shuffles its data to the output (using items in
+  // indices) so the gradient propagated to a given data input simply
+  // selects the gradient for its output position.
+  for (int32_t i = 0; i < num_values; i++) {
+    // index has the destination positions for the i'th data
+    // element. We cast it into an int32 if necessary, so we can use
+    // it from a Gather op.
+    // i = 0: index = 2
+    // i = 1: index = [1, 0]
+    auto index = op.input(i);
+    if (index.type() != DT_INT32) {
+      index = Cast(scope, index, DT_INT32);
+    }
+    // Gather the index specified locations in the gradient and
+    // propagate it as the gradient for the i'th data item.
+    // i = 0: gather(grad, 2) = [g_5, g_6]
+    // i = 1: gather(grad, [1, 0]) = [[g_3, g_4], [g_1, g_2]]
+    grad_outputs->push_back(Gather(scope, grad_inputs[0], index));
+  }
+
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DynamicStitch", DynamicStitchGrad);
+REGISTER_GRADIENT_OP("ParallelDynamicStitch", DynamicStitchGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/data_flow_grad_test.cc

@@ -0,0 +1,76 @@
+/* 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 <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/random/random.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Const;
+using ops::DynamicPartition;
+using ops::DynamicStitch;
+using ops::Placeholder;
+
+class DataFlowGradTest : public ::testing::Test {
+ protected:
+  DataFlowGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-4);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(DataFlowGradTest, DynamicPartitionGrad) {
+  TensorShape data_shape({2, 3, 2});
+  auto data = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(data_shape));
+  auto partitions = Const(scope_, {{2, 1, 0}, {1, 2, 0}});
+  auto y = DynamicPartition(scope_, data, partitions, 3);
+  TensorShape partition_shape({2, 2});
+  RunTest({data}, {data_shape}, y.outputs,
+          {partition_shape, partition_shape, partition_shape});
+}
+
+TEST_F(DataFlowGradTest, DynamicStitchGrad) {
+  TensorShape d1_shape({2});
+  TensorShape d2_shape({2, 2});
+  std::vector<Output> indices = {Const(scope_, 2), Const(scope_, {1, 0})};
+  std::vector<Output> data = {
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(d1_shape)),
+      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(d2_shape))};
+  auto y = DynamicStitch(scope_, indices, data);
+  TensorShape y_shape({3, 2});
+  RunTest(data, {d1_shape, d2_shape}, {y}, {y_shape});
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/functional_grad.cc

@@ -0,0 +1,66 @@
+/* Copyright 2020 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 <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/functional_ops.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status PartitionedCallGrad(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  NameAttrList f;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "f", &f));
+  for (const auto& attr : op.node()->attrs()) {
+    (*f.mutable_attr())[attr.first] = attr.second;
+  }
+
+  std::vector<Output> func_inputs;
+  std::vector<DataType> input_dtypes;
+  const int num_inputs = op.num_inputs();
+  func_inputs.reserve(num_inputs + grad_inputs.size());
+  input_dtypes.reserve(num_inputs);
+
+  for (int i = 0; i < num_inputs; i++) {
+    func_inputs.push_back(op.input(i));
+    input_dtypes.push_back(op.input_type(i));
+  }
+
+  func_inputs.insert(std::end(func_inputs), std::begin(grad_inputs),
+                     std::end(grad_inputs));
+
+  auto grad = SymbolicGradient(scope, func_inputs, input_dtypes, f);
+  if (!scope.ok()) return scope.status();
+  for (int i = 0; i < num_inputs; i++) {
+    grad_outputs->push_back(grad[i]);
+  }
+
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("PartitionedCall", PartitionedCallGrad);
+REGISTER_GRADIENT_OP("StatefulPartitionedCall", PartitionedCallGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/functional_grad_test.cc

@@ -0,0 +1,91 @@
+/* Copyright 2020 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 <initializer_list>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/functional_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+class FunctionGradTest : public ::testing::Test {
+ protected:
+  FunctionGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    auto result = (ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error));
+    TF_CHECK_OK(result);
+    TF_ASSERT_OK(result);
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(FunctionGradTest, PartitionedCallGrad) {
+  FunctionDefLibrary f_lib_proto;
+  *(f_lib_proto.add_function()) = test::function::XTimesTwo();
+
+  // Construct a graph:
+  //   A = Placeholder[dtype=int32]
+  //   B = XTimesTwo[_tpu_replicate="cluster"](A)
+  //   C = XTimesTwo[_xla_compile_id="cluster"](A)
+  TF_ASSERT_OK(scope_.graph()->AddFunctionLibrary(f_lib_proto));
+
+  Output x = Placeholder(scope_, DT_FLOAT);
+  NameAttrList f;
+  f.set_name("XTimesTwo");
+  (*f.mutable_attr())["T"].set_type(DT_FLOAT);
+  auto results =
+      PartitionedCall(scope_, std::initializer_list<Input>{x}, {DT_FLOAT}, f);
+  RunTest(x, {}, results[0], {});
+
+  auto stateful_results = StatefulPartitionedCall(
+      scope_, std::initializer_list<Input>{x}, {DT_FLOAT}, f);
+  RunTest(x, {}, stateful_results[0], {});
+}
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/grad_helper.cc

@@ -0,0 +1,82 @@
+/* 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/gradients/grad_helper.h"
+
+#include <vector>
+
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+
+namespace tensorflow {
+
+using tensorflow::ops::Add;
+using tensorflow::ops::Const;
+using tensorflow::ops::DynamicStitch;
+using tensorflow::ops::Mod;
+using tensorflow::ops::OnesLike;
+using tensorflow::ops::Range;
+using tensorflow::ops::Size;
+
+Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
+                          const Output& reduction_axes) {
+  auto zero = Const(scope, 0);
+  auto one = Const(scope, 1);
+
+  // Running example in comments
+  // input_shape = [2, 3, 5, 7]
+  // axes = [1, 2]
+  // The result (a shape after a reduction with keep_dims=True)
+  // [2, 1, 1, 7]
+  //
+  // We can treat each entry in axes as an index into input_shape that
+  // should be replaced by 1.
+  // We use DynamicStitch to do this.
+
+  // input_rank = 4
+  auto input_rank = Size(scope, input_shape);
+
+  // Normalize any negative indices in the reduction_axes to positive
+  // values.
+  auto axes = Mod(scope, Add(scope, reduction_axes, input_rank), input_rank);
+
+  // This [0..input_rank) range of integers is used in DynamicStitch to
+  // first copy input_shape to the result.
+  // input_rank_range = [0, 1, 2, 3]
+  auto input_rank_range = Range(scope, zero, input_rank, one);
+
+  // A 1-filled tensor with the same shape as axes. DynamicStitch will
+  // merge these 1s (using axes for indices) to the correct
+  // position in the result.
+  // axes_ones = [1, 1]
+  auto axes_ones = OnesLike(scope, axes);
+
+  // using DynamicStitch:
+  // indices = { input_rank_range, axes }
+  //         = { [0, 1, 2, 3], [1, 2] }
+  // data = { input_shape, axes_ones }
+  //      = { [2, 3, 5, 7], [1, 1] }
+  // The input_rank_range entry in indices first replicates the
+  // input_shape to the result.
+  // The axes entry in indices then moves a 1 to each of its entries,
+  // resulting in
+  // [2, 1, 1, 7]
+  std::vector<Output> indices = {input_rank_range, axes};
+  std::vector<Output> data = {input_shape, axes_ones};
+  return DynamicStitch(scope, indices, data);
+}
+
+}  // namespace tensorflow

tensorflow/cc/gradients/grad_helper.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_GRADIENTS_GRAD_HELPER_H_
+#define TENSORFLOW_CC_GRADIENTS_GRAD_HELPER_H_
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+
+// Helper function for reduction ops.
+//
+// input_shape: 1-D Tensor, the shape of the Tensor being reduced.
+// axes: 1-D Tensor, the reduction axes.
+//   Note that the reduction indices are in the range
+//   -rank(input_shape), rank(input_shape)
+// returns a 1-D Tensor, the output shape as if keep_dims were set to True.
+Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
+                          const Output& reduction_axes);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_GRADIENTS_GRAD_HELPER_H_

tensorflow/cc/gradients/grad_testutil.cc

@@ -0,0 +1,38 @@
+/* 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/gradients/grad_testutil.h"
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+
+namespace tensorflow {
+namespace test {
+
+absl::Status CallGradFunction(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  ops::GradFunc grad_fn;
+  TF_RETURN_IF_ERROR(ops::GradOpRegistry::Global()->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();
+}
+
+}  // end namespace test
+}  // end namespace tensorflow

tensorflow/cc/gradients/grad_testutil.h

@@ -0,0 +1,38 @@
+/* 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_GRADIENTS_GRAD_TESTUTIL_H_
+#define TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+namespace test {
+
+/// Calls the gradient function registered for 'op', adding gradient operations
+/// to the graph associated with 'scope'. Gradient outputs for each 'op' input
+/// are returned in 'grad_outputs'.
+absl::Status CallGradFunction(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs);
+
+}  // namespace test
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_GRADIENTS_GRAD_TESTUTIL_H_

tensorflow/cc/gradients/image_grad.cc

@@ -0,0 +1,138 @@
+/* Copyright 2018 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 "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/image_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+REGISTER_NO_GRADIENT_OP("NonMaxSuppression");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV2");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV3");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV4");
+REGISTER_NO_GRADIENT_OP("NonMaxSuppressionV5");
+
+absl::Status ResizeNearestNeighborGradHelper(
+    const Scope& scope, const Operation& op,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
+  bool align_corners;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners));
+  bool half_pixel_centers;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers",
+                                 &half_pixel_centers));
+  // The internal gradient implementation needs the shape of the input image.
+  // x_shape = shape(x)[1:3]
+  //         = slice(shape(x), {1}, {3 - 1})
+  auto x_shape = Slice(scope, Shape(scope, op.input(0)), {1}, {2});
+  grad_outputs->push_back(internal::ResizeNearestNeighborGrad(
+      scope, grad_inputs[0], x_shape,
+      internal::ResizeNearestNeighborGrad::AlignCorners(align_corners)
+          .HalfPixelCenters(half_pixel_centers)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ResizeNearestNeighbor", ResizeNearestNeighborGradHelper);
+
+absl::Status ResizeBilinearGradHelper(const Scope& scope, const Operation& op,
+                                      const std::vector<Output>& grad_inputs,
+                                      std::vector<Output>* grad_outputs) {
+  bool align_corners;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners));
+  bool half_pixel_centers;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers",
+                                 &half_pixel_centers));
+  grad_outputs->push_back(internal::ResizeBilinearGrad(
+      scope, grad_inputs[0], op.input(0),
+      internal::ResizeBilinearGrad::AlignCorners(align_corners)
+          .HalfPixelCenters(half_pixel_centers)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ResizeBilinear", ResizeBilinearGradHelper);
+
+absl::Status ResizeBicubicGradHelper(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  bool align_corners;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners));
+  bool half_pixel_centers;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers",
+                                 &half_pixel_centers));
+
+  grad_outputs->push_back(internal::ResizeBicubicGrad(
+      scope, grad_inputs[0], op.input(0),
+      internal::ResizeBicubicGrad::AlignCorners(align_corners)
+          .HalfPixelCenters(half_pixel_centers)));
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ResizeBicubic", ResizeBicubicGradHelper);
+
+absl::Status ScaleAndTranslateGradHelper(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  std::string kernel_type;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "kernel_type", &kernel_type));
+  bool antialias;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "antialias", &antialias));
+  grad_outputs->push_back(internal::ScaleAndTranslateGrad(
+      scope, grad_inputs[0], op.input(0), op.input(2), op.input(3),
+      internal::ScaleAndTranslateGrad::KernelType(kernel_type)
+          .Antialias(antialias)));
+
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("ScaleAndTranslate", ScaleAndTranslateGradHelper);
+
+absl::Status CropAndResizeGradHelper(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  DataType input_type;
+  std::string method;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "method", &method));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "T", &input_type));
+  auto image_shape = Shape(scope, op.input(0));
+  grad_outputs->push_back(CropAndResizeGradImage(
+      scope, grad_inputs[0], op.input(1), op.input(2), image_shape, input_type,
+      CropAndResizeGradImage::Method(method)));
+  grad_outputs->push_back(CropAndResizeGradBoxes(
+      scope, grad_inputs[0], op.input(0), op.input(1), op.input(2)));
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("CropAndResize", CropAndResizeGradHelper);
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/image_grad_test.cc

@@ -0,0 +1,352 @@
+/* Copyright 2018 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 <cassert>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/image_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Const;
+using ops::CropAndResize;
+using ops::ResizeBicubic;
+using ops::ResizeBilinear;
+using ops::ResizeNearestNeighbor;
+using ops::ScaleAndTranslate;
+
+class ImageGradTest : public ::testing::Test {
+ protected:
+  ImageGradTest() : scope_(Scope::NewRootScope()) {}
+
+  enum OpType { RESIZE_NEAREST, RESIZE_BILINEAR, RESIZE_BICUBIC };
+
+  template <typename T>
+  Tensor MakeData(const TensorShape& data_shape) {
+    DataType data_type = DataTypeToEnum<T>::v();
+    Tensor data(data_type, data_shape);
+    auto data_flat = data.flat<T>();
+    for (int i = 0; i < data_flat.size(); ++i) {
+      data_flat(i) = T(i);
+    }
+    return data;
+  }
+
+  template <typename T>
+  void MakeOp(const OpType op_type, const Tensor& x_data, const Input& y_shape,
+              const bool align_corners, const bool half_pixel_centers,
+              Output* x, Output* y) {
+    *x = Const<T>(scope_, x_data);
+    switch (op_type) {
+      case RESIZE_NEAREST:
+        *y = ResizeNearestNeighbor(
+            scope_, *x, y_shape,
+            ResizeNearestNeighbor::AlignCorners(align_corners));
+        return;
+      case RESIZE_BILINEAR:
+        *y = ResizeBilinear(scope_, *x, y_shape,
+                            ResizeBilinear::AlignCorners(align_corners)
+                                .HalfPixelCenters(half_pixel_centers));
+        return;
+      case RESIZE_BICUBIC:
+        *y = ResizeBicubic(scope_, *x, y_shape,
+                           ResizeBicubic::AlignCorners(align_corners)
+                               .HalfPixelCenters(half_pixel_centers));
+        return;
+    }
+    assert(false);
+  }
+
+  template <typename T>
+  void TestResizedShapeForType(const OpType op_type, const bool align_corners,
+                               const bool half_pixel_centers) {
+    TensorShape x_shape({1, 2, 2, 1});
+    Tensor x_data = MakeData<T>(x_shape);
+    Output x, y;
+    MakeOp<T>(op_type, x_data, {4, 6}, align_corners, half_pixel_centers, &x,
+              &y);
+
+    ClientSession session(scope_);
+    std::vector<Tensor> outputs;
+    TF_ASSERT_OK(session.Run({y}, &outputs));
+    EXPECT_EQ(outputs.size(), 1);
+    EXPECT_EQ(outputs[0].shape(), TensorShape({1, 4, 6, 1}));
+  }
+
+  void TestResizedShape(OpType op_type) {
+    for (const bool half_pixel_centers : {true, false}) {
+      for (const bool align_corners : {true, false}) {
+        if (half_pixel_centers && align_corners) {
+          continue;
+        }
+        TestResizedShapeForType<Eigen::half>(op_type, align_corners,
+                                             half_pixel_centers);
+        TestResizedShapeForType<float>(op_type, align_corners,
+                                       half_pixel_centers);
+        TestResizedShapeForType<double>(op_type, align_corners,
+                                        half_pixel_centers);
+      }
+    }
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestResizeToSmallerAndAlign(const OpType op_type,
+                                   const bool align_corners,
+                                   const bool half_pixel_centers) {
+    TensorShape x_shape({1, 4, 6, 1});
+    Tensor x_data = MakeData<X_T>(x_shape);
+    Output x, y;
+    MakeOp<X_T>(op_type, x_data, {2, 3}, align_corners, half_pixel_centers, &x,
+                &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, 2, 3, 1}, &max_error)));
+    EXPECT_LT(max_error, 1.5e-3);
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestResizeToLargerAndAlign(const OpType op_type,
+                                  const bool align_corners,
+                                  const bool half_pixel_centers) {
+    TensorShape x_shape({1, 2, 3, 1});
+    Tensor x_data = MakeData<X_T>(x_shape);
+    Output x, y;
+    MakeOp<X_T>(op_type, x_data, {4, 6}, align_corners, half_pixel_centers, &x,
+                &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, 4, 6, 1}, &max_error)));
+    EXPECT_LT(max_error, 1.5e-3);
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestResize(OpType op_type) {
+    for (const bool half_pixel_centers : {true, false}) {
+      for (const bool align_corners : {true, false}) {
+        // if (!half_pixel_centers) continue;
+        if (half_pixel_centers && align_corners) {
+          continue;
+        }
+        TestResizeToSmallerAndAlign<X_T, Y_T, JAC_T>(op_type, align_corners,
+                                                     half_pixel_centers);
+        TestResizeToLargerAndAlign<X_T, Y_T, JAC_T>(op_type, align_corners,
+                                                    half_pixel_centers);
+      }
+    }
+  }
+
+  Scope scope_;
+};
+
+TEST_F(ImageGradTest, TestNearestNeighbor) {
+  TestResizedShape(RESIZE_NEAREST);
+  TestResize<float, float, float>(RESIZE_NEAREST);
+  TestResize<double, double, double>(RESIZE_NEAREST);
+}
+
+TEST_F(ImageGradTest, TestBilinear) {
+  TestResizedShape(RESIZE_BILINEAR);
+  TestResize<float, float, float>(RESIZE_BILINEAR);
+  // Note that Y_T is always float for this op. We choose
+  // double for the jacobian to capture the higher precision
+  // between X_T and Y_T.
+  TestResize<double, float, double>(RESIZE_BILINEAR);
+}
+
+TEST_F(ImageGradTest, TestBicubic) {
+  TestResizedShape(RESIZE_BICUBIC);
+  TestResize<float, float, float>(RESIZE_BICUBIC);
+  // Note that Y_T is always float for this op. We choose
+  // double for the jacobian to capture the higher precision
+  // between X_T and Y_T.
+  TestResize<double, float, double>(RESIZE_BICUBIC);
+}
+
+class ScaleAndTranslateGradTest : public ::testing::Test {
+ protected:
+  ScaleAndTranslateGradTest() : scope_(Scope::NewRootScope()) {}
+
+  template <typename T>
+  Tensor MakeData(const TensorShape& data_shape) {
+    DataType data_type = DataTypeToEnum<T>::v();
+    Tensor data(data_type, data_shape);
+    auto data_flat = data.flat<T>();
+    for (int i = 0; i < data_flat.size(); ++i) {
+      data_flat(i) = T(i);
+    }
+    return data;
+  }
+
+  template <typename T>
+  void MakeOp(const Tensor& x_data, const Input& y_shape, Input scale,
+              Input translation, const std::string& kernel_type, bool antialias,
+              Output* x, Output* y) {
+    *x = Const<T>(scope_, x_data);
+    *y = ScaleAndTranslate(scope_, *x, y_shape, scale, translation,
+                           ScaleAndTranslate::KernelType(kernel_type)
+                               .Antialias(antialias)
+                               .Antialias(antialias));
+    TF_ASSERT_OK(scope_.status());
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestScaleAndTranslate(const TensorShape x_shape, const int out_height,
+                             const int out_width, Input scale,
+                             Input translation, const std::string& kernel_type,
+                             bool antialias) {
+    Tensor x_data = MakeData<X_T>(x_shape);
+    Output x, y;
+    MakeOp<X_T>(x_data, {out_height, out_width}, scale, translation,
+                kernel_type, antialias, &x, &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, out_height, out_width, 1}, &max_error)));
+    EXPECT_LT(max_error, 2e-3);
+  }
+
+  const std::vector<Input> kScales = {Input{1.0f, 1.0f}, Input{0.37f, 0.47f},
+                                      Input{2.1f, 2.1f}};
+  const std::vector<Input> kTranslations = {
+      Input{0.0f, 0.0f}, Input{3.14f, 1.19f}, Input{2.1f, 3.1f},
+      Input{100.0f, 200.0f}};
+  Scope scope_;
+};
+
+TEST_F(ScaleAndTranslateGradTest, TestGrads) {
+  const std::vector<std::string> kKernelTypes = {"lanczos1", "lanczos3",
+                                                 "lanczos5", "gaussian"};
+  constexpr int kOutHeight = 4;
+  constexpr int kOutWidth = 6;
+
+  const TensorShape kXShape = TensorShape({1, 2, 3, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      for (const std::string& kernel_type : kKernelTypes) {
+        TestScaleAndTranslate<float, float, float>(
+            kXShape, kOutHeight, kOutWidth, scale, translation, kernel_type,
+            true);
+      }
+    }
+  }
+}
+
+TEST_F(ScaleAndTranslateGradTest, TestGradsWithoutAntialias) {
+  constexpr int kOutHeight = 4;
+  constexpr int kOutWidth = 6;
+
+  const TensorShape kXShape = TensorShape({1, 2, 3, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      TestScaleAndTranslate<float, float, float>(kXShape, kOutHeight, kOutWidth,
+                                                 scale, translation, "lanczos3",
+                                                 false);
+    }
+  }
+}
+
+TEST_F(ScaleAndTranslateGradTest, TestGradsWithSameShape) {
+  const std::vector<std::string> kKernelTypes = {"lanczos3", "gaussian"};
+
+  constexpr int kOutHeight = 2;
+  constexpr int kOutWidth = 3;
+
+  const TensorShape kXShape = TensorShape({1, 2, 3, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      for (const std::string& kernel_type : kKernelTypes) {
+        TestScaleAndTranslate<float, float, float>(
+            kXShape, kOutHeight, kOutWidth, scale, translation, kernel_type,
+            true);
+      }
+    }
+  }
+}
+
+TEST_F(ScaleAndTranslateGradTest, TestGradsWithSmallerShape) {
+  const std::vector<std::string> kKernelTypes = {"lanczos3", "gaussian"};
+  constexpr int kOutHeight = 2;
+  constexpr int kOutWidth = 3;
+
+  const TensorShape kXShape = TensorShape({1, 4, 6, 1});
+  for (const Input scale : kScales) {
+    for (const Input translation : kTranslations) {
+      for (const std::string& kernel_type : kKernelTypes) {
+        TestScaleAndTranslate<float, float, float>(
+            kXShape, kOutHeight, kOutWidth, scale, translation, kernel_type,
+            true);
+      }
+    }
+  }
+}
+
+class CropAndResizeGradTest : public ::testing::Test {
+ protected:
+  CropAndResizeGradTest() : scope_(Scope::NewRootScope()) {}
+
+  template <typename T>
+  Tensor MakeData(const TensorShape& data_shape) {
+    DataType data_type = DataTypeToEnum<T>::v();
+    Tensor data(data_type, data_shape);
+    auto data_flat = data.flat<T>();
+    for (int i = 0; i < data_flat.size(); ++i) {
+      data_flat(i) = T(i);
+    }
+    return data;
+  }
+
+  template <typename T>
+  void MakeOp(const Tensor& x_data, const Input& boxes, const Input& box_ind,
+              const Input& crop_size, Output* x, Output* y) {
+    *x = Const<T>(scope_, x_data);
+    *y = CropAndResize(scope_, *x, boxes, box_ind, crop_size,
+                       CropAndResize::Method("bilinear"));
+    TF_ASSERT_OK(scope_.status());
+  }
+
+  template <typename X_T, typename Y_T, typename JAC_T>
+  void TestCropAndResize() {
+    TensorShape x_shape({1, 4, 2, 1});
+    Tensor x_data = MakeData<X_T>(x_shape);
+    TensorShape box_shape({1, 4});
+    Tensor boxes = MakeData<X_T>(box_shape);
+    Output x, y;
+    MakeOp<X_T>(x_data, boxes, {0}, {1, 1}, &x, &y);
+    JAC_T max_error;
+    TF_ASSERT_OK((ComputeGradientError<X_T, Y_T, JAC_T>(
+        scope_, x, x_data, y, {1, 1, 1, 1}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(CropAndResizeGradTest, TestCrop) {
+  TestCropAndResize<float, float, float>();
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/linalg_grad.cc

@@ -0,0 +1,481 @@
+/* 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 <algorithm>
+#include <optional>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/btree_set.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/gradients/grad_helper.h"
+#include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/math_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+constexpr absl::string_view kEllipsis = "...";
+
+// Returns the axis (possibly negative) corresponding to a label.
+//
+// Returns the axis index of the axis label if it is before an ellipsis (or if
+// the ellipsis is not present), and the negative index if it occurs after the
+// ellipsis. E.g. index of `b` in `ab...cd`, is `1`, but that of `c` is `-2`.
+//
+// For multiple occurrences, returns the leftmost one. If not found, returns
+// absl::nullopt.
+//
+// Parameters:
+//   subscripts: A string denoting the einsum subscript (e.g. `ab...cd`)
+//   label: The single character axis label.
+std::optional<int> EinsumGetAxisFromLabel(absl::string_view subscripts,
+                                          char label) {
+  std::vector<absl::string_view> splits = absl::StrSplit(subscripts, kEllipsis);
+  auto index = splits[0].find(label);
+  if (index != splits[0].npos) {
+    return index;
+  }
+  if (splits.size() < 2) {
+    return std::nullopt;
+  }
+  index = splits[1].find(label);
+  if (index != splits[1].npos) {
+    return index - splits[1].length();
+  }
+  return std::nullopt;
+}
+
+// Returns a tuple denoting the slice mapping to ellipsis.
+//
+// For a given subscript, returns a tuple (start, end) denoting the start
+// axis index and the (negative) end axis index respectively. For any input
+// Tensor `x` described by the subscript, `x[start:end]` would be the slice
+// represented by the ellipsis. E.g. For `ab...cd` returns `[1, -2]`.
+//
+// If ellipsis is not present in `subscripts`, returns `(0, 0)`.
+//
+// Parameters:
+//   subscripts: A string denoting the einsum subscript.
+//   start: Output for the start index
+//   end: Output for the end index (or nullopt to go to the end).
+std::tuple<int, std::optional<int>> EinsumGetBcastSubshape(
+    absl::string_view subscripts) {
+  int start = subscripts.find(kEllipsis);
+  if (start == subscripts.npos) {
+    return std::make_tuple(0, 0);
+  }
+  int remaining = subscripts.length() - (start + kEllipsis.length());
+  std::optional<int> end;
+  if (remaining > 0) {
+    end = -remaining;
+  } else {
+    end = std::nullopt;
+  }
+  return std::make_tuple(start, end);
+}
+
+// Slices elements of a 1d tensor from [start,end].
+// If end is nullopt, it goes to the end of the tensor.
+// Supports negative values for end.
+// This attempts to give the same result as tenspr[start:end] would give in
+// Python.
+Output Slice1dHelper(const Scope& scope, Output tensor, int start,
+                     std::optional<int> end) {
+  if (end.has_value() && *end > 0) {
+    return Slice(scope, tensor, Const(scope, start, TensorShape({1})),
+                 Const(scope, *end - start, TensorShape({1})));
+  } else {
+    return Slice(scope, tensor, Const(scope, start, TensorShape({1})),
+                 Add(scope, Shape(scope, tensor), end.value_or(0) - start));
+  }
+}
+
+// Returns reduced subscripts and their corresponding dimensions and axes.
+//
+// Given a set of axis labels, returns their concatenated subscript, their
+// corresponding dimensions from input_shape, and their corresponding axes.
+// Note that the concatenated subscript `reduced_subs` may have axis labels
+// from `reduced_label_set` in any order. For example, for the reduced label
+// set `{b, d}`, subscripts `aabbcd` and input shape `[2,2,5,5,3,4]`, returns
+// subscripts `bd`, dimensions `[5,4]` and axes `[2,5]`.
+//
+// Args:
+//   reduced_label_set: Set of axis labels which appear in `subscripts`.
+//   input_shape: A `Tensor` representing the shape of the einsum operand
+//     corresponding to `subscripts`.
+//   subscripts: A string denoting the einsum subscript.
+//
+// Returns:
+//   reduced_subs: Subscripts formed by a concatenation of labels in
+//     `reduced_label_set`.
+//   reduced_dims: Dimensions from `input_shape` corresponding to each label
+//     in `reduced_subs`.
+//   reduced_axes: Axes described by `subscripts` corresponding to each label
+//     in `reduced_subs`. If there are multiple occurrences in `subscripts`,
+//     we consider only the leftmost one.
+std::tuple<std::string, Output, Output> EinsumGetReducedSubscripts(
+    const Scope& scope, const absl::btree_set<char>& reduced_label_set,
+    Output input_shape, absl::string_view subscripts) {
+  // Concatenate the sequence of reduced axis labels.
+  const std::string reduced_subs =
+      std::string(reduced_label_set.begin(), reduced_label_set.end());
+  // Get the axis (may be positive, negative or zero) for each of the reduced
+  // labels. If the same label appears multiple times, get the left-most axis.
+  std::vector<int> reduced_axes;
+  reduced_axes.reserve(reduced_subs.size());
+  for (const char s : reduced_subs) {
+    auto axis = EinsumGetAxisFromLabel(subscripts, s);
+    if (!axis.has_value()) {
+      // Should never happen.
+      scope.UpdateStatus(absl::InternalError(
+          absl::StrCat("Missing axis", absl::string_view(&s, 1))));
+    } else {
+      reduced_axes.push_back(*axis);
+    }
+  }
+  // Get the corresponding dimensions for each reduced axis.
+  std::vector<Output> reduced_dims_inputs;
+  reduced_dims_inputs.reserve(reduced_axes.size());
+  for (const int i : reduced_axes) {
+    if (i < 0) {
+      reduced_dims_inputs.push_back(
+          Gather(scope, input_shape, Add(scope, Size(scope, input_shape), i)));
+    } else {
+      reduced_dims_inputs.push_back(Gather(scope, input_shape, i));
+    }
+  }
+  const Output reduced_dims = Stack(scope, reduced_dims_inputs);
+  Tensor reduced_axes_tensor(
+      DataType::DT_INT32, TensorShape({static_cast<int>(reduced_axes.size())}));
+  std::copy_n(reduced_axes.begin(), reduced_axes.size(),
+              reduced_axes_tensor.flat<int>().data());
+  return std::make_tuple(reduced_subs, reduced_dims,
+                         Const(scope, reduced_axes_tensor));
+}
+
+// Returns the gradient wrt input for a unary einsum with reductions.
+//
+//  scope: Scope for grad operations.
+//  output_grad: The gradient wrt the output of a unary einsum operation.
+//  output_subs: The output subscript. (E.g. `ac` for equation `abc->ac`).
+//  input_subs: The input subscript. (E.g. `abc` for equation `abc->ac`).
+//  input_shape: The shape of the input operand.
+//  reduced_label_set: The set of axis labels appearing in `input_subs` but
+//    not in `output_subs`.
+Output EinsumGradReducedHelper(const Scope& scope, const Output& output_grad,
+                               absl::string_view output_subs,
+                               absl::string_view input_subs,
+                               const Output& input_shape,
+                               const absl::btree_set<char>& reduced_label_set) {
+  // Let's say the einsum operation was "aabbcd->ca", where axis labels 'b' and
+  // 'd' are reduced with input_shape [2,2,5,5,3,4]. Then obtain the reduced
+  // subscripts "bd", corresponding dimensions [5,4] and axes [2,5].
+  std::string reduced_subs;
+  Output reduced_dims, reduced_axes;
+  std::tie(reduced_subs, reduced_dims, reduced_axes) =
+      EinsumGetReducedSubscripts(scope, reduced_label_set, input_shape,
+                                 input_subs);
+  // Whether either the input or the output subscripts have a repeated label.
+  // This is true for "aabbcd->ca" or "abd->cca" but false for "abcd->ca".
+  const int distinct_input_labels =
+      absl::flat_hash_set<char>(input_subs.begin(), input_subs.end()).size();
+  const int distinct_output_labels =
+      absl::flat_hash_set<char>(output_subs.begin(), output_subs.end()).size();
+  const bool has_repeated_labels =
+      (distinct_input_labels + distinct_output_labels) <
+      input_subs.length() + output_subs.length();
+  // Compute the input subscripts without the reduced axis labels, e.g. "aac"
+  // for the equation "aabbcd->ca".
+  std::string input_subs_without_reduced_labels;
+  for (const char s : input_subs) {
+    if (!absl::c_linear_search(reduced_label_set, s)) {
+      input_subs_without_reduced_labels.push_back(s);
+    }
+  }
+
+  // The gradient wrt the input for the equation "abc->ac" (or, equivalently
+  // reduce_sum(..., axis=1)) is just the gradient of the output tiled N times
+  // along axis 1, where label 'b' represents a dimension of size N.
+  //
+  // If we're not dealing with repeated labels, and the non-reduced labels
+  // doesn't need to be transposed, then just tiling is enough and there is no
+  // need to call another einsum. For example, tiling is sufficient for
+  // "abcd->ac". But for equations like "aabbcd->ac" (generalized traces) or
+  // "abc->ca" (transpose), we'd need another einsum operation after tiling.
+  if (!has_repeated_labels &&
+      input_subs_without_reduced_labels == output_subs) {
+    // Obtain the shape of the output, as if keepdims=True on reduce sum. E.g.
+    // for the equation "abcd->ac" with input shape [2,5,3,4], we get the
+    // reduced shape [2,1,3,1].
+    auto reduced_shape = ReducedShapeHelper(scope, input_shape, reduced_axes);
+    // Reshaping the gradient (wrt "ac") to [2,1,3,1] and broadcasting it to
+    // the shape [2,5,3,4] results in the gradient wrt "abcd".
+    return BroadcastTo(scope, Reshape(scope, output_grad, reduced_shape),
+                       input_shape);
+  }
+
+  // If we *do* have traces or transpose operations, then prepend the extra
+  // reduced dimensions to the front. E.g. Given the equation "aabbcd->ca" we'd
+  // first obtain the VJP for "bdca->ca", and then the VJP for "aabbcd->bdca".
+  //
+  // Obtain the input shape with reduced dimensions prepended, viz. [5,4,3,2].
+  // This is the shape of the intermediate "bdca".
+  Output output_grad_shape = Shape(scope, output_grad);
+  auto grad_shape_with_reduced_labels =
+      Concat(scope, {reduced_dims, output_grad_shape}, /*axis=*/0);
+
+  // Obtain the output shape of the reduction-only equation "bdca->ca" as if
+  // keepdims=True; viz. [1,1,3,2]. Since we prepended the reduced labels,
+  // we just have to prepend that many 1s to the output shape.
+
+  auto reduced_shape = Concat(
+      scope,
+      {Const(scope, 1, TensorShape{static_cast<int>(reduced_label_set.size())}),
+       output_grad_shape},
+      /*axis=*/0);
+  // Compute the VJP for the intermediate (viz. "bdca->ca") for which
+  // broadcasting is sufficient.
+  Output broadcasted_grad =
+      BroadcastTo(scope, Reshape(scope, output_grad, reduced_shape),
+                  grad_shape_with_reduced_labels);
+  // Compute the VJP for the final step (viz. "aabbcd->bdca"). We can
+  // use einsum with the input and output subscripts reversed (viz.
+  // "bdca->aabbcd") since the output axis labels now appear in the
+  // input subscripts.
+  return Einsum(scope, {broadcasted_grad},
+                absl::StrCat(reduced_subs, output_subs, "->", input_subs));
+}
+
+// Returns the gradient wrt an input operand for a binary einsum.
+//
+// This function does not handle (un)broadcasting. This must be done separately
+// on the returned gradient.
+//
+// Args:
+//   output_grad: The gradient wrt the output of a binary einsum operation.
+//   other_operand: The complementary `Tensor` operand i.e. which is not the
+//     input operand.
+//   input_shape: A `Tensor` representing the shape of input operand.
+//   input_subs: The subscripts of the input operand.
+//   other_subs: The subscripts of the complementary operand.
+//   output_subs: The output subscripts.
+Output EinsumGradWrt(const Scope& scope, Output output_grad,
+                     Output other_operand, Output input_shape,
+                     absl::string_view input_subs, absl::string_view other_subs,
+                     absl::string_view output_subs) {
+  // Claim: For the einsum operation z = einsum("{eq_x},{eq_y}->{eq_z}", x, y),
+  //   where the equation involves only Tensor contractions, generalized traces
+  //   and transposes, the input gradients are given by the vector-jacobian
+  //   products (VJPs):
+  //
+  //     grad_wrt_x = einsum("{eq_y},{eq_z}->{eq_x}", y, grad_wrt_z)
+  //     grad_wrt_y = einsum("{eq_x},{eq_z}->{eq_y}", x, grad_wrt_z}
+  //
+  //   where grad_wrt_x and grad_wrt_y are the gradients with respect to inputs
+  //   x and y and grad_wrt_z is the given gradient with respect to output z.
+  //
+  // Proof: For unary einsum equations involving only transpose ("ij->ji") and
+  //   traces ("ii->i"), the linear mapping's Jacobian at input x is given
+  //   by the function itself. We can verify that the linear map given by the
+  //   VJP are einsums with the equations "ji->ij" and "i->ii" respectively,
+  //   where the latter represents 'un-tracing', or filling the diagonal with
+  //   the input axis and non-diagonal entries are zeros.
+  //        Furthermore, recall that matrix multiplication, which is
+  //   represented by the equation "ab,bc->ac", has its VJPs given by the
+  //   einsum equations "ac,bc->ab" and "ab,ac->bc" (see, for example
+  //   https://math.stackexchange.com/a/2755680). Combined with transposes and
+  //   traces we can rewrite Tensor contractions as regular matrix
+  //   multiplication. Since each of these operations have their VJPs described
+  //   by einsums of the required pattern, the result follows.
+  //
+  // Accordingly, einsum operations except for those with reductions, e.g.
+  // "abc,cd->ad" have their VJPs defined by:
+  //   "{output_subs},{other_subs}->{input_subs}".
+  //
+  // But if there is a reduction, this would lead to the equation "ad,cd->abc"
+  // which is invalid because the reduced axis label 'b' is present in the
+  // output but not in any of the inputs. Therefore, we compute the VJP in two
+  // steps: first we obtain VJP for "ac,cd->ad" and then we compute the VJP of
+  // "abc->ac" or, equivalently, reduce_sum(..., axis=1).
+  //
+  // Compute the set of input axis labels which doesn't appear in either the
+  // output subscripts or the other operand's subscript. E.g. the set {'b'} for
+  // the equation "abc,cd->ad".
+  absl::btree_set<char> reduced_label_set(input_subs.begin(), input_subs.end());
+  for (const char x : output_subs) {
+    reduced_label_set.erase(x);
+  }
+  for (const char x : other_subs) {
+    reduced_label_set.erase(x);
+  }
+  reduced_label_set.erase('.');
+
+  // Obtain the input subscripts with the reduced axis labels removed. E.g.
+  // "ac" in the above example.
+  std::string left_subs;
+  for (const char s : input_subs) {
+    if (!reduced_label_set.contains(s)) {
+      left_subs.push_back(s);
+    }
+  }
+
+  // Compute the gradient wrt the input, without accounting for the operation
+  // "abc->ac". So, now we have the VJP of the operation "ac,cd->ad".
+  Output grad_reduced =
+      Einsum(scope, {output_grad, other_operand},
+             absl::StrCat(output_subs, ",", other_subs, "->", left_subs));
+
+  // If the reduced_label_set is empty, then we already have the gradient
+  // wrt the input.
+  if (reduced_label_set.empty()) {
+    return grad_reduced;
+  }
+  // Otherwise, we currently have the gradient wrt the output of the reduction
+  // operation "abc->ac". Invoke the subroutine for the gradient for unary
+  // einsum with reductions.
+  return EinsumGradReducedHelper(scope, grad_reduced, left_subs, input_subs,
+                                 input_shape, reduced_label_set);
+}
+
+absl::Status EinsumGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  if (grad_inputs.size() != 1) {
+    return absl::InvalidArgumentError("Expect 1 grad input.");
+  }
+  const Output& grad = grad_inputs[0];
+
+  std::string equation;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "equation", &equation));
+  std::vector<absl::string_view> equation_split =
+      absl::StrSplit(equation, "->");
+  if (equation_split.size() != 2) {
+    return absl::InvalidArgumentError("Equation must contain a single ->");
+  }
+
+  const absl::string_view input_subs = equation_split[0];
+  const absl::string_view output_subs = equation_split[1];
+  if (op.num_inputs() == 1) {
+    // For the unary einsum z = einsum("{eq_x}->{eq_z}", x), the gradient wrt
+    // the input (VJP) is given by the reversed equation:
+    //   grad_wrt_x = einsum("{eq_z}->{eq_x}", grad_wrt_z)
+    // (See the justification in _GetGradWrt). This is valid unless there are
+    // reduced axis labels; i.e. axis labels appearing in the input but not in
+    // the output subscripts.
+    auto input_shape = Shape(scope, op.input(0));
+    // Find the axis labels which appear only in the input.
+    absl::btree_set<char> reduced_label_set(input_subs.begin(),
+                                            input_subs.end());
+    for (const char x : output_subs) {
+      reduced_label_set.erase(x);
+    }
+    reduced_label_set.erase('.');
+    if (reduced_label_set.empty()) {
+      grad_outputs->push_back(Einsum(
+          scope, grad_inputs, absl::StrCat(output_subs, "->", input_subs)));
+      return scope.status();
+    }
+    // We do have reduced axes, so we invoke the subroutine for reduced unary
+    // einsums.
+    grad_outputs->push_back(EinsumGradReducedHelper(
+        scope, grad, output_subs, input_subs, input_shape, reduced_label_set));
+    return scope.status();
+  }
+
+  std::vector<absl::string_view> subs = absl::StrSplit(input_subs, ',');
+  if (subs.size() != 2) {
+    return absl::InvalidArgumentError("Only 2 inputs are supported");
+  }
+  std::string x_subs(subs[0]);
+  std::string y_subs(subs[1]);
+  // Add ellipsis for broadcasted dimensions if any operand does not have it.
+  // This is because the equation "...ij,jk->ik" may be valid if the 0th input's
+  // batch shape is empty, but the VJP equation "jk,ik->...ij" is not valid
+  // because only the output subscripts contain ellipsis.
+  if (absl::StrContains(output_subs, kEllipsis)) {
+    if (!absl::StrContains(x_subs, kEllipsis)) {
+      absl::StrAppend(&x_subs, kEllipsis);
+    }
+    if (!absl::StrContains(y_subs, kEllipsis)) {
+      absl::StrAppend(&y_subs, kEllipsis);
+    }
+  }
+
+  // Obtain the gradients wrt the inputs x and y, without taking into account
+  // the unbroadcasting.
+  tensorflow::Output x = op.input(0);
+  tensorflow::Output y = op.input(1);
+  if (DataTypeIsComplex(grad.type())) {
+    x = Conj(scope, x);
+    y = Conj(scope, y);
+  }
+
+  const auto x_shape = Shape(scope, x);
+  const auto y_shape = Shape(scope, y);
+  Output grad_x =
+      EinsumGradWrt(scope, grad, y, x_shape, x_subs, y_subs, output_subs);
+  Output grad_y =
+      EinsumGradWrt(scope, grad, x, y_shape, y_subs, x_subs, output_subs);
+
+  if (!absl::StrContains(output_subs, kEllipsis)) {
+    // If no ellipsis in the output; then no need to unbroadcast.
+    grad_outputs->push_back(grad_x);
+    grad_outputs->push_back(grad_y);
+    return scope.status();
+  }
+
+  // Below we handle the case that broadcasting between x and y was necessary,
+  // with x and y having possibly different batch shapes.
+
+  // Obtain the range of axes which map to ellipsis. E.g. for subscripts
+  // 'ab...c' and shape of rank 10; the range [3:-1] denotes the broadcasted
+  // axes.
+  int bx_start, by_start;
+  std::optional<int> bx_end, by_end;
+  std::tie(bx_start, bx_end) = EinsumGetBcastSubshape(x_subs);
+  std::tie(by_start, by_end) = EinsumGetBcastSubshape(y_subs);
+
+  // Sum the gradient across the broadcasted axes.
+  auto args = internal::BroadcastGradientArgs(
+      scope, Slice1dHelper(scope, x_shape, bx_start, bx_end),
+      Slice1dHelper(scope, y_shape, by_start, by_end));
+  grad_x = Reshape(
+      scope, ReduceSum(scope, grad_x, Add(scope, bx_start, args.r0)), x_shape);
+  grad_y = Reshape(
+      scope, ReduceSum(scope, grad_y, Add(scope, by_start, args.r1)), y_shape);
+  grad_outputs->push_back(grad_x);
+  grad_outputs->push_back(grad_y);
+  return scope.status();
+}
+
+REGISTER_GRADIENT_OP("Einsum", EinsumGrad);
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/linalg_grad_test.cc

@@ -0,0 +1,161 @@
+/* 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 <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using tensorflow::ops::Einsum;
+using tensorflow::ops::Placeholder;
+
+class LinalgGradTest : public ::testing::Test {
+ protected:
+  LinalgGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(LinalgGradTest, Einsum_Transpose) {
+  TensorShape x_shape({2, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Einsum(scope_, {x}, "ij->ji");
+  TensorShape y_shape({3, 2});
+  RunTest({x}, {x_shape}, {y}, {y_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_TransposeBroadcast) {
+  TensorShape x_shape({3, 2, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = Einsum(scope_, {x}, "...ij->...ji");
+  TensorShape y_shape({3, 3, 2});
+  RunTest({x}, {x_shape}, {y}, {y_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_MatMul) {
+  TensorShape x_shape({2, 3});
+  TensorShape y_shape({3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "ij,jk->ik");
+  TensorShape z_shape({2, 3});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_MatMulComplex) {
+  TensorShape x_shape({2, 3});
+  TensorShape y_shape({3, 3});
+  Output x = Placeholder(scope_, DT_COMPLEX64, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_COMPLEX64, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "ij,jk->ik");
+  TensorShape z_shape({2, 3});
+  TF_ASSERT_OK(scope_.status());
+  float max_error;
+  TF_ASSERT_OK((ComputeGradientError<complex64, complex64, float>(
+      scope_, {x, y}, {x_shape, y_shape}, {z}, {z_shape}, &max_error)));
+  EXPECT_LT(max_error, 1e-3);
+}
+
+TEST_F(LinalgGradTest, Einsum_MatMulBroadcast) {
+  TensorShape x_shape({3, 2, 3});
+  TensorShape y_shape({3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "...ij,...jk->...ik");
+  TensorShape z_shape({3, 2, 3});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_Trace) {
+  TensorShape x_shape({3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Note: In Python this could just be "ii" becuase tf.einsum normalizes the
+  // equation, but c++ doesn't do that.
+  auto z = Einsum(scope_, {x}, "ii->");
+  TensorShape z_shape({});
+  RunTest({x}, {x_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_TraceBroadcast) {
+  TensorShape x_shape({4, 3, 3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Note: In Python this could just be "ii" becuase tf.einsum normalizes the
+  // equation, but c++ doesn't do that.
+  auto z = Einsum(scope_, {x}, "...ii->...");
+  TensorShape z_shape({4});
+  RunTest({x}, {x_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_DotProduct) {
+  TensorShape x_shape({3});
+  TensorShape y_shape({3});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "i,i->");
+  TensorShape z_shape({});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_OuterProduct) {
+  TensorShape x_shape({3});
+  TensorShape y_shape({5});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "i,j->ij");
+  TensorShape z_shape({3, 5});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+TEST_F(LinalgGradTest, Einsum_TwoInputReduction) {
+  TensorShape x_shape({3, 2, 4});
+  TensorShape y_shape({4, 5});
+  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  Output y = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(y_shape));
+  auto z = Einsum(scope_, {x, y}, "abc,cd->ad");
+  TensorShape z_shape({3, 5});
+  RunTest({x, y}, {x_shape, y_shape}, {z}, {z_shape});
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/manip_grad.cc

@@ -0,0 +1,43 @@
+/* Copyright 2020 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 <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/manip_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status RollGrad(const Scope& scope, const Operation& op,
+                      const std::vector<Output>& grad_inputs,
+                      std::vector<Output>* grad_outputs) {
+  auto shift = op.input(1);
+  auto axis = op.input(2);
+  auto grad_op = Roll(scope, grad_inputs[0], Neg(scope, shift), axis);
+  grad_outputs->push_back(grad_op);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Roll", RollGrad);
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/manip_grad_test.cc

@@ -0,0 +1,53 @@
+/* Copyright 2020 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 <gtest/gtest.h>
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/manip_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Placeholder;
+using ops::Roll;
+
+class ManipGradTest : public ::testing::Test {
+ protected:
+  ManipGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-4);
+  }
+
+  Scope scope_;
+};
+
+TEST_F(ManipGradTest, RollGrad) {
+  TensorShape shape({5, 4, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Roll(scope_, x, {2, 1}, {0, 1});
+  RunTest(x, shape, y, shape);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/nn_grad.cc

@@ -0,0 +1,610 @@
+/* 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 <cstdint>
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/cc/ops/nn_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status SoftmaxGrad(const Scope& scope, const Operation& op,
+                         const std::vector<Output>& grad_inputs,
+                         std::vector<Output>* grad_outputs) {
+  // Softmax gradient function.
+  // p = softmax(x) maps from [batch, n] to [batch, m]
+  // dp/dx = [dp0/dx0   ... dp0/dxn-1  ]
+  //         [  ...           ...      ]
+  //         [dpm-1/dx0 ... dpm-1/dxn-1]
+  // dL/dx = dp/dx * dL/dy
+  //
+  // Using alternative formula:
+  // dL/dx = dL/dy * y - sum(dL/dy * y) * y
+  //    = (dL/dy - sum(dL/dy * y)) * y
+  auto y = op.output(0);
+  auto dyy = Mul(scope, grad_inputs[0], y);
+  auto sum = Sum(scope, dyy, /*axis=*/-1, Sum::KeepDims(true));
+  auto sub = Sub(scope, grad_inputs[0], sum);
+  auto dx = Mul(scope, sub, y);
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Softmax", SoftmaxGrad);
+
+bool IsZero(const Scope& scope, const Output& grad) {
+  std::string op_type_name = grad.op().node()->type_string();
+  if (op_type_name == "ZerosLike" || op_type_name == "Zeros") {
+    return true;
+  }
+  // The Operation we were provided is not named something obvious so
+  // we need to actually look at its contents.
+  // The original python code did this by calling a utility function called
+  // tensor_util.constant_value.
+  // There is no C++ equivalent to tensor_util.constant_value so we do nothing
+  // for the moment.
+  return false;
+}
+
+// Multiply after broadcasting vec to match dimensions of mat.
+//   Args:
+//     vec: A 1-D tensor of dimension [D0]
+//     mat: A 2-D tensor of dimension [D0, D1]
+//
+//   Returns:
+//     A tensor of dimension [D0, D1], the result for vec * mat.
+Output BroadcastMul(const Scope& scope, const Output& vec, const Output& mat) {
+  auto reshaped = ExpandDims(scope, vec, -1);
+  return Multiply(scope, reshaped, mat);
+}
+
+absl::Status SoftmaxCrossEntropyWithLogitsGrad(
+    const Scope& scope, const Operation& op,
+    const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
+  // Softmax gradient with cross entropy logits function.
+  // We multiply the backprop for cost with the gradients - op.output[1].
+  // There is no gradient for labels.
+
+  // The outputs of the network are at input index 0.
+  auto logits = op.input(0);
+  // The "truth" labels are at index 1.
+  auto softmax_grad = op.output(1);
+
+  // The loss is the output at index 0, and backprop is the output at index 1.
+  auto grad_loss = grad_inputs[0];
+  auto grad_grad = grad_inputs[1];
+
+  auto grad = BroadcastMul(scope, grad_loss, softmax_grad);
+  if (!IsZero(scope, grad_grad)) {
+    std::vector<int> axis;
+    auto logits_softmax = Softmax(scope, logits);
+
+    auto grad_grad_expand = ExpandDims(scope, grad_grad, 1);
+    auto logits_softmax_expand = ExpandDims(scope, logits_softmax, 2);
+    auto matmul_result =
+        BatchMatMul(scope, grad_grad_expand, logits_softmax_expand);
+    axis.push_back(1);
+    auto squeeze_result = Squeeze(scope, matmul_result, Squeeze::Axis(axis));
+    auto subtraction_result = Subtract(scope, grad_grad, squeeze_result);
+    auto multiply_result = Multiply(scope, subtraction_result, logits_softmax);
+    grad = Add(scope, grad, multiply_result);
+  }
+  auto minus_log_softmax = Multiply(scope, LogSoftmax(scope, logits), -1.0f);
+  grad_outputs->push_back(grad);
+  grad_outputs->push_back(BroadcastMul(scope, grad_loss, minus_log_softmax));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("SoftmaxCrossEntropyWithLogits",
+                     SoftmaxCrossEntropyWithLogitsGrad);
+
+absl::Status LogSoftmaxGrad(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto softmax = Exp(scope, op.output(0));
+  auto sum = Sum(scope, grad_inputs[0], {1}, Sum::KeepDims(true));
+  auto mul = Mul(scope, sum, softmax);
+  auto dx = Sub(scope, grad_inputs[0], mul);
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LogSoftmax", LogSoftmaxGrad);
+
+absl::Status ReluGradHelper(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto dx = internal::ReluGrad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Relu", ReluGradHelper);
+
+absl::Status Relu6GradHelper(const Scope& scope, const Operation& op,
+                             const std::vector<Output>& grad_inputs,
+                             std::vector<Output>* grad_outputs) {
+  auto dx = internal::Relu6Grad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Relu6", Relu6GradHelper);
+
+absl::Status LeakyReluGradHelper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  float alpha;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "alpha", &alpha));
+  internal::LeakyReluGrad::Attrs attrs;
+  auto dx = internal::LeakyReluGrad(scope, grad_inputs[0], op.input(0),
+                                    attrs.Alpha(alpha));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LeakyRelu", LeakyReluGradHelper);
+
+absl::Status LeakyReluGradGradHelper(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  float alpha;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "alpha", &alpha));
+  internal::LeakyReluGrad::Attrs attrs;
+  auto dx = internal::LeakyReluGrad(scope, grad_inputs[0], op.input(1),
+                                    attrs.Alpha(alpha));
+  grad_outputs->push_back(dx);
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LeakyReluGrad", LeakyReluGradGradHelper);
+
+absl::Status EluGradHelper(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto dx = internal::EluGrad(scope, grad_inputs[0], op.output(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Elu", EluGradHelper);
+
+absl::Status SeluGradHelper(const Scope& scope, const Operation& op,
+                            const std::vector<Output>& grad_inputs,
+                            std::vector<Output>* grad_outputs) {
+  auto dx = internal::SeluGrad(scope, grad_inputs[0], op.output(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Selu", SeluGradHelper);
+
+absl::Status L2LossGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Mul(scope, op.input(0), grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("L2Loss", L2LossGrad);
+
+absl::Status BiasAddGradHelper(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.output(0).node()->attrs(), "data_format", &data_format));
+  auto dx_1 =
+      BiasAddGrad(scope, grad_inputs[0], BiasAddGrad::DataFormat(data_format));
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  grad_outputs->push_back(dx_1);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("BiasAdd", BiasAddGradHelper);
+
+absl::Status Conv2DGrad(const Scope& scope, const Operation& op,
+                        const std::vector<Output>& grad_inputs,
+                        std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  std::string padding;
+  std::vector<int32_t> strides;
+  bool use_cudnn_on_gpu;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "use_cudnn_on_gpu", &use_cudnn_on_gpu));
+  auto dx_1 = Conv2DBackpropInput(scope, Shape(scope, op.input(0)), op.input(1),
+                                  grad_inputs[0], strides, padding,
+                                  Conv2DBackpropInput::DataFormat(data_format)
+                                      .UseCudnnOnGpu(use_cudnn_on_gpu));
+  grad_outputs->push_back(dx_1);
+  auto dx_2 =
+      Conv2DBackpropFilter(scope, op.input(0), Shape(scope, op.input(1)),
+                           grad_inputs[0], strides, padding,
+                           Conv2DBackpropFilter::DataFormat(data_format)
+                               .UseCudnnOnGpu(use_cudnn_on_gpu));
+  grad_outputs->push_back(dx_2);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Conv2D", Conv2DGrad);
+
+absl::Status MaxPoolGradHelper(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  std::string padding;
+  std::vector<int32_t> strides;
+  std::vector<int32_t> ksize;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  auto dx = internal::MaxPoolGrad(
+      scope, op.input(0), op.output(0), grad_inputs[0], ksize, strides, padding,
+      internal::MaxPoolGrad::DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPool", MaxPoolGradHelper);
+
+absl::Status MaxPoolGradV2Helper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  std::string data_format;
+  std::string padding;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  auto dx = MaxPoolGradV2(scope, op.input(0), op.output(0), grad_inputs[0],
+                          op.input(1), op.input(2), padding,
+                          MaxPoolGradV2::DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPoolV2", MaxPoolGradV2Helper);
+
+absl::Status MaxPool3DGradHelper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  std::vector<int32_t> ksize;
+  std::vector<int32_t> strides;
+  std::string padding;
+  std::string data_format;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  MaxPool3DGrad::Attrs grad_attrs;
+  auto dx =
+      MaxPool3DGrad(scope, op.input(0), op.output(0), grad_inputs[0], ksize,
+                    strides, padding, grad_attrs.DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPool3D", MaxPool3DGradHelper);
+
+absl::Status AvgPoolGradHelper(const Scope& scope, const Operation& op,
+                               const std::vector<Output>& grad_inputs,
+                               std::vector<Output>* grad_outputs) {
+  std::vector<int32_t> ksize;
+  std::vector<int32_t> strides;
+  std::string padding;
+  std::string data_format;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  internal::AvgPoolGrad::Attrs grad_attrs;
+  auto dx = internal::AvgPoolGrad(scope, Shape(scope, op.input(0)),
+                                  grad_inputs[0], ksize, strides, padding,
+                                  grad_attrs.DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("AvgPool", AvgPoolGradHelper);
+
+absl::Status AvgPool3DGradHelper(const Scope& scope, const Operation& op,
+                                 const std::vector<Output>& grad_inputs,
+                                 std::vector<Output>* grad_outputs) {
+  std::vector<int32_t> ksize;
+  std::vector<int32_t> strides;
+  std::string padding;
+  std::string data_format;
+  auto attrs = op.output(0).node()->attrs();
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+  TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+  AvgPool3DGrad::Attrs grad_attrs;
+  auto dx =
+      AvgPool3DGrad(scope, Shape(scope, op.input(0)), grad_inputs[0], ksize,
+                    strides, padding, grad_attrs.DataFormat(data_format));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("AvgPool3D", AvgPool3DGradHelper);
+
+absl::Status LRNGradHelper(const Scope& scope, const Operation& op,
+                           const std::vector<Output>& grad_inputs,
+                           std::vector<Output>* grad_outputs) {
+  auto dx = internal::LRNGrad(scope, grad_inputs[0], op.input(0), op.output(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("LRN", LRNGradHelper);
+
+absl::Status SoftplusGradHelper(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  auto dx = internal::SoftplusGrad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Softplus", SoftplusGradHelper);
+
+absl::Status SoftsignGradHelper(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  auto dx = internal::SoftsignGrad(scope, grad_inputs[0], op.input(0));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Softsign", SoftsignGradHelper);
+
+absl::Status FractionalAvgPoolGradHelper(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  bool overlapping;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.output(0).node()->attrs(), "overlapping", &overlapping));
+  auto dx = internal::FractionalAvgPoolGrad(
+      scope, Shape(scope, op.input(0), Shape::OutType(DT_INT64)),
+      grad_inputs[0], op.output(1), op.output(2),
+      internal::FractionalAvgPoolGrad::Overlapping(overlapping));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("FractionalAvgPool", FractionalAvgPoolGradHelper);
+
+absl::Status FractionalMaxPoolGradHelper(const Scope& scope,
+                                         const Operation& op,
+                                         const std::vector<Output>& grad_inputs,
+                                         std::vector<Output>* grad_outputs) {
+  bool overlapping;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.output(0).node()->attrs(), "overlapping", &overlapping));
+  auto dx = internal::FractionalMaxPoolGrad(
+      scope, op.input(0), op.output(0), grad_inputs[0], op.output(1),
+      op.output(2), internal::FractionalMaxPoolGrad::Overlapping(overlapping));
+  grad_outputs->push_back(dx);
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("FractionalMaxPool", FractionalMaxPoolGradHelper);
+
+// Templated constructor for FusedBatchNormGrad[..]::Attrs.
+template <typename T>
+T FusedBatchNormGradAttrs(float epsilon, absl::string_view data_format,
+                          bool is_training) {
+  T result;
+  result.epsilon_ = epsilon;
+  result.data_format_ = data_format;
+  result.is_training_ = is_training;
+  return result;
+}
+
+using BatchNormGradFn = std::function<absl::Status(
+    const Scope&, Output x, Output grad_y, Output scale,
+    const std::vector<Output>& reserve_spaces, float epsilon,
+    absl::string_view data_format, bool is_training,
+    std::vector<Output>* grad_outputs)>;
+
+absl::Status BaseFusedBatchNormGrad(const Scope& scope, const Operation& op,
+                                    const std::vector<Output>& grad_inputs,
+                                    BatchNormGradFn grad_fn,
+                                    std::vector<Output>* grad_outputs) {
+  if (op.num_outputs() < 5) {
+    return absl::InvalidArgumentError(
+        "FusedBatchNorm requires at least 5 outputs");
+  }
+  if (grad_inputs.empty()) {
+    return absl::InvalidArgumentError(
+        "FusedBatchNorm grad requires 1 grad input");
+  }
+  if (op.num_inputs() < 3) {
+    return absl::InvalidArgumentError("FusedBatchNorm has too few inputs");
+  }
+
+  Output x = op.input(0);
+  Output grad_y = grad_inputs[0];
+  Output scale = op.input(1);
+  float epsilon;
+  std::string data_format;
+  bool is_training;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "epsilon", &epsilon));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "is_training", &is_training));
+
+  std::vector<Output> reserve_spaces;
+  reserve_spaces.push_back(op.output(3));
+  reserve_spaces.push_back(op.output(4));
+  if (op.num_outputs() > 5) {
+    reserve_spaces.push_back(op.output(5));
+  }
+
+  if (is_training) {
+    return grad_fn(scope, x, grad_y, scale, reserve_spaces, epsilon,
+                   data_format, is_training, grad_outputs);
+  } else {
+    if (op.num_inputs() < 5) {
+      return absl::InvalidArgumentError(
+          "FusedBatchNorm requires 5 inputs in eval mode");
+    }
+
+    reserve_spaces[0] = op.input(3);  // pop_mean
+    reserve_spaces[1] = op.input(4);  // pop_var
+    if (data_format == "NCHW") {
+      x = Transpose(scope, x, {0, 2, 3, 1});
+      grad_y = Transpose(scope, grad_y, {0, 2, 3, 1});
+    } else if (data_format == "NCDHW") {
+      x = Transpose(scope, x, {0, 2, 3, 4, 1});
+      grad_y = Transpose(scope, grad_y, {0, 2, 3, 4, 1});
+    }
+
+    absl::string_view target_data_format;
+    if (data_format == "NCHW" || data_format == "NHWC") {
+      target_data_format = "NHWC";
+    } else {
+      target_data_format = "NDHWC";
+    }
+
+    TF_RETURN_IF_ERROR(grad_fn(scope, x, grad_y, scale, reserve_spaces, epsilon,
+                               target_data_format, is_training, grad_outputs));
+    if (data_format == "NCHW") {
+      (*grad_outputs)[0] = Transpose(scope, (*grad_outputs)[0], {0, 3, 1, 2});
+    } else if (data_format == "NCDHW") {
+      (*grad_outputs)[0] =
+          Transpose(scope, (*grad_outputs)[0], {0, 4, 1, 2, 3});
+    }
+    return scope.status();
+  }
+}
+
+absl::Status FusedBatchNormV3Grad(const Scope& scope, const Operation& op,
+                                  const std::vector<Output>& grad_inputs,
+                                  std::vector<Output>* grad_outputs) {
+  return BaseFusedBatchNormGrad(
+      scope, op, grad_inputs,
+      [](const Scope& scope, Output x, Output grad_y, Output scale,
+         const std::vector<Output>& reserve_spaces, float epsilon,
+         absl::string_view data_format, bool is_training,
+         std::vector<Output>* grad_outputs) {
+        FusedBatchNormGradV3 grad(
+            scope, grad_y, x, scale, reserve_spaces[0], reserve_spaces[1],
+            reserve_spaces[2],
+            FusedBatchNormGradAttrs<FusedBatchNormGradV3::Attrs>(
+                epsilon, data_format, is_training));
+        grad_outputs->push_back(grad.x_backprop);
+        grad_outputs->push_back(grad.scale_backprop);
+        grad_outputs->push_back(grad.offset_backprop);
+        grad_outputs->push_back(NoGradient());
+        grad_outputs->push_back(NoGradient());
+        return scope.status();
+      },
+      grad_outputs);
+}
+
+REGISTER_GRADIENT_OP("FusedBatchNormV3", FusedBatchNormV3Grad);
+
+absl::Status Conv2DBackpropInputGrad(const Scope& scope, const Operation& op,
+                                     const std::vector<Output>& grad_inputs,
+                                     std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 3) {
+    return absl::InvalidArgumentError("Conv2DBackpropInput requires 3 inputs.");
+  }
+  if (grad_inputs.empty()) {
+    return absl::InvalidArgumentError(
+        "Conv2DBackpropInput grad requires 1 grad input");
+  }
+
+  std::vector<int> dilations, strides, explicit_paddings;
+  bool use_cudnn_on_gpu;
+  std::string data_format, padding;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "dilations", &dilations));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "strides", &strides));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "explicit_paddings", &explicit_paddings));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "use_cudnn_on_gpu", &use_cudnn_on_gpu));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "padding", &padding));
+
+  grad_outputs->push_back(NoGradient());
+
+  Conv2DBackpropFilter::Attrs filter_attrs;
+  filter_attrs.use_cudnn_on_gpu_ = use_cudnn_on_gpu;
+  filter_attrs.explicit_paddings_ = explicit_paddings;
+  filter_attrs.data_format_ = data_format;
+  filter_attrs.dilations_ = dilations;
+  grad_outputs->push_back(
+      Conv2DBackpropFilter(scope, grad_inputs[0], Shape(scope, op.input(1)),
+                           op.input(2), strides, padding, filter_attrs));
+
+  Conv2D::Attrs conv_attrs;
+  conv_attrs.use_cudnn_on_gpu_ = use_cudnn_on_gpu;
+  conv_attrs.explicit_paddings_ = explicit_paddings;
+  conv_attrs.data_format_ = data_format;
+  conv_attrs.dilations_ = dilations;
+  grad_outputs->push_back(
+      Conv2D(scope, grad_inputs[0], op.input(1), strides, padding, conv_attrs));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("Conv2DBackpropInput", Conv2DBackpropInputGrad);
+
+absl::Status DepthwiseConv2dNativeGrad(const Scope& scope, const Operation& op,
+                                       const std::vector<Output>& grad_inputs,
+                                       std::vector<Output>* grad_outputs) {
+  if (op.num_inputs() != 2) {
+    return absl::InvalidArgumentError(
+        "DepthwiseConv2dNative requires 2 inputs.");
+  }
+  if (grad_inputs.empty()) {
+    return absl::InvalidArgumentError(
+        "DepthwiseConv2dNative grad requires 1 grad input");
+  }
+
+  std::vector<int> dilations, strides, explicit_paddings;
+  std::string data_format, padding;
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "dilations", &dilations));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "strides", &strides));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "explicit_paddings", &explicit_paddings));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "padding", &padding));
+
+  DepthwiseConv2dNativeBackpropInput::Attrs input_attrs;
+  input_attrs.explicit_paddings_ = explicit_paddings;
+  input_attrs.data_format_ = data_format;
+  input_attrs.dilations_ = dilations;
+  grad_outputs->push_back(DepthwiseConv2dNativeBackpropInput(
+      scope, Shape(scope, op.input(0)), op.input(1), grad_inputs[0], strides,
+      padding, input_attrs));
+
+  DepthwiseConv2dNativeBackpropFilter::Attrs filter_attrs;
+  filter_attrs.explicit_paddings_ = explicit_paddings;
+  filter_attrs.data_format_ = data_format;
+  filter_attrs.dilations_ = dilations;
+  grad_outputs->push_back(DepthwiseConv2dNativeBackpropFilter(
+      scope, op.input(0), Shape(scope, op.input(1)), grad_inputs[0], strides,
+      padding, filter_attrs));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("DepthwiseConv2dNative", DepthwiseConv2dNativeGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/nn_grad_test.cc

@@ -0,0 +1,416 @@
+/* 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 <cstddef>
+#include <tuple>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/nn_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/random/random.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::AvgPool;
+using ops::AvgPool3D;
+using ops::BiasAdd;
+using ops::Conv2D;
+using ops::Conv2DBackpropInput;
+using ops::DepthwiseConv2dNative;
+using ops::Elu;
+using ops::FractionalAvgPool;
+using ops::FractionalMaxPool;
+using ops::FusedBatchNormV3;
+using ops::L2Loss;
+using ops::LogSoftmax;
+using ops::LRN;
+using ops::MaxPool;
+using ops::MaxPool3D;
+using ops::MaxPoolV2;
+using ops::Placeholder;
+using ops::Relu;
+using ops::Relu6;
+using ops::Selu;
+using ops::Softmax;
+using ops::Softplus;
+using ops::Softsign;
+
+class NNGradTest : public ::testing::Test {
+ protected:
+  NNGradTest() : scope_(Scope::NewRootScope()) {}
+
+  void RunTest(const Output& x, const TensorShape& x_shape, const Output& y,
+               const TensorShape& y_shape) {
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const Output& x, const Tensor& x_init_value, const Output& y,
+               const TensorShape& y_shape) {
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, x, x_init_value, y, y_shape, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  void RunTest(const OutputList& xs, const std::vector<TensorShape>& x_shapes,
+               const OutputList& ys, const std::vector<TensorShape>& y_shapes) {
+    TF_ASSERT_OK(scope_.status());
+    float max_error;
+    TF_ASSERT_OK((ComputeGradientError<float, float, float>(
+        scope_, xs, x_shapes, ys, y_shapes, &max_error)));
+    EXPECT_LT(max_error, 1e-3);
+  }
+
+  // Sets tensor with random values, ensuring that every pair of elements are at
+  // least a reasonable amount apart.
+  // This is an issue for max pooling operations, in which perturbations by the
+  // numeric gradient computation in the gradient checker can change the max
+  // value if a pool has values that are too close together.
+  template <typename T>
+  void SetRandomValuesForMaxPooling(Tensor* tensor) {
+    auto tensor_flat = tensor->flat<T>();
+    // First set the array to an increasing sequence of values spaced
+    // a reasonable amount apart
+    T cur = 0;
+    for (size_t i = 0; i < tensor->NumElements(); i++) {
+      tensor_flat(i) = cur;
+      cur += 5e-2;
+    }
+    // Fischer-Yates shuffle the array
+    for (size_t i = tensor->NumElements() - 1; i >= 1; i--) {
+      // j <- random integer 0 <= j <= i
+      size_t j = random::New64() % (i + 1);
+      // swap values at i, j
+      T tmp = tensor_flat(i);
+      tensor_flat(i) = tensor_flat(j);
+      tensor_flat(j) = tmp;
+    }
+  }
+
+  Scope scope_;
+};
+
+TEST_F(NNGradTest, SoftmaxGrad) {
+  TensorShape shape({32, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softmax(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, SoftmaxRank3Grad) {
+  TensorShape shape({32, 1, 10});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softmax(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, SoftmaxCrossEntropyWithLogitsGrad) {
+  TensorShape logits_shape({5, 3});
+  TensorShape loss_shape({5});
+
+  auto logits = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(logits_shape));
+  auto labels = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(logits_shape));
+  auto y =
+      tensorflow::ops::SoftmaxCrossEntropyWithLogits(scope_, logits, labels);
+  // Note the reversal of the backprop and loss orders. Issue #18734 has been
+  // opened for this.
+  RunTest({logits, labels}, {logits_shape, logits_shape}, {y.backprop, y.loss},
+          {logits_shape, loss_shape});
+}
+
+TEST_F(NNGradTest, LogSoftmaxGrad) {
+  TensorShape shape({5, 3});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = LogSoftmax(scope_, x);
+  // Avoid numerical instability when computing finite differences.
+  Tensor x_init_value =
+      test::AsTensor<float>({-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f,
+                             0.5f, 0.7f, 0.8f, -0.1f, 0.1f, 0.1f, 0.1f, 1.2f},
+                            {5, 3});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, ReluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Relu(scope_, x);
+  // Avoid input values where ReLU gradient is not well defined (around zero).
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, Relu6Grad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Relu6(scope_, x);
+  // Avoid input values where ReLU gradient is not well defined (around zero
+  // and six).
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 6.1f, 6.3f, 6.5f, 6.7f, 6.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, LeakyReluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = ops::internal::LeakyRelu(scope_, x);
+  // Avoid input values where Leaky ReLU gradient is not well defined (around
+  // zero).
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, LeakyReluGradGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  // Avoid input values where Leaky ReLU gradient is not well defined (around
+  // zero).
+  Tensor x_init_value = test::AsTensor<float>(
+      {2.3f, 1.9f, 1.5f, 1.1f, 0.7f, 0.3f, -0.1f, -0.5f, -0.9f, -1.3f}, {5, 2});
+  Tensor features = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  auto y = ops::internal::LeakyReluGrad(scope_, x, features);
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, EluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Elu(scope_, x);
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, SeluGrad) {
+  TensorShape shape({5, 2});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Selu(scope_, x);
+  Tensor x_init_value = test::AsTensor<float>(
+      {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f},
+      {5, 2});
+  RunTest(x, x_init_value, y, shape);
+}
+
+TEST_F(NNGradTest, L2LossGrad) {
+  TensorShape x_shape({5, 2});
+  TensorShape y_shape({1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = L2Loss(scope_, x);
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, BiasAddGradHelper) {
+  TensorShape shape({4, 5});
+  TensorShape bias_shape({5});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto bias = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(bias_shape));
+  auto y = BiasAdd(scope_, x, bias);
+  RunTest({x, bias}, {shape, bias_shape}, {y}, {shape});
+}
+
+TEST_F(NNGradTest, Conv2DGrad) {
+  TensorShape shape({1, 2, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  Tensor filter = test::AsTensor<float>({0.5f}, {1, 1, 1, 1});
+  const std::vector<int> strides{1, 1, 1, 1};
+  auto y = Conv2D(scope_, x, filter, strides, "SAME");
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, MaxPoolGradHelper) {
+  TensorShape x_shape({1, 2, 2, 1});
+  TensorShape y_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one MaxPool.
+  const std::vector<int> ksize{1, 2, 2, 1};
+  const std::vector<int> strides{1, 2, 2, 1};
+  auto y = MaxPool(scope_, x, ksize, strides, "VALID");
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, MaxPoolGradV2Helper) {
+  TensorShape x_shape({1, 2, 2, 1});
+  TensorShape y_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one MaxPool.
+  Tensor ksize = test::AsTensor<int>({1, 2, 2, 1}, {4});
+  Tensor strides = test::AsTensor<int>({1, 2, 2, 1}, {4});
+  auto y = MaxPoolV2(scope_, x, ksize, strides, "VALID");
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, MaxPool3DGradHelper) {
+  TensorShape x_shape({1, 3, 3, 3, 1});
+  TensorShape y_shape({1, 1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one MaxPool3D.
+  const std::vector<int> ksize{1, 3, 3, 3, 1};
+  const std::vector<int> strides{1, 3, 3, 3, 1};
+  auto y = MaxPool3D(scope_, x, ksize, strides, "VALID");
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, AvgPoolGradHelper) {
+  TensorShape x_shape({1, 2, 2, 1});
+  TensorShape y_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one AvgPool.
+  const std::vector<int> ksize{1, 2, 2, 1};
+  const std::vector<int> strides{1, 2, 2, 1};
+  auto y = AvgPool(scope_, x, ksize, strides, "SAME");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, AvgPool3DGradHelper) {
+  TensorShape x_shape({1, 3, 3, 3, 1});
+  TensorShape y_shape({1, 1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Setup window and strides so that we only do one AvgPool3D.
+  const std::vector<int> ksize{1, 3, 3, 3, 1};
+  const std::vector<int> strides{1, 3, 3, 3, 1};
+  auto y = AvgPool3D(scope_, x, ksize, strides, "SAME");
+  RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, LRN) {
+  TensorShape x_shape({1, 1, 2, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  auto y = LRN(scope_, x);
+  RunTest(x, x_shape, y, x_shape);
+}
+
+TEST_F(NNGradTest, SoftplusGrad) {
+  TensorShape shape({3, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softplus(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, SoftsignGrad) {
+  TensorShape shape({3, 7});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto y = Softsign(scope_, x);
+  RunTest(x, shape, y, shape);
+}
+
+TEST_F(NNGradTest, FractionalAvgPoolGradHelper) {
+  TensorShape x_shape({1, 3, 7, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Force consistent pooling regions for unit testing.
+  auto y = FractionalAvgPool(
+      scope_, x, {1, 1.2, 1.9, 1},
+      FractionalAvgPool::Deterministic(true).Overlapping(true).Seed(1).Seed2(
+          2));
+  TensorShape y_shape({1, 2, 3, 1});
+  RunTest(x, x_shape, y.output, y_shape);
+}
+
+TEST_F(NNGradTest, FractionalMaxPoolGradHelper) {
+  TensorShape x_shape({1, 3, 7, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+  // Force consistent pooling regions for unit testing.
+  auto y = FractionalMaxPool(
+      scope_, x, {1, 1.2, 1.9, 1},
+      FractionalMaxPool::Deterministic(true).Overlapping(true).Seed(1).Seed2(
+          2));
+  Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+  SetRandomValuesForMaxPooling<float>(&x_init_value);
+  TensorShape y_shape({1, 2, 3, 1});
+  RunTest(x, x_init_value, y.output, y_shape);
+}
+
+class FusedBatchNormGradTest : public NNGradTest,
+                               public ::testing::WithParamInterface<
+                                   std::tuple<bool, bool, TensorShape>> {};
+
+TEST_P(FusedBatchNormGradTest, FusedBatchNormV3Grad) {
+  FusedBatchNormV3::Attrs attrs;
+  attrs.is_training_ = std::get<0>(GetParam());
+  bool channel_first = std::get<1>(GetParam());
+  TensorShape shape = std::get<2>(GetParam());
+  int channel_dim = (channel_first) ? 1 : shape.dims() - 1;
+  TensorShape scale_shape({shape.dim_size(channel_dim)});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto scale = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
+  auto offset = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
+  auto mean = ops::ZerosLike(scope_, scale);
+  auto var = ops::OnesLike(scope_, scale);
+
+  if (!channel_first) {
+    attrs.data_format_ = (shape.dims() == 5) ? "NDHWC" : "NHWC";
+  } else {
+    attrs.data_format_ = (shape.dims() == 5) ? "NCDHW" : "NCHW";
+  }
+
+  auto y = FusedBatchNormV3(scope_, x, scale, offset, mean, var, attrs);
+  RunTest({x, scale, offset}, {shape, scale_shape, scale_shape}, {y.y},
+          {shape});
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    FusedBatchNormGrad, FusedBatchNormGradTest,
+    ::testing::Combine(::testing::Bool(), ::testing::Bool(),
+                       ::testing::Values(TensorShape({2, 3, 4, 5}),
+                                         TensorShape({2, 3, 2, 2, 2}))));
+
+TEST_F(NNGradTest, Conv2DBackpropInputGrad) {
+  TensorShape shape({1, 2, 2, 1});
+  TensorShape filter_shape({1, 1, 1, 1});
+  auto out = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto filter = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(filter_shape));
+  const std::vector<int> strides{1, 1, 1, 1};
+  auto y = Conv2DBackpropInput(scope_, ops::Shape(scope_, out), filter, out,
+                               strides, "SAME");
+  RunTest({out, filter}, {shape, filter_shape}, {y}, {shape});
+}
+
+TEST_F(NNGradTest, DepthwiseConv2dNativeGrad) {
+  TensorShape shape({1, 2, 2, 1});
+  TensorShape filter_shape({1, 1, 1, 1});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
+  auto filter = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(filter_shape));
+  const std::vector<int> strides{1, 1, 1, 1};
+  auto y = DepthwiseConv2dNative(scope_, x, filter, strides, "SAME");
+  RunTest({x, filter}, {shape, filter_shape}, {y}, {shape});
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/gradients/resource_variable_grad.cc

@@ -0,0 +1,37 @@
+/* Copyright 2020 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 <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/ops/array_ops.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+absl::Status ReadVariableOpGrad(const Scope& scope, const Operation& op,
+                                const std::vector<Output>& grad_inputs,
+                                std::vector<Output>* grad_outputs) {
+  grad_outputs->push_back(Identity(scope, grad_inputs[0]));
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("ReadVariableOp", ReadVariableOpGrad);
+
+}  // anonymous namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/gradients/resource_variable_grad_test.cc

@@ -0,0 +1,68 @@
+/* Copyright 2020 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 <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/framework/grad_op_registry.h"
+#include "tensorflow/cc/framework/gradient_checker.h"
+#include "tensorflow/cc/framework/gradients.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/testutil.h"
+#include "tensorflow/cc/gradients/grad_testutil.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/resource_variable_ops.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace tensorflow {
+namespace ops {
+namespace {
+
+TEST(ResourceVariableGradTest, ReadVariableOpGrad) {
+  TensorShape shape({});
+  auto scope = Scope::NewRootScope();
+  auto x = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+
+  auto var = VarHandleOp(scope, DT_FLOAT, shape);
+  auto init = AssignVariableOp(scope, var, Const(scope, 2.0f, shape));
+
+  auto temp = ReadVariableOp(scope, var, DT_FLOAT);
+
+  auto y = Mul(scope, temp, x);
+
+  auto dy = Placeholder(scope, DT_FLOAT, Placeholder::Shape(shape));
+
+  OutputList dxs;
+  TF_ASSERT_OK(AddSymbolicGradients(scope, {y}, {var}, {dy}, &dxs));
+
+  ClientSession::FeedType feed_list;
+  feed_list.insert({x, 5.0f});
+  feed_list.insert({dy, 1.0f});
+
+  std::vector<Tensor> dxout;
+  ClientSession session(scope);
+  TF_ASSERT_OK(session.Run(feed_list, dxs, &dxout));
+
+  auto grad = dxout[0].scalar<float>()();
+  EXPECT_EQ(grad, 5.0f);
+}
+
+}  // namespace
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/ops/const_op.cc

@@ -0,0 +1,86 @@
+/* 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/ops/const_op.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+namespace ops {
+
+namespace {
+template <typename T>
+Output ConstHelper(const Scope& scope, const T& value, DataType dtype) {
+  if (!scope.ok()) return Output();
+
+  Node* ret;
+  Graph* graph = scope.graph();
+  const std::string unique_name = scope.GetUniqueNameForOp("Const");
+  auto builder = NodeBuilder(unique_name, "Const")
+                     .Attr("value", value)
+                     .Attr("dtype", dtype);
+  scope.UpdateBuilder(&builder);
+  scope.UpdateStatus(builder.Finalize(graph, &ret));
+  if (!scope.ok()) return Output();
+
+  scope.UpdateStatus(scope.DoShapeInference(ret));
+  if (!scope.ok()) return Output();
+
+  return Output(ret);
+}
+}  // namespace
+
+Output Const(const Scope& scope, const Input::Initializer& val) {
+  if (!val.status.ok()) {
+    scope.UpdateStatus(val.status);
+    return Output();
+  }
+  return ConstHelper(scope, val.tensor, val.tensor.dtype());
+}
+
+Output ConstFromProto(const Scope& scope, const TensorProto& proto) {
+  return ConstHelper(scope, proto, proto.dtype());
+}
+
+NodeBuilder::NodeOut AsNodeOut(const Scope& scope, const Input& inp) {
+  if (!inp.status().ok()) {
+    scope.UpdateStatus(inp.status());
+    return NodeBuilder::NodeOut(inp.node(), inp.index());
+  }
+  if (inp.node()) {
+    return NodeBuilder::NodeOut(inp.node(), inp.index());
+  }
+  if (!inp.node_name().empty()) {
+    return NodeBuilder::NodeOut(inp.node_name(), inp.index(), inp.data_type());
+  }
+  auto transformed = Input{
+      Const(scope.NewSubScope("Const"), Input::Initializer(inp.tensor()))};
+  return NodeBuilder::NodeOut{transformed.node(), transformed.index()};
+}
+
+std::vector<NodeBuilder::NodeOut> AsNodeOutList(const Scope& scope,
+                                                const InputList& inp) {
+  std::vector<NodeBuilder::NodeOut> out;
+  for (const auto& i : inp) {
+    const auto node_out = AsNodeOut(scope, i);
+    if (!scope.ok()) {
+      return {};
+    }
+    out.push_back(node_out);
+  }
+  return out;
+}
+
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/ops/const_op.h

@@ -0,0 +1,87 @@
+/* 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_OPS_CONST_OP_H_
+#define TENSORFLOW_CC_OPS_CONST_OP_H_
+
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/graph/node_builder.h"
+
+namespace tensorflow {
+namespace ops {
+
+/// @defgroup const_op Const Op
+/// @{
+
+Output Const(const Scope& scope, const Input::Initializer& val);
+
+Output ConstFromProto(const Scope& scope, const TensorProto& proto);
+
+NodeBuilder::NodeOut AsNodeOut(const Scope& scope, const Input& inp);
+
+template <typename T>
+Output Const(const Scope& scope, const Input::Initializer& val) {
+  auto orig_const_output = Const(scope, val);
+  if (!scope.ok()) return Output();
+
+  typedef typename Input::Initializer::RealType<T>::type DstT;
+
+  if (val.tensor.dtype() == DataTypeToEnum<DstT>::v()) {
+    return orig_const_output;
+  }
+  if (val.tensor.NumElements() == 0) {
+    Tensor t(DataTypeToEnum<DstT>::v(), val.tensor.shape());
+    return Const(scope, Input::Initializer(t));
+  }
+
+  // TODO(keveman): Refactor Cast op's kernel implementation such that the code
+  // can be directly called here instead of adding the Cast op to the graph.
+  auto orig_const = AsNodeOut(scope, orig_const_output);
+  const auto cast_op_name = scope.GetUniqueNameForOp("Cast");
+
+  auto cast_builder = NodeBuilder(cast_op_name, "Cast")
+                          .Input(orig_const)
+                          .Attr("DstT", DataTypeToEnum<DstT>::v());
+  scope.UpdateBuilder(&cast_builder);
+  Node* ret;
+  scope.UpdateStatus(cast_builder.Finalize(scope.graph(), &ret));
+  if (!scope.ok()) return Output();
+  scope.UpdateStatus(scope.DoShapeInference(ret));
+  return Output(ret, 0);
+}
+
+template <typename T>
+Output Const(const Scope& scope, const T& v, const TensorShape shape) {
+  return Const(scope, Input::Initializer(v, shape));
+}
+
+template <typename T>
+Output Const(const Scope& scope, const std::initializer_list<T>& v,
+             const TensorShape shape) {
+  return Const(scope, Input::Initializer(v, shape));
+}
+
+std::vector<NodeBuilder::NodeOut> AsNodeOutList(const Scope& scope,
+                                                const InputList& inp);
+
+/// }@
+
+}  // namespace ops
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_OPS_CONST_OP_H_

tensorflow/cc/ops/const_op_test.cc

@@ -0,0 +1,151 @@
+/* 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/ops/const_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"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+template <typename T>
+void ExpectNodeEqual(const Node* n, gtl::ArraySlice<T> values,
+                     TensorShape shape) {
+  EXPECT_TRUE(n->IsConstant());
+  Tensor tensor;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "value", &tensor));
+  DataType dtype;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "dtype", &dtype));
+  EXPECT_EQ(tensor.dtype(), dtype);
+  test::ExpectTensorEqual<T>(tensor, test::AsTensor(values, shape));
+}
+
+void ExpectTypeAndShape(const Node* n, DataType expected_dtype,
+                        TensorShape expected_shape) {
+  EXPECT_TRUE(n->IsConstant());
+  Tensor tensor;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "value", &tensor));
+  DataType dtype;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "dtype", &dtype));
+  EXPECT_EQ(dtype, expected_dtype);
+  EXPECT_EQ(expected_shape, TensorShape(tensor.shape()));
+}
+
+}  // namespace
+
+TEST(ConstOpTest, Basic) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, 42.0f);
+  TF_EXPECT_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_FLOAT);
+  ExpectNodeEqual<float>(c.node(), {42.0f}, {});
+}
+
+TEST(ConstOpTest, MultiDim) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, {{2.0}, {3.0}});
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_DOUBLE);
+  ExpectNodeEqual<double>(c.node(), {2.0, 3.0}, {2, 1});
+}
+
+TEST(ConstOpTest, Empty) {
+  Scope root = Scope::NewRootScope();
+
+  auto c1 = ops::Const(root, {});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c1.node(), DT_FLOAT, {0});
+
+  auto c2 = ops::Const(root, {{}});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c2.node(), DT_FLOAT, {1, 0});
+
+  auto c3 = ops::Const(root, {{{}, {}}});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c3.node(), DT_FLOAT, {1, 2, 0});
+
+  auto c4 = ops::Const<int>(root, {{{}}});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c4.node(), DT_INT32, {1, 1, 0});
+
+  ops::Const(root, {{}, {{}}});
+  EXPECT_FALSE(root.status().ok());
+}
+
+TEST(ConstOpTest, WithExplicitShape) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, 42.0, {2, 2});
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_DOUBLE);
+  ExpectNodeEqual<double>(c.node(), {42.0, 42.0, 42.0, 42.0}, {2, 2});
+
+  auto d = ops::Const(root, {"1", "2", "3", "4", "5", "6"}, {2, 3});
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(d.op().output_type(0), DT_STRING);
+  ExpectNodeEqual<tstring>(d.node(), {"1", "2", "3", "4", "5", "6"}, {2, 3});
+}
+
+TEST(ConstOpTest, FromProto) {
+  Scope root = Scope::NewRootScope();
+  TensorProto proto;
+  proto.set_dtype(DT_DOUBLE);
+  TensorShape({2, 2}).AsProto(proto.mutable_tensor_shape());
+  for (int i = 0; i < 4; ++i) {
+    proto.add_double_val(static_cast<double>(i));
+  }
+  auto c = ops::ConstFromProto(root, proto);
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_DOUBLE);
+  ExpectNodeEqual<double>(c.node(), {0.0, 1.0, 2.0, 3.0}, {2, 2});
+}
+
+TEST(ConstOpTest, InvalidInitializer) {
+  Scope root = Scope::NewRootScope();
+  ops::Const(root, {{2.0}, {"df"}});
+  EXPECT_FALSE(root.status().ok());
+}
+
+TEST(ConstOpTest, Names) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, {{2.0}, {3.0}});
+  EXPECT_EQ(c.node()->name(), "Const");
+  auto c_1 = ops::Const(root, {{2.0}, {3.0}});
+  EXPECT_EQ(c_1.node()->name(), "Const_1");
+
+  auto x = ops::Const(root.WithOpName("x"), 1);
+  EXPECT_EQ(x.node()->name(), "x");
+  auto x_1 = ops::Const(root.WithOpName("x"), 1);
+  EXPECT_EQ(x_1.node()->name(), "x_1");
+
+  Scope child = root.NewSubScope("c");
+  auto c_y = ops::Const(child.WithOpName("y"), 1);
+  EXPECT_EQ(c_y.node()->name(), "c/y");
+  auto c_y_1 = ops::Const(child.WithOpName("y"), 1);
+  EXPECT_EQ(c_y_1.node()->name(), "c/y_1");
+}
+
+TEST(ConstOpTest, TemplatedConst) {
+  Scope root = Scope::NewRootScope();
+  auto c1 = ops::Const<int>(root, {1, 2});
+  ExpectTypeAndShape(c1.node(), DT_INT32, {2});
+
+  auto c2 = ops::Const<tstring>(root, {{"this"}, {"is"}, {"a"}, {"constant"}});
+  ExpectTypeAndShape(c2.node(), DT_STRING, {4, 1});
+}
+
+}  // namespace tensorflow

tensorflow/cc/ops/standard_ops.h

@@ -0,0 +1,40 @@
+/* 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_OPS_STANDARD_OPS_H_
+#define TENSORFLOW_CC_OPS_STANDARD_OPS_H_
+
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/candidate_sampling_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/control_flow_ops.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/image_ops.h"
+#include "tensorflow/cc/ops/io_ops.h"
+#include "tensorflow/cc/ops/linalg_ops.h"
+#include "tensorflow/cc/ops/logging_ops.h"
+#include "tensorflow/cc/ops/lookup_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/cc/ops/no_op.h"
+#include "tensorflow/cc/ops/parsing_ops.h"
+#include "tensorflow/cc/ops/random_ops.h"
+#include "tensorflow/cc/ops/sparse_ops.h"
+#include "tensorflow/cc/ops/state_ops.h"
+#include "tensorflow/cc/ops/string_ops.h"
+#include "tensorflow/cc/ops/training_ops.h"
+#include "tensorflow/cc/ops/user_ops.h"
+
+#endif  // TENSORFLOW_CC_OPS_STANDARD_OPS_H_

tensorflow/cc/ops/while_loop.cc

@@ -0,0 +1,251 @@
+/* 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/ops/while_loop.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/core/common_runtime/shape_refiner.h"
+#include "tensorflow/core/graph/node_builder.h"
+
+namespace tensorflow {
+namespace ops {
+
+namespace {
+
+// Utility function for converting to internal C++ datatypes.
+OutputTensor ToOutputTensor(const Output& output) {
+  return OutputTensor(output.node(), output.index());
+}
+
+// Utility function for converting to internal C++ datatypes.
+std::vector<OutputTensor> ToOutputTensors(const std::vector<Output>& outputs) {
+  std::vector<OutputTensor> result(outputs.size());
+  for (int i = 0; i < outputs.size(); ++i) {
+    result[i] = ToOutputTensor(outputs[i]);
+  }
+  return result;
+}
+
+// Utility function for converting to internal C++ datatypes.
+std::vector<Node*> ToNodes(const std::vector<Output>& outputs) {
+  std::vector<Node*> result(outputs.size());
+  for (int i = 0; i < outputs.size(); ++i) {
+    result[i] = outputs[i].node();
+  }
+  return result;
+}
+
+// Manually generates the name of the `loop_var_idx`-th NextIteration node of a
+// loop being constructed with `scope`. This is used to define the backedge
+// before the NextIteration node is created.
+std::string NextIterationName(const Scope& scope, int loop_var_idx) {
+  std::string result;
+  const std::string& prefix = scope.impl()->name();
+  if (!prefix.empty()) absl::StrAppend(&result, prefix, "/");
+  absl::StrAppend(&result, "NextIteration");
+  if (loop_var_idx > 0) absl::StrAppend(&result, "_", loop_var_idx);
+  return result;
+}
+
+// Creates the `loop_var_idx`-th Merge node of a loop being constructed with
+// `scope`. `enter_output` is the `loop_var_idx`-th Enter node's output.
+absl::Status CreateMerge(const Scope& scope, int loop_var_idx,
+                         const Output& enter_output, Output* merge_output) {
+  // The merge nodes accept the while loop's back edges as an input (i.e. the
+  // not-yet-created next iteration nodes). Use the underlying NodeBuilder API
+  // directly to create the back edge.
+  NodeBuilder::NodeOut enter_input(enter_output.node(), enter_output.index());
+
+  const int next_output_index = 0;
+  DataType dtype = enter_output.node()->output_type(0);
+  NodeBuilder::NodeOut next_input(NextIterationName(scope, loop_var_idx),
+                                  next_output_index, dtype);
+
+  std::vector<NodeBuilder::NodeOut> input_list({enter_input, next_input});
+  const std::string unique_name = scope.GetUniqueNameForOp("Merge");
+  NodeBuilder builder = NodeBuilder(unique_name, "Merge").Input(input_list);
+  scope.UpdateBuilder(&builder);
+
+  Node* merge_node;
+  TF_RETURN_IF_ERROR(builder.Finalize(scope.graph(), &merge_node));
+  TF_RETURN_IF_ERROR(scope.DoShapeInference(merge_node));
+  *merge_output = Output(merge_node, 0);
+  return absl::OkStatus();
+}
+
+// Creates the condition subgraph defined by `cond`.
+absl::Status CreateCond(const Scope& scope, const CondGraphBuilderFn& cond,
+                        const std::vector<Output>& inputs, Output* output) {
+  // The control dependency is for constants in the cond graph, and other ops
+  // that do not depend on the loop variables. This ensures that these ops are
+  // in the while loop frame (since they will indirectly depend on an Enter node
+  // defining the frame) and that they are executed once per loop iteration.
+  //
+  // TODO(skyewm): the control dep will be added to all nodes in the cond graph.
+  // This is at best unnecessary, and at worst may prevent different parts of
+  // different loop iterations from executing in parallel.
+  Scope cond_scope =
+      scope.NewSubScope("cond").WithControlDependencies(inputs[0]);
+  Output raw_cond_out;
+  TF_RETURN_IF_ERROR(cond(cond_scope, inputs, &raw_cond_out));
+
+  TF_RETURN_IF_ERROR(scope.graph()->IsValidOutputTensor(raw_cond_out.node(),
+                                                        raw_cond_out.index()));
+  if (raw_cond_out.type() != DT_BOOL) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "BuildWhileLoop: 'cond' argument must return a boolean output, got ",
+        DataTypeString(raw_cond_out.type())));
+  }
+  // TODO(skyewm): check that raw_cond_out is scalar
+
+  *output = LoopCond(scope, raw_cond_out).output;
+  return absl::OkStatus();
+}
+
+// Create the body subgraph defined by `body`. `outputs` must be non-null and
+// empty.
+absl::Status CreateBody(const Scope& scope, const BodyGraphBuilderFn& body,
+                        const std::vector<Output>& inputs,
+                        std::vector<Output>* outputs) {
+  DCHECK(outputs != nullptr);
+  DCHECK(outputs->empty());
+
+  // The control dependency is analogous to that in CreateCond().
+  Scope body_scope =
+      scope.NewSubScope("body").WithControlDependencies(inputs[0]);
+  TF_RETURN_IF_ERROR(body(body_scope, inputs, outputs));
+
+  const size_t num_loop_vars = inputs.size();
+  if (outputs->size() != num_loop_vars) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("BuildWhileLoop: 'body' argument expected to return ",
+                     num_loop_vars, " output(s), got ", outputs->size()));
+  }
+  for (const Output& output : *outputs) {
+    TF_RETURN_IF_ERROR(
+        scope.graph()->IsValidOutputTensor(output.node(), output.index()));
+    // TODO(skyewm): check output types/shapes
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+// A while loop with a single loop variable looks like this:
+//
+// (output)
+//     ^    +---------------+
+//     |    | body subgraph +-------------+
+//    Exit  +---------------+             |
+//      ^    ^                            |
+//      |    |                            |
+//      Switch<--------+                  v
+//        ^            |             NextIteration
+//        |     +------+--------+         |
+//        +---->| cond subgraph |         |
+//        |     +---------------+         |
+//       Merge<---------------------------+
+//       ^
+//       |
+//    Enter
+//      ^
+//      |
+//   (input)
+//
+// If there are multiple loop variables, each of the control flow ops is
+// duplicated for each loop variable.
+// TODO(skyewm): link to public version of design doc
+absl::Status BuildWhileLoop(const Scope& scope,
+                            const std::vector<Output>& inputs,
+                            const CondGraphBuilderFn& cond,
+                            const BodyGraphBuilderFn& body,
+                            const std::string& frame_name, OutputList* outputs,
+                            bool create_while_ctx, Output* cond_output) {
+  DCHECK(!inputs.empty());
+  DCHECK(outputs != nullptr);
+  DCHECK(outputs->empty());
+
+  TF_RETURN_IF_ERROR(scope.status());
+  const size_t num_loop_vars = inputs.size();
+
+  std::vector<Output> enter_outputs(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    enter_outputs[i] = internal::Enter(scope, inputs[i], frame_name);
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  std::vector<Output> merge_outputs(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    TF_RETURN_IF_ERROR(
+        CreateMerge(scope, i, enter_outputs[i], &merge_outputs[i]));
+  }
+
+  Output cond_out;
+  TF_RETURN_IF_ERROR(CreateCond(scope, cond, merge_outputs, &cond_out));
+  if (cond_output != nullptr) *cond_output = cond_out;
+
+  std::vector<Output> switch_trues(num_loop_vars);
+  std::vector<Output> switch_falses(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    auto switch_i = Switch(scope, merge_outputs[i], cond_out);
+    switch_trues[i] = switch_i.output_true;
+    switch_falses[i] = switch_i.output_false;
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  std::vector<Output> body_outputs;
+  TF_RETURN_IF_ERROR(CreateBody(scope, body, switch_trues, &body_outputs));
+
+  std::vector<Output> next_outputs(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    next_outputs[i] = NextIteration(scope, body_outputs[i]);
+    DCHECK_EQ(next_outputs[i].node()->name(), NextIterationName(scope, i));
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  // Create the backedges from the NextIteration nodes to the Merge nodes.
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    const int merge_backedge_output_index = 1;
+    scope.graph()->AddEdge(next_outputs[i].node(), next_outputs[i].index(),
+                           merge_outputs[i].node(),
+                           merge_backedge_output_index);
+  }
+
+  outputs->resize(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    (*outputs)[i] = internal::Exit(scope, switch_falses[i]);
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  if (create_while_ctx) {
+    WhileContext* while_ctx;
+    TF_RETURN_IF_ERROR(scope.graph()->AddWhileContext(
+        frame_name, ToNodes(enter_outputs), ToNodes(*outputs),
+        ToOutputTensor(cond_out), ToOutputTensors(switch_trues),
+        ToOutputTensors(body_outputs), &while_ctx));
+
+    // Set while_ctx for all exit nodes. We currently don't require knowing the
+    // while_ctx for any other nodes.
+    for (size_t i = 0; i < num_loop_vars; ++i) {
+      (*outputs)[i].node()->set_while_ctx(while_ctx);
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/ops/while_loop.h

@@ -0,0 +1,80 @@
+/* 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_OPS_WHILE_LOOP_H_
+#define TENSORFLOW_CC_OPS_WHILE_LOOP_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+namespace ops {
+
+// Function that takes cond graph inputs and returns cond graph boolean output.
+// 'output' need not be set if an error is returned.
+typedef std::function<absl::Status(
+    const Scope&, const std::vector<Output>& inputs, Output* output)>
+    CondGraphBuilderFn;
+
+// Function that takes body graph inputs and returns body graph outputs.
+// 'outputs' need not be populated if an error is returned.
+typedef std::function<absl::Status(const Scope&,
+                                   const std::vector<Output>& inputs,
+                                   std::vector<Output>* outputs)>
+    BodyGraphBuilderFn;
+
+// Constructs a while loop.
+//
+// Arguments:
+// * scope: used to construct the while loop.
+// * inputs: the initial values of the loop variables. Must be non-empty.
+// * cond: a function that builds the condition graph of the loop. Takes the
+//     current loop variables as inputs and returns a scalar boolean Output
+//     indicating whether the loop should continue.
+// * body: a function that builds the body graph of the loop. Takes the current
+//     loop variables as inputs and returns the updated loop variables.
+// * frame_name: the frame name to use for this while loop. This should be a
+//     unique name. This will be used as a prefix for created operations.
+// * outputs: output param that returns final loop variable outputs in non-error
+//     case. Must be non-null and empty.
+// * create_while_ctx: if true, a WhileContext is created and populated for this
+//     loop. See core/graph/while_context.h for more details on
+//     WhileContexts. This is set to false for loops used as part of gradient
+//     computations, since they're part of the gradient for a loop in the
+//     forward-pass.
+//     TODO(skyewm): revisit this. Should we create WhileContexts for all loops,
+//     even if we don't need them?
+// * cond_output: if non-null, the output of the predicate is returned. This
+//     will always be a LoopCond node.
+//
+// Returns an error if the while loop could not be fully constructed.
+//
+// TODO(skyewm): clean up partially-constructed loop in error case
+// TODO(skyewm): create public interface to this method
+absl::Status BuildWhileLoop(const Scope& scope,
+                            const std::vector<Output>& inputs,
+                            const CondGraphBuilderFn& cond,
+                            const BodyGraphBuilderFn& body,
+                            const std::string& frame_name, OutputList* outputs,
+                            bool create_while_ctx = true,
+                            Output* cond_output = nullptr);
+
+}  // namespace ops
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_OPS_WHILE_LOOP_H_

tensorflow/cc/ops/while_loop_test.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/ops/while_loop.h"
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/graph/while_context.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+class WhileLoopTest : public ::testing::Test {
+ protected:
+  WhileLoopTest() : 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,
+                  error::Code error_code = error::OK,
+                  const std::string& error_msg = "") {
+    absl::Status s =
+        ops::BuildWhileLoop(scope_, inputs_, cond, body, kFrameName, &outputs_);
+    EXPECT_EQ(s.code(), error_code);
+    EXPECT_EQ(s.message(), error_msg);
+  }
+
+  template <typename T>
+  void Run(const std::vector<Input::Initializer>& input_values,
+           const std::vector<T>& expected_output_values) {
+    ClientSession session(scope_);
+
+    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<Tensor> out_tensors;
+    TF_ASSERT_OK(session.Run(feeds, outputs_, &out_tensors));
+    ASSERT_EQ(out_tensors.size(), outputs_.size());
+
+    DCHECK_EQ(expected_output_values.size(), out_tensors.size());
+    for (int i = 0; i < out_tensors.size(); ++i) {
+      test::ExpectTensorEqual<T>(
+          out_tensors[i], test::AsTensor<T>({expected_output_values[i]}, {}));
+    }
+  }
+
+  Scope scope_;
+  std::vector<Output> inputs_;
+  std::vector<Output> outputs_;
+
+  static const char* const kFrameName;
+};
+
+const char* const WhileLoopTest::kFrameName = "test_loop";
+
+absl::Status LessThanTenCond(const Scope& s, const std::vector<Output>& inputs,
+                             Output* output) {
+  *output = ops::Less(s, inputs[0], 10);
+  return s.status();
+}
+
+absl::Status AddOneBody(const Scope& s, const std::vector<Output>& inputs,
+                        std::vector<Output>* outputs) {
+  outputs->push_back(ops::Add(s, inputs[0], 1));
+  return s.status();
+}
+
+TEST_F(WhileLoopTest, Basic) {
+  // Create loop: while (i < 10) i += 1
+  Init(1);
+  CreateLoop(LessThanTenCond, AddOneBody);
+
+  // Verify some output invariants
+  WhileContext* while_ctx;
+  for (int i = 0; i < outputs_.size(); ++i) {
+    Node* node = outputs_[i].node();
+    ASSERT_TRUE(node->IsExit()) << "Output node " << i << ":\n"
+                                << node->DebugString();
+    ASSERT_TRUE(node->while_ctx() != nullptr) << i;
+    if (i == 0) {
+      while_ctx = node->while_ctx();
+      EXPECT_EQ(while_ctx->frame_name(), kFrameName);
+    } else {
+      EXPECT_EQ(node->while_ctx(), while_ctx) << i;
+    }
+  }
+
+  // Run the loop and test we get the expected results
+  Run<int>({1}, {10});
+  Run<int>({11}, {11});
+}
+
+TEST_F(WhileLoopTest, WrongCondOutputType) {
+  Init(1);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = ops::Placeholder(s, DT_FLOAT);
+        return s.status();
+      },
+      AddOneBody, error::INVALID_ARGUMENT,
+      "BuildWhileLoop: 'cond' argument must return a boolean output, got "
+      "float");
+}
+
+// TODO(skyewm): test bad cond output shape
+
+TEST_F(WhileLoopTest, NullCondOutputNode) {
+  Init(1);
+  // TODO(skyewm): improve error message
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = {nullptr, 0};
+        return s.status();
+      },
+      AddOneBody, error::INVALID_ARGUMENT, "Node is null");
+}
+
+TEST_F(WhileLoopTest, InvalidCondOutputIndex) {
+  Init(1);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        auto less = ops::Less(s, inputs[0], 10);
+        *output = {less.node(), 100};
+        return s.status();
+      },
+      AddOneBody, error::OUT_OF_RANGE,
+      "Node 'cond/Less' (type: 'Less', num of outputs: 1) does not have output "
+      "100");
+}
+
+TEST_F(WhileLoopTest, UnsetCondOutput) {
+  Init(1);
+  CreateLoop([](const Scope& s, const std::vector<Output>& inputs,
+                Output* output) { return s.status(); },
+             AddOneBody, error::INVALID_ARGUMENT, "Node is null");
+}
+
+// TODO(skyewm): test bad body output type
+// TODO(skyewm): test bad body output shape
+
+TEST_F(WhileLoopTest, NullBodyOutputNode) {
+  Init(1);
+  // TODO(skyewm): improve error message
+  CreateLoop(LessThanTenCond,
+             [](const Scope& s, const std::vector<Output>& inputs,
+                std::vector<Output>* outputs) {
+               outputs->push_back({nullptr, 0});
+               return s.status();
+             },
+             error::INVALID_ARGUMENT, "Node is null");
+}
+
+TEST_F(WhileLoopTest, InvalidBodyOutputIndex) {
+  Init(1);
+  CreateLoop(LessThanTenCond,
+             [](const Scope& s, const std::vector<Output>& inputs,
+                std::vector<Output>* outputs) {
+               auto add = ops::Add(s, inputs[0], 1);
+               outputs->emplace_back(add.node(), 100);
+               return s.status();
+             },
+             error::OUT_OF_RANGE,
+             "Node 'body/Add' (type: 'Add', num of outputs: 1) does not have "
+             "output 100");
+}
+
+TEST_F(WhileLoopTest, UnsetBodyOutputs) {
+  Init(1);
+  CreateLoop(
+      LessThanTenCond,
+      [](const Scope& s, const std::vector<Output>& inputs,
+         std::vector<Output>* outputs) { return s.status(); },
+      error::INVALID_ARGUMENT,
+      "BuildWhileLoop: 'body' argument expected to return 1 output(s), got 0");
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/saved_model/BUILD

@@ -0,0 +1,682 @@
+#Description:
+# TensorFlow SavedModel.
+
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_android",
+    "if_google",
+    "if_mobile",
+    "if_not_mobile",
+    "if_not_windows_or_mac",
+    "tf_cc_test",
+)
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "if_static",
+)
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+load(
+    "//tensorflow/security/fuzzing:tf_fuzzing.bzl",
+    "tf_cc_fuzz_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+exports_files([
+    "loader.h",
+    "testdata/chunked_saved_model/chunked_model/saved_model.cpb",
+    "testdata/chunked_saved_model/chunked_model/saved_model.pbtxt",
+])
+
+cc_library(
+    name = "constants",
+    hdrs = ["constants.h"],
+    deps = ["@com_google_absl//absl/strings"],
+)
+
+cc_library(
+    name = "signature_constants",
+    hdrs = ["signature_constants.h"],
+)
+
+cc_library(
+    name = "tag_constants",
+    hdrs = ["tag_constants.h"],
+)
+
+# copybara:uncomment_begin(google-only)
+# cc_library(
+#     name = "mobile_only_deps",
+#     visibility = ["//visibility:private"],
+#     deps = if_mobile(["//tensorflow/core:portable_tensorflow_lib"]),
+# )
+# copybara:uncomment_end
+
+cc_library(
+    name = "reader",
+    srcs = ["reader.cc"],
+    hdrs = ["reader.h"],
+    deps = [
+        ":constants",
+        ":metrics",
+        ":util",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status:statusor",
+    ] + if_google([
+        "//tensorflow/tools/proto_splitter:merge",
+    ]) + if_not_mobile([
+        # TODO(b/111634734): :lib and :protos_all contain dependencies that
+        # cannot be built on mobile platforms. Instead, include the appropriate
+        # tf_lib depending on the build platform.
+        "@com_google_absl//absl/memory:memory",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/util/tensor_bundle:byteswaptensor",
+    ]),
+)
+
+tf_cc_test(
+    name = "reader_test",
+    srcs = ["reader_test.cc"],
+    data = [
+        ":saved_model_test_files",
+    ],
+    linkstatic = 1,
+    deps = [
+        ":constants",
+        ":metrics",
+        ":reader",
+        ":tag_constants",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:resource_loader",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "loader_default_deps",
+    visibility = ["//visibility:private"],
+    deps = if_static([
+        "//tensorflow/core:all_kernels",
+        "//tensorflow/core:direct_session",
+    ]) + if_google(if_static(["//tensorflow/core/platform/default/build_config:tensorflow_platform_specific"])),
+)
+
+cc_library(
+    name = "loader",
+    hdrs = ["loader.h"],
+    deps = [
+        ":loader_lite",
+    ] + select({
+        "//tensorflow:android": [],
+        "//tensorflow:ios": [],
+        "//tensorflow:tflite_converter": ["//tensorflow/core:direct_session"] + if_google([
+            "//tensorflow/core/platform/default/build_config:tensorflow_platform_specific",
+        ]),
+        "//conditions:default": [":loader_default_deps"],
+    }) + if_not_mobile([
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:ops",
+        "//tensorflow/core:protos_all_cc",
+    ]) + if_android([
+        "//tensorflow/core:portable_tensorflow_lib",
+    ]),
+)
+
+cc_library(
+    name = "loader_lite",
+    hdrs = ["loader.h"],
+    deps = if_static([
+        ":loader_lite_impl",
+    ]) + if_not_mobile([
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+    ]),
+)
+
+cc_library(
+    name = "loader_lite_impl",
+    srcs = ["loader.cc"],
+    hdrs = ["loader.h"],
+    deps = [
+        ":constants",
+        ":fingerprinting",
+        ":loader_util",
+        ":reader",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ] + if_not_mobile([
+        ":metrics",
+        ":util",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/util/tensor_bundle:naming",
+    ]),
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "bundle_v2",
+    srcs = ["bundle_v2.cc"],
+    hdrs = ["bundle_v2.h"],
+    deps = [
+        ":constants",
+        ":fingerprinting",
+        ":metrics",
+        ":reader",
+        ":util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:strcat",
+        "//tensorflow/core/platform:tstring",
+        "//tensorflow/core/util/tensor_bundle",
+        "//tensorflow/core/util/tensor_bundle:byteswaptensor",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@jsoncpp_git//:jsoncpp",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:strcat",
+    ],
+)
+
+cc_library(
+    name = "loader_util",
+    srcs = ["loader_util.cc"],
+    hdrs = ["loader_util.h"],
+    deps = [":constants"] + if_not_mobile([
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:protos_all_cc",
+    ]),
+)
+
+tf_cc_test(
+    name = "bundle_v2_test",
+    srcs = ["bundle_v2_test.cc"],
+    data = [
+        ":saved_model_test_files",
+    ],
+    linkstatic = 1,
+    deps = [
+        ":bundle_v2",
+        ":metrics",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:test",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@jsoncpp_git//:jsoncpp",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "saved_model_bundle_test",
+    srcs = ["saved_model_bundle_test.cc"],
+    data = [
+        ":saved_model_test_files",
+    ],
+    linkstatic = 1,
+    deps = [
+        ":constants",
+        ":loader",
+        ":metrics",
+        ":reader",
+        ":signature_constants",
+        ":tag_constants",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+tf_cc_test(
+    name = "saved_model_bundle_lite_test",
+    srcs = ["saved_model_bundle_lite_test.cc"],
+    data = [
+        ":saved_model_test_files",
+    ],
+    linkstatic = 1,
+    deps = [
+        ":constants",
+        ":loader",
+        ":signature_constants",
+        ":tag_constants",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+# A subset of the TF2 saved models can be generated with this tool.
+py_binary(
+    name = "testdata/generate_saved_models",
+    srcs = ["testdata/generate_saved_models.py"],
+    data = [
+        ":saved_model_asset_data",
+        ":saved_model_static_hashtable_asset_data",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:io_ops",
+        "//tensorflow/python/ops:lookup_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model",
+        "//tensorflow/python/saved_model:save_options",
+        "//tensorflow/python/trackable:asset",
+        "@absl_py//absl:app",
+    ],
+)
+
+# copybara:uncomment_begin(google-only)
+#
+# py_binary(
+#     name = "testdata/generate_chunked_models",
+#     srcs = ["testdata/generate_chunked_models.py"],
+#     strict_deps = True,
+#     deps = [
+#         "//tensorflow/python/compat:v2_compat",
+#         "//tensorflow/python/eager:def_function",
+#         "//tensorflow/python/framework:constant_op",
+#         "//tensorflow/python/lib/io:file_io",
+#         "//tensorflow/python/module",
+#         "//tensorflow/python/saved_model:loader",
+#         "//tensorflow/python/saved_model:save",
+#         "//tensorflow/python/saved_model:save_options",
+#         "//tensorflow/python/util:compat",
+#         "//tensorflow/tools/proto_splitter:constants",
+#         "//tensorflow/tools/proto_splitter/python:saved_model",
+#         "//third_party/py/numpy",
+#         "@absl_py//absl:app",
+#         "@absl_py//absl/flags",
+#     ],
+# )
+#
+# copybara:uncomment_end
+
+# TODO(b/32673259): add a test to continuously validate these files.
+filegroup(
+    name = "saved_model_test_files",
+    srcs = glob([
+        "testdata/AssetModule/**",
+        "testdata/half_plus_two_pbtxt/**",
+        "testdata/half_plus_two_main_op/**",
+        "testdata/half_plus_two/**",
+        "testdata/half_plus_two_v2/**",
+        "testdata/x_plus_y_v2_debuginfo/**",
+        "testdata/CyclicModule/**",
+        "testdata/StaticHashTableModule/**",
+        "testdata/VarsAndArithmeticObjectGraph/**",
+        "testdata/fuzz_generated/**",
+        "testdata/SimpleV1Model/**",
+        "testdata/OptimizerSlotVariableModule/**",
+        "testdata/chunked_saved_model/**",
+    ]),
+)
+
+filegroup(
+    name = "saved_model_fingerprinting_test_files",
+    srcs = glob([
+        "testdata/bert2/**",
+        "testdata/bert1/**",
+        "testdata/half_plus_two_pbtxt/**",
+    ]),
+)
+
+alias(
+    name = "saved_model_half_plus_two",
+    actual = ":saved_model_test_files",
+)
+
+filegroup(
+    name = "saved_model_asset_data",
+    srcs = [
+        "testdata/test_asset.txt",
+    ],
+)
+
+filegroup(
+    name = "saved_model_static_hashtable_asset_data",
+    srcs = [
+        "testdata/static_hashtable_asset.txt",
+    ],
+)
+
+exports_files(
+    glob([
+        "testdata/half_plus_two_pbtxt/**",
+        "testdata/half_plus_two_main_op/**",
+        "testdata/half_plus_two/**",
+        "testdata/half_plus_two_v2/**",
+        "testdata/x_plus_y_v2_debuginfo/**",
+        "testdata/CyclicModule/**",
+        "testdata/VarsAndArithmeticObjectGraph/**",
+        "testdata/fuzz_generated/**",
+    ]),
+)
+
+# Linked directly into ":tensorflow_framework".
+cc_library(
+    name = "metrics_impl",
+    srcs = [
+        "metrics.cc",
+        "metrics.h",
+    ],
+    visibility = [
+        "//tensorflow:__pkg__",
+        "//tensorflow/python:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@jsoncpp_git//:jsoncpp",
+    ] + if_not_mobile([
+        "//tensorflow/core:lib",
+    ]) + if_android([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]),
+    alwayslink = True,
+)
+
+cc_library(
+    name = "metrics",
+    hdrs = ["metrics.h"],
+    visibility = [
+        "//tensorflow/cc/saved_model/image_format:__subpackages__",
+        "//tensorflow/python/saved_model:__subpackages__",
+    ],
+    deps = if_static([
+        ":metrics_impl",
+    ]) + if_not_mobile([
+        "//tensorflow/core:lib",
+    ]) + if_android([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + [
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+tf_cc_test(
+    name = "metrics_test",
+    size = "small",
+    srcs = ["metrics_test.cc"],
+    deps = [
+        ":metrics",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:status_matchers",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest",
+        "@jsoncpp_git//:jsoncpp",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:tensor_proto_cc",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+    ] + if_not_mobile(["//tensorflow/core:lib"]) + if_android(["//tensorflow/core:portable_tensorflow_lib_lite"]),
+)
+
+cc_library(
+    name = "test_utils",
+    testonly = True,
+    hdrs = ["test_utils.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:protobuf",
+    ],
+)
+
+tf_cc_test(
+    name = "util_test",
+    size = "small",
+    srcs = ["util_test.cc"],
+    deps = [
+        ":test_utils",
+        ":util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/framework:tensor_shape_proto_cc",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@tsl//tsl/platform:status_matchers",
+    ],
+)
+
+# Linked directly into ":tensorflow_framework".
+cc_library(
+    name = "fingerprinting_impl",
+    srcs = [
+        "fingerprinting.cc",
+        "fingerprinting.h",
+    ],
+    visibility = [
+        "//tensorflow:__pkg__",
+        "//tensorflow/python:__pkg__",
+    ],
+    deps = [
+        ":constants",
+        ":fingerprinting_x_platform_utils",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/graph/regularization:simple_delete",
+        "//tensorflow/core/graph/regularization:util",
+        "//tensorflow/core/util/tensor_bundle:naming",
+        "@com_google_absl//absl/container:btree",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:types",
+    ] + if_not_mobile([
+        "//tensorflow/core:lib",
+    ]) + if_android([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_windows_or_mac([
+        ":fingerprinting_utils",
+        "//tensorflow/tools/proto_splitter/cc:util",
+    ]),
+    alwayslink = True,
+)
+
+cc_library(
+    name = "fingerprinting",
+    hdrs = ["fingerprinting.h"],
+    visibility = [
+        "//learning/brain/contrib/hub/server/ingestion:__subpackages__",
+        "//learning/brain/contrib/tpu_modeling:__subpackages__",
+        "//learning/gemini/deployment/disaggregation/wiz_cc:__subpackages__",
+        "//learning/metadata/artifactoid/cc:__subpackages__",
+        "//learning/tfx/pipeline/util:__subpackages__",
+        "//tensorflow/core/tfrt:__subpackages__",
+        "//tensorflow/python/saved_model:__subpackages__",
+    ],
+    deps = if_static([
+        ":fingerprinting_impl",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/status:statusor",
+        "//tensorflow/core:protos_all_cc",
+    ]) + if_not_mobile([
+        "//tensorflow/core:lib",
+    ]) + if_android([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]),
+)
+
+cc_library(
+    name = "fingerprinting_utils",
+    srcs = ["fingerprinting_utils.cc"],
+    hdrs = ["fingerprinting_utils.h"],
+    visibility = ["//visibility:private"],
+    deps = [
+        ":constants",
+        ":fingerprinting_x_platform_utils",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/util/tensor_bundle:naming",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "//tensorflow/tools/proto_splitter:merge",
+        "//tensorflow/tools/proto_splitter/cc:util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@riegeli//riegeli/bytes:fd_reader",
+        "@riegeli//riegeli/records:record_reader",
+    ],
+    alwayslink = True,
+)
+
+cc_library(
+    name = "fingerprinting_x_platform_utils",
+    srcs = ["fingerprinting_x_platform_utils.cc"],
+    hdrs = ["fingerprinting_x_platform_utils.h"],
+    deps = [
+        "@com_google_absl//absl/numeric:int128",
+        "@com_google_absl//absl/strings:str_format",
+        "@tsl//tsl/platform:random",
+    ],
+)
+
+tf_cc_test(
+    name = "fingerprinting_utils_test",
+    srcs = ["fingerprinting_utils_test.cc"],
+    data = [
+        "//tensorflow/tools/proto_splitter/testdata:many-field.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:split-standard.cpb",
+    ],
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":fingerprinting_utils",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "//tensorflow/tools/proto_splitter/cc:util",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:status_matchers",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:test",
+        "@xla//xla/tsl/lib/core:status_test_util",
+    ],
+)
+
+tf_cc_test(
+    name = "fingerprinting_chunked_test",
+    size = "small",
+    srcs = ["fingerprinting_chunked_test.cc"],
+    data = [
+        ":saved_model_fingerprinting_test_files",
+        ":saved_model_test_files",
+    ],
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":fingerprinting",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:path",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+# copybara:uncomment_end
+
+tf_cc_test(
+    name = "fingerprinting_test",
+    size = "small",
+    srcs = ["fingerprinting_test.cc"],
+    data = [
+        ":saved_model_fingerprinting_test_files",
+        ":saved_model_test_files",
+    ],
+    deps = [
+        ":fingerprinting",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/numeric:int128",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
+tf_cc_fuzz_test(
+    name = "saved_model_fuzz",
+    srcs = ["saved_model_fuzz.cc"],
+    deps = [
+        ":constants",
+        ":loader",
+        ":tag_constants",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+    ],
+)

tensorflow/cc/saved_model/README.md

@@ -0,0 +1 @@
+<!--#include file="../../python/saved_model/README.md"-->

tensorflow/cc/saved_model/bundle_v2.cc

@@ -0,0 +1,234 @@
+/* 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/saved_model/bundle_v2.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/fingerprinting.h"
+#include "tensorflow/cc/saved_model/metrics.h"
+#include "tensorflow/cc/saved_model/reader.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/byte_order.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
+#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
+#include "tensorflow/core/util/tensor_bundle/byte_swap_tensor.h"
+#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/strcat.h"
+
+namespace tensorflow {
+namespace {
+
+using strings::StrCat;
+
+// `tensorflow::SavedModelV2Bundle::Load` API label.
+constexpr char kCCLoadBundleV2Label[] = "cc_load_bundle_v2";
+
+absl::Status ReadCheckpointObjectGraph(BundleReader* bundle_reader,
+                                       TrackableObjectGraph* object_graph) {
+  Tensor object_graph_tensor;
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      bundle_reader->Lookup(kObjectGraphProtoKey, &object_graph_tensor),
+      "SavedModel checkpoint does not contain object graph.");
+  if (object_graph_tensor.dtype() != DT_STRING ||
+      object_graph_tensor.dims() != 0 ||
+      object_graph_tensor.NumElements() != 1) {
+    return absl::Status(
+        absl::StatusCode::kFailedPrecondition,
+        "SavedModel checkpoint object graph was not the correct type.");
+  }
+
+  if (!object_graph->ParseFromString(object_graph_tensor.scalar<tstring>()())) {
+    return absl::Status(
+        absl::StatusCode::kFailedPrecondition,
+        "SavedModel checkpoint object graph could not be deserialized.");
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status SavedModelV2Bundle::Load(const std::string& export_dir,
+                                      SavedModelV2Bundle* const bundle) {
+  metrics::SavedModelReadApi(kCCLoadBundleV2Label).IncrementBy(1);
+  SavedModel saved_model_proto;
+  TF_RETURN_IF_ERROR(ReadSavedModel(export_dir, &saved_model_proto));
+  metrics::SavedModelReadPath().Set(export_dir);
+
+  // Load MetaGraphDef.
+  // In version 2 SavedModels, there is only one MetaGraphDef.
+  if (saved_model_proto.meta_graphs_size() != 1) {
+    return absl::Status(
+        absl::StatusCode::kInvalidArgument,
+        strings::StrCat(
+            "SavedModelV2 should have exactly one MetaGraphDef but actually ",
+            "contains ", saved_model_proto.meta_graphs_size()));
+  }
+  bundle->meta_graph_def_ =
+      std::move(*saved_model_proto.mutable_meta_graphs(0));
+
+  // Correct the endiness of Tensor content on big-endian system
+  if (!port::kLittleEndian) {
+    TF_RETURN_IF_ERROR(
+        ByteSwapTensorContentInMetaGraphDef(&(bundle->meta_graph_def_)));
+  }
+
+  // Load GraphDebugInfo.
+  TF_RETURN_IF_ERROR(
+      ReadSavedModelDebugInfoIfPresent(export_dir, &bundle->debug_info_));
+
+  const std::string variables_dir =
+      io::JoinPath(export_dir, kSavedModelVariablesDirectory);
+  if (!Env::Default()->FileExists(variables_dir).ok()) {
+    LOG(INFO)
+        << "No checkpoint found, assuming this is a program-only SavedModel";
+  } else {
+    // Load the variables checkpoint reader.
+    const std::string variables_prefix =
+        io::JoinPath(variables_dir, kSavedModelVariablesFilename);
+    bundle->variable_reader_ =
+        std::make_unique<BundleReader>(Env::Default(), variables_prefix);
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        bundle->variable_reader_->status(),
+        "Unable to load SavedModel variables checkpoint from ",
+        variables_prefix);
+
+    // Deserialize the object graph proto from the tensor bundle.
+    TF_RETURN_IF_ERROR(ReadCheckpointObjectGraph(
+        bundle->variable_reader_.get(), &bundle->trackable_object_graph_));
+  }
+  // Read the fingerprint.
+  auto fingerprint_proto =
+      saved_model::fingerprinting::ReadSavedModelFingerprint(export_dir);
+  if (fingerprint_proto.ok()) {
+    metrics::SavedModelReadFingerprint().Set(
+        metrics::MakeFingerprintJson(fingerprint_proto.value()));
+
+    TF_ASSIGN_OR_RETURN(
+        std::string path_and_singleprint,
+        metrics::MakeSavedModelPathAndSingleprint(
+            export_dir, saved_model::fingerprinting::Singleprint(
+                            fingerprint_proto.value())));
+    metrics::SavedModelReadPathAndSingleprint().Set(path_and_singleprint);
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status SavedModelV2Bundle::VisitObjectsToRestore(
+    RestoreObjectsCallback callback) {
+  if (saved_object_graph().nodes_size() == 0 ||
+      trackable_object_graph().nodes_size() == 0) {
+    return absl::OkStatus();
+  }
+
+  // Start from root nodes of both the SavedObjectGraph and TrackableObjectGraph
+  // and descend to leaves. Note that the TrackableObjectGraph can have cycles
+  // (as can the SavedObjectGraph).
+  // This is detected and cycle edges are skipped.
+  const SavedObject* root_saved_object = &saved_object_graph().nodes(0);
+  const TrackableObjectGraph::TrackableObject* root_trackable_object =
+      &trackable_object_graph().nodes(0);
+  absl::flat_hash_set<int> trackable_node_ids;
+  return RecurseObjectsToRestore(root_saved_object, 0, root_trackable_object,
+                                 std::string(), &trackable_node_ids,
+                                 std::move(callback));
+}
+
+absl::Status SavedModelV2Bundle::RecurseObjectsToRestore(
+    const SavedObject* saved_object, int saved_object_node_id,
+    const TrackableObjectGraph::TrackableObject* trackable_object,
+    std::string object_name, absl::flat_hash_set<int>* seen_trackable_node_ids,
+    RestoreObjectsCallback callback) {
+  // Callback if any attributes or slot variables.
+  // Note that the root is always excluded from the search (it can never
+  // be a restorable object). This matches some logic on the Python side.
+  if (saved_object_node_id != 0 &&
+      (trackable_object->attributes_size() > 0 ||
+       trackable_object->slot_variables_size() > 0)) {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(
+        callback(saved_object_node_id, *trackable_object), "Unable to restore ",
+        object_name);
+  }
+
+  for (const auto& trackable_child_ref : trackable_object->children()) {
+    const auto& local_name = trackable_child_ref.local_name();
+
+    // Compute the full child name.
+    std::string child_name;
+    if (object_name.empty()) {
+      child_name = local_name;
+    } else {
+      child_name = strings::StrCat(object_name, ".", local_name);
+    }
+
+    // Descend down the trackable graph.
+    int trackable_child_node_id = trackable_child_ref.node_id();
+    if (!seen_trackable_node_ids->insert(trackable_child_node_id).second) {
+      // Cycle or duplicate detected - ignore this branch.
+      continue;
+    }
+    if (trackable_child_node_id < 0 ||
+        trackable_child_node_id >= trackable_object_graph().nodes_size()) {
+      return absl::FailedPreconditionError(
+          strings::StrCat("Illegal trackable child node id for ", child_name));
+    }
+    const auto* trackable_child =
+        &trackable_object_graph().nodes(trackable_child_node_id);
+
+    // Descend down the saved object graph.
+    int saved_child_node_id = -1;
+    const SavedObject* saved_child = nullptr;
+    for (const auto& saved_child_ref : saved_object->children()) {
+      if (saved_child_ref.local_name() == local_name) {
+        // Found.
+        saved_child_node_id = saved_child_ref.node_id();
+        if (saved_child_node_id >= 0 &&
+            saved_child_node_id < saved_object_graph().nodes_size()) {
+          saved_child = &saved_object_graph().nodes(saved_child_node_id);
+        }
+        break;
+      }
+    }
+
+    if (!saved_child) {
+      return absl::Status(
+          absl::StatusCode::kFailedPrecondition,
+          strings::StrCat("Could not find saved object to restore for ",
+                          child_name));
+    }
+
+    TF_RETURN_IF_ERROR(RecurseObjectsToRestore(
+        saved_child, saved_child_node_id, trackable_child, child_name,
+        seen_trackable_node_ids, callback));
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/cc/saved_model/bundle_v2.h

@@ -0,0 +1,90 @@
+/* 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.
+==============================================================================*/
+
+// Helpers for loading the persistent representation of a SavedModelV2.
+// Please note that this is depended on by code that does not make use of
+// the full runtime and its dependencies should be restricted.
+
+#ifndef TENSORFLOW_CC_SAVED_MODEL_BUNDLE_V2_H_
+#define TENSORFLOW_CC_SAVED_MODEL_BUNDLE_V2_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "absl/container/flat_hash_set.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
+#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
+#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
+
+namespace tensorflow {
+
+/// Represents a version 2 SavedModel that is loaded from storage (but not yet
+/// loaded into an executable in-memory representation).
+class SavedModelV2Bundle {
+ public:
+  using RestoreObjectsCallback = std::function<absl::Status(
+      int, const TrackableObjectGraph::TrackableObject&)>;
+
+  /// Loads persistent representations for a SavedModelV2 from the specified
+  /// export directory.
+  static absl::Status Load(const std::string& export_dir,
+                           SavedModelV2Bundle* bundle);
+
+  /// MetaGraphDef from the loaded SavedModel.
+  MetaGraphDef& meta_graph_def() { return meta_graph_def_; }
+
+  /// SavedObjectGraph from the MetaGraphDef.
+  const SavedObjectGraph& saved_object_graph() {
+    return meta_graph_def().object_graph_def();
+  }
+
+  /// TrackableObjectGraph loaded from the variable_reader() checkpoint.
+  TrackableObjectGraph& trackable_object_graph() {
+    return trackable_object_graph_;
+  }
+
+  /// BundleReader for accessing the variables bundle.
+  BundleReader* variable_reader() { return variable_reader_.get(); }
+
+  /// The GraphDebugInfo (or nullptr if none).
+  GraphDebugInfo* debug_info() { return debug_info_.get(); }
+
+  /// Restores objects, invoking the callback with the node id in the
+  /// saved_object_graph() and the corresponding TrackableObject from the
+  /// trackable_object_graph(). The callback may use the variable_reader() but
+  /// must not modify the underlying saved_object_graph().
+  absl::Status VisitObjectsToRestore(RestoreObjectsCallback callback);
+
+ private:
+  absl::Status RecurseObjectsToRestore(
+      const SavedObject* saved_object, int saved_object_node_id,
+      const TrackableObjectGraph::TrackableObject* trackable_object,
+      std::string object_name,
+      absl::flat_hash_set<int>* seen_trackable_node_ids,
+      RestoreObjectsCallback callback);
+
+  MetaGraphDef meta_graph_def_;
+  TrackableObjectGraph trackable_object_graph_;
+  std::unique_ptr<BundleReader> variable_reader_;
+  std::unique_ptr<GraphDebugInfo> debug_info_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_BUNDLE_V2_H_

tensorflow/cc/saved_model/bundle_v2_test.cc

@@ -0,0 +1,158 @@
+/* Copyright 2018 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/saved_model/bundle_v2.h"
+
+#include <algorithm>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "json/json.h"
+#include "json/reader.h"
+#include "json/value.h"
+#include "tensorflow/cc/saved_model/metrics.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace {
+
+constexpr char kTestData[] = "cc/saved_model/testdata";
+
+class BundleV2Test : public ::testing::Test {
+ protected:
+  BundleV2Test() {}
+
+  void RestoreVarsAndVerify(SavedModelV2Bundle* bundle,
+                            std::vector<std::string> expected_names) {
+    // Collect saved_node_id, full_name, checkpoint_key into a vector.
+    using RestoredVarType = std::tuple<int, std::string, std::string>;
+    std::vector<RestoredVarType> restored_vars;
+    TF_ASSERT_OK(bundle->VisitObjectsToRestore(
+        [&](int saved_node_id,
+            const TrackableObjectGraph::TrackableObject& trackable_object)
+            -> absl::Status {
+          for (const auto& attr : trackable_object.attributes()) {
+            if (attr.name() == "VARIABLE_VALUE") {
+              restored_vars.emplace_back(saved_node_id, attr.full_name(),
+                                         attr.checkpoint_key());
+            }
+          }
+          return absl::OkStatus();
+        }));
+
+    // Should be one of each var name restored.
+    for (const auto& expected_name : expected_names) {
+      EXPECT_EQ(1, std::count_if(restored_vars.begin(), restored_vars.end(),
+                                 [&](RestoredVarType t) {
+                                   return std::get<1>(t) == expected_name;
+                                 }));
+    }
+
+    for (const auto& restored_var : restored_vars) {
+      // Each restored var should match a SavedObjectGraph node with the same
+      // variable name.
+      const auto& saved_node =
+          bundle->saved_object_graph().nodes(std::get<0>(restored_var));
+      EXPECT_EQ(std::get<1>(restored_var), saved_node.variable().name());
+
+      // And should be able to load it from the tensor_bundle.
+      Tensor value;
+      TF_ASSERT_OK(
+          bundle->variable_reader()->Lookup(std::get<2>(restored_var), &value));
+    }
+  }
+};
+
+TEST_F(BundleV2Test, LoadsVarsAndArithmeticObjectGraph) {
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), kTestData, "VarsAndArithmeticObjectGraph");
+
+  SavedModelV2Bundle bundle;
+  TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
+
+  // Ensure that there are nodes in the trackable_object_graph.
+  EXPECT_GT(bundle.trackable_object_graph().nodes_size(), 0);
+
+  RestoreVarsAndVerify(&bundle, {"variable_x", "variable_y", "child_variable"});
+}
+
+TEST_F(BundleV2Test, LoadsCyclicModule) {
+  const std::string export_dir =
+      io::JoinPath(testing::TensorFlowSrcRoot(), kTestData, "CyclicModule");
+
+  SavedModelV2Bundle bundle;
+  TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
+
+  // Ensure that there are nodes in the trackable_object_graph.
+  EXPECT_GT(bundle.trackable_object_graph().nodes_size(), 0);
+
+  RestoreVarsAndVerify(&bundle, {"MyVariable"});
+}
+
+TEST_F(BundleV2Test, UpdatesMetrics) {
+  const std::string kCCLoadBundleV2Label = "cc_load_bundle_v2";
+  const int read_count = metrics::SavedModelReadCount("2").value();
+  const int api_count =
+      metrics::SavedModelReadApi(kCCLoadBundleV2Label).value();
+  const std::string export_dir = io::JoinPath(
+      testing::TensorFlowSrcRoot(), kTestData, "VarsAndArithmeticObjectGraph");
+
+  SavedModelV2Bundle bundle;
+  TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
+
+  EXPECT_EQ(metrics::SavedModelReadCount("2").value(), read_count + 1);
+  EXPECT_EQ(metrics::SavedModelReadApi(kCCLoadBundleV2Label).value(),
+            api_count + 1);
+  // Check that the gauge contains the path and fingerprint.
+  EXPECT_EQ(metrics::SavedModelReadPath().value(), export_dir);
+
+  Json::Value fingerprint = Json::objectValue;
+  Json::Reader reader = Json::Reader();
+  reader.parse(metrics::SavedModelReadFingerprint().value(), fingerprint);
+  EXPECT_EQ(fingerprint["saved_model_checksum"].asUInt64(),
+            15788619162413586750ULL);
+  EXPECT_EQ(fingerprint["graph_def_program_hash"].asUInt64(),
+            706963557435316516ULL);
+  EXPECT_EQ(fingerprint["signature_def_hash"].asUInt64(),
+            5693392539583495303ULL);
+  EXPECT_EQ(fingerprint["saved_object_graph_hash"].asUInt64(),
+            12074714563970609759ULL);
+  EXPECT_EQ(fingerprint["checkpoint_hash"].asUInt64(), 10788359570789890102ULL);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto path_and_singleprint,
+      metrics::ParseSavedModelPathAndSingleprint(
+          metrics::SavedModelReadPathAndSingleprint().value()));
+  auto [path, singleprint] = path_and_singleprint;
+  EXPECT_TRUE(absl::StrContains(
+      path, absl::StrCat(kTestData, "/VarsAndArithmeticObjectGraph")));
+  EXPECT_EQ(singleprint,
+            "706963557435316516/"     // graph_def_program_hash
+            "5693392539583495303/"    // signature_def_hash
+            "12074714563970609759/"   // saved_object_graph_hash
+            "10788359570789890102");  // checkpoint_hash
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/saved_model/constants.h

@@ -0,0 +1,82 @@
+/* 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_SAVED_MODEL_CONSTANTS_H_
+#define TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_
+
+namespace tensorflow {
+
+// SavedModel assets directory.
+inline constexpr char kSavedModelAssetsDirectory[] = "assets";
+
+// SavedModel assets.extra directory.
+inline constexpr char kSavedModelAssetsExtraDirectory[] = "assets.extra";
+
+// SavedModel assets key for graph collection-def.
+inline constexpr char kSavedModelAssetsKey[] = "saved_model_assets";
+
+/// SavedModel legacy init op collection key. Used in v1 SavedModels.
+inline constexpr char kSavedModelLegacyInitOpKey[] = "legacy_init_op";
+
+/// SavedModel main op collection key. Used in v1 SavedModels.
+inline constexpr char kSavedModelMainOpKey[] = "saved_model_main_op";
+
+// CollectionDef key for the SavedModel train op.
+// Not exported while export_all_saved_models is experimental.
+inline constexpr char kSavedModelTrainOpKey[] = "saved_model_train_op";
+
+// Schema version for SavedModel.
+inline constexpr int kSavedModelSchemaVersion = 1;
+
+// SavedModel proto filename prefix.
+inline constexpr char kSavedModelFilenamePrefix[] = "saved_model";
+// SavedModel proto filename.
+inline constexpr char kSavedModelFilenamePb[] = "saved_model.pb";
+
+// SavedModel chunked proto filename.
+inline constexpr char kSavedModelFilenameCpb[] = "saved_model.cpb";
+
+// SavedModel text format proto filename.
+inline constexpr char kSavedModelFilenamePbTxt[] = "saved_model.pbtxt";
+
+// Subdirectory where debugging related files are written.
+inline constexpr char kSavedModelDebugDirectory[] = "debug";
+
+// File name for GraphDebugInfo protocol buffer which corresponds to the
+// SavedModel.
+inline constexpr char kSavedModelDebugInfoFilenamePb[] =
+    "saved_model_debug_info.pb";
+
+// Directory in which to save the SavedModel variables.
+inline constexpr char kSavedModelVariablesDirectory[] = "variables";
+
+// SavedModel variables filename.
+inline constexpr char kSavedModelVariablesFilename[] = "variables";
+
+// SavedModel SignatureDef keys for the initialization and train ops. Used in
+// V2 SavedModels.
+inline constexpr char kSavedModelInitOpSignatureKey[] = "__saved_model_init_op";
+inline constexpr char kSavedModelTrainOpSignatureKey[] =
+    "__saved_model_train_op";
+
+// Key in the TensorBundle for the object graph proto.
+inline constexpr char kObjectGraphProtoKey[] = "_CHECKPOINTABLE_OBJECT_GRAPH";
+
+// Filename for the FingerprintDef protocol buffer.
+inline constexpr char kFingerprintFilenamePb[] = "fingerprint.pb";
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_

tensorflow/cc/saved_model/experimental/public/BUILD

@@ -0,0 +1,85 @@
+# Experimental C++ SavedModel Header Only APIs. See RFC
+# https://github.com/tensorflow/community/pull/207
+
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    # This is intentionally public
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "concrete_function",
+    hdrs = [
+        "concrete_function.h",
+    ],
+    deps = [
+        ":function_metadata",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/experimental/saved_model/public:concrete_function",
+        "//tensorflow/cc/experimental/base/public:status",
+    ],
+)
+
+cc_library(
+    name = "concrete_function_list",
+    hdrs = [
+        "concrete_function_list.h",
+    ],
+    deps = [
+        ":concrete_function",
+        "//tensorflow/c/experimental/saved_model/public:concrete_function_list",
+    ],
+)
+
+cc_library(
+    name = "function_metadata",
+    hdrs = [
+        "function_metadata.h",
+    ],
+    deps = [
+        "//tensorflow/c/experimental/saved_model/public:function_metadata",
+    ],
+)
+
+cc_library(
+    name = "saved_model_api",
+    hdrs = [
+        "saved_model_api.h",
+    ],
+    deps = [
+        ":concrete_function",
+        ":concrete_function_list",
+        ":signature_def_function",
+        "//tensorflow/c/experimental/saved_model/public:saved_model_api",
+        "//tensorflow/cc/experimental/base/public:runtime",
+        "//tensorflow/cc/experimental/base/public:status",
+    ],
+)
+
+cc_library(
+    name = "signature_def_function",
+    hdrs = [
+        "signature_def_function.h",
+    ],
+    deps = [
+        ":signature_def_function_metadata",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/experimental/saved_model/public:signature_def_function",
+        "//tensorflow/cc/experimental/base/public:status",
+    ],
+)
+
+cc_library(
+    name = "signature_def_function_metadata",
+    hdrs = [
+        "signature_def_function_metadata.h",
+    ],
+    deps = [
+        "//tensorflow/c/experimental/saved_model/public:signature_def_function_metadata",
+    ],
+)

tensorflow/cc/saved_model/experimental/public/concrete_function.h

@@ -0,0 +1,61 @@
+/* Copyright 2020 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_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_H_
+#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_H_
+
+#include <vector>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/experimental/saved_model/public/concrete_function.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/cc/saved_model/experimental/public/function_metadata.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// ConcreteFunction is an executable "function" loaded from a SavedModelAPI.
+class ConcreteFunction final {
+ public:
+  // TODO(bmzhao): Adding ConcreteFunction::Run in subsequent CL, since
+  // it depends on tensorflow::cc::Tensor and tensorflow::cc::TensorHandle
+
+  // Returns FunctionMetadata associated with this ConcreteFunction.
+  const FunctionMetadata* GetFunctionMetadata();
+
+ private:
+  friend class SavedModelAPI;
+  friend class ConcreteFunctionList;
+
+  // TODO(bmzhao): Consider adding a macro for wrapping/unwrapping
+  // when moving out of experimental.
+  static ConcreteFunction* wrap(TF_ConcreteFunction* p) {
+    return reinterpret_cast<ConcreteFunction*>(p);
+  }
+  static TF_ConcreteFunction* unwrap(ConcreteFunction* p) {
+    return reinterpret_cast<TF_ConcreteFunction*>(p);
+  }
+};
+
+inline const FunctionMetadata* ConcreteFunction::GetFunctionMetadata() {
+  return FunctionMetadata::wrap(TF_ConcreteFunctionGetMetadata(unwrap(this)));
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_H_

tensorflow/cc/saved_model/experimental/public/concrete_function_list.h

@@ -0,0 +1,63 @@
+/* Copyright 2020 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_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
+#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
+
+#include <vector>
+
+#include "tensorflow/c/experimental/saved_model/public/concrete_function_list.h"
+#include "tensorflow/cc/saved_model/experimental/public/concrete_function.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// ConcreteFunctionList helps convert an opaque pointer to an array of
+// ConcreteFunction pointers to a std::vector.
+class ConcreteFunctionList {
+ public:
+  // Converts this object to a std::vector<ConcreteFunction*>
+  std::vector<ConcreteFunction*> ToVector();
+
+ private:
+  friend class SavedModelAPI;
+  // Wraps a TF_ConcreteFunctionList. Takes ownership of list.
+  explicit ConcreteFunctionList(TF_ConcreteFunctionList* list) : list_(list) {}
+
+  struct TFConcreteFunctionListDeleter {
+    void operator()(TF_ConcreteFunctionList* p) const {
+      TF_DeleteConcreteFunctionList(p);
+    }
+  };
+  std::unique_ptr<TF_ConcreteFunctionList, TFConcreteFunctionListDeleter> list_;
+};
+
+inline std::vector<ConcreteFunction*> ConcreteFunctionList::ToVector() {
+  int size = TF_ConcreteFunctionListSize(list_.get());
+  std::vector<ConcreteFunction*> result;
+  result.reserve(size);
+  for (int i = 0; i < size; ++i) {
+    result.push_back(
+        ConcreteFunction::wrap(TF_ConcreteFunctionListGet(list_.get(), i)));
+  }
+  return result;
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_LIST_H_

tensorflow/cc/saved_model/experimental/public/function_metadata.h

@@ -0,0 +1,47 @@
+/* Copyright 2020 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_SAVED_MODEL_EXPERIMENTAL_PUBLIC_FUNCTION_METADATA_H_
+#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_FUNCTION_METADATA_H_
+
+#include <memory>
+
+#include "tensorflow/c/experimental/saved_model/public/function_metadata.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// FunctionMetadata stores additional function information, including
+// optional signaturedef feeds/fetches (for TF1-based ConcreteFunctions),
+// a valid function path (for TF2-based ConcreteFunctions), and
+// the types + number of inputs and outputs.
+class FunctionMetadata final {
+  // TODO(bmzhao): Add getters here as necessary.
+ private:
+  friend class ConcreteFunction;
+  static FunctionMetadata* wrap(TF_FunctionMetadata* p) {
+    return reinterpret_cast<FunctionMetadata*>(p);
+  }
+  static TF_FunctionMetadata* unwrap(FunctionMetadata* p) {
+    return reinterpret_cast<TF_FunctionMetadata*>(p);
+  }
+};
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_FUNCTION_METADATA_H_

tensorflow/cc/saved_model/experimental/public/saved_model_api.h

@@ -0,0 +1,155 @@
+/* Copyright 2020 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_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SAVED_MODEL_API_H_
+#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SAVED_MODEL_API_H_
+
+#include <memory>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "tensorflow/c/experimental/saved_model/public/saved_model_api.h"
+#include "tensorflow/cc/experimental/base/public/runtime.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/cc/saved_model/experimental/public/concrete_function.h"
+#include "tensorflow/cc/saved_model/experimental/public/concrete_function_list.h"
+#include "tensorflow/cc/saved_model/experimental/public/signature_def_function.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// SavedModelAPI offers a way to load Tensorflow Saved Models
+// (https://www.tensorflow.org/guide/saved_model) and execute saved
+// tf.functions or legacy SignatureDefs in a TF2-idiomatic fashion.
+// See RFC 207
+// (https://github.com/tensorflow/community/blob/master/rfcs/20200218-tf-c-saved-model.md)
+// TODO(bmzhao): Add an e2e example here, once ConcreteFunction::Run is added.
+class SavedModelAPI {
+ public:
+  // Load a SavedModel from `dirname`.
+  //
+  // Params:
+  //  saved_model_path - A directory filepath that the SavedModel is at.
+  //  runtime - A runtime used to load SavedModelAPI. `runtime` must outlive the
+  //            returned TF_SavedModel pointer.
+  //  tags - Optional set of tags. If tags = nullptr, we expect the SavedModel
+  //         to contain a single Metagraph (as for those exported from TF2's
+  //         `tf.saved_model.save`). If tags != nullptr, we load the metagraph
+  //         matching the tags:
+  //         https://github.com/tensorflow/tensorflow/blob/428cdeda09aef81e958eeb274b83d27ad635b57b/tensorflow/core/protobuf/meta_graph.proto#L50-L56
+  //  status - Set to OK on success and an appropriate error on failure.
+  // Returns:
+  //  If status is not OK, returns nullptr.
+  static std::unique_ptr<SavedModelAPI> Load(
+      const std::string& saved_model_path, const Runtime& runtime,
+      Status* status, const std::unordered_set<std::string>* tags = nullptr);
+
+  // Retrieve a function from the TF2 SavedModel via function path.
+  //
+  // Params:
+  //  function_path - A string containing the path from the root saved python
+  //                  object to a tf.function method.
+  //  status - Set to OK on success and an appropriate error on failure.
+  // Returns:
+  //  If status is not OK, returns nullptr. Otherwise, returns a
+  //  tensorflow::cc::ConcreteFunction pointer. The lifetime of this pointer
+  //  is bound to SavedModelAPI it was loaded from.
+  ConcreteFunction* GetConcreteFunction(const std::string& function_path,
+                                        Status* status);
+
+  // Retrieve a function from the TF SavedModel via a SignatureDef key.
+  //
+  // Params:
+  //  signature_def_key - String key of SignatureDef map of a SavedModel:
+  //                      https://github.com/tensorflow/tensorflow/blob/69b08900b1e991d84bce31f3b404f5ed768f339f/tensorflow/core/protobuf/meta_graph.proto#L89
+  //  status - Set to OK on success and an appropriate error on failure.
+  // Returns:
+  //  If status is not OK, returns nullptr. Otherwise, returns a
+  //  tensorflow::cc::ConcreteFunction pointer. The lifetime of this pointer
+  //  is bound to SavedModelAPI it was loaded from.
+  SignatureDefFunction* GetSignatureDefFunction(
+      const std::string& function_path, Status* status);
+
+  // SavedModelAPI is movable, but not copyable.
+  SavedModelAPI(SavedModelAPI&&) = default;
+  SavedModelAPI& operator=(SavedModelAPI&&) = default;
+
+ private:
+  SavedModelAPI(const SavedModelAPI&) = delete;
+  SavedModelAPI& operator=(const SavedModelAPI&) = delete;
+
+  explicit SavedModelAPI(TF_SavedModel* model) : saved_model_(model) {}
+  struct TFSavedModelDeleter {
+    void operator()(TF_SavedModel* p) const { TF_DeleteSavedModel(p); }
+  };
+  std::unique_ptr<TF_SavedModel, TFSavedModelDeleter> saved_model_;
+};
+
+inline std::unique_ptr<SavedModelAPI> SavedModelAPI::Load(
+    const std::string& saved_model_path, const Runtime& runtime, Status* status,
+    const std::unordered_set<std::string>* tags) {
+  TF_SavedModel* saved_model = nullptr;
+
+  if (tags == nullptr) {
+    saved_model =
+        TF_LoadSavedModel(saved_model_path.c_str(), runtime.GetTFEContext(),
+                          status->GetTFStatus());
+  } else {
+    std::vector<const char*> tags_vector;
+    tags_vector.reserve(tags->size());
+    for (const std::string& tag : *tags) {
+      tags_vector.push_back(tag.c_str());
+    }
+    saved_model = TF_LoadSavedModelWithTags(
+        saved_model_path.c_str(), runtime.GetTFEContext(), tags_vector.data(),
+        tags_vector.size(), status->GetTFStatus());
+  }
+
+  if (!status->ok()) {
+    return nullptr;
+  }
+
+  // We can't use std::make_unique here because of its interaction with a
+  // private constructor: https://abseil.io/tips/134
+  return std::unique_ptr<SavedModelAPI>(new SavedModelAPI(saved_model));
+}
+
+inline ConcreteFunction* SavedModelAPI::GetConcreteFunction(
+    const std::string& function_path, Status* status) {
+  TF_ConcreteFunction* function = TF_GetSavedModelConcreteFunction(
+      saved_model_.get(), function_path.c_str(), status->GetTFStatus());
+  if (!status->ok()) {
+    return nullptr;
+  }
+  return ConcreteFunction::wrap(function);
+}
+
+inline SignatureDefFunction* SavedModelAPI::GetSignatureDefFunction(
+    const std::string& function_path, Status* status) {
+  TF_SignatureDefFunction* function = TF_GetSavedModelSignatureDefFunction(
+      saved_model_.get(), function_path.c_str(), status->GetTFStatus());
+  if (!status->ok()) {
+    return nullptr;
+  }
+  return SignatureDefFunction::wrap(function);
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SAVED_MODEL_API_H_

tensorflow/cc/saved_model/experimental/public/signature_def_function.h

@@ -0,0 +1,89 @@
+/* Copyright 2020 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_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
+#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
+
+#include <vector>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/experimental/saved_model/public/signature_def_function.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/cc/saved_model/experimental/public/signature_def_function_metadata.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// SignatureDefFunctions are functions that correspond to either:
+// "signatures" saved from a TF2 SavedModel APIs:
+// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/python/saved_model/save.py#L830-L854
+// Or the "SignatureDefMap" saved from TF1 SavedModel APIs:
+// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/python/saved_model/load_v1_in_v2_test.py#L170-L174
+// In both cases, a SignatureDef is serialized as a SignatureDef protobuf:
+// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/core/protobuf/meta_graph.proto#L260-L330
+// and represents a computation defined by a TF subgraph.
+// These Signatures were primarily designed to be interoperable with the legacy
+// TF 1 Session-based C++ SavedModelBundle loading APIs:
+// https://github.com/tensorflow/tensorflow/blob/26c4ee0c833e74f94d0102d8b005c41a28b44445/tensorflow/cc/saved_model/loader.h#L96-L108
+// SignatureDefFunctions have different semantics from regular TF2
+// ConcreteFunctions, and are mainly intended provide a serving-friendly
+// transition point from the TF1 Session API.
+// First, SignatureDefFunctions have different calling conventions.
+// SignatureDefFunctions' inputs and outputs are constrained to **flattened
+// lists of TensorHandles only**. They do not support more exotic input/output
+// types (like optionals, generators, etc). Additionally, this flattening means
+// they will not preserve the exact interface of the original tf.function they
+// were traced from, as things like composite tensors decay into their
+// internal dense tensor representation.
+// Second, all inputs and outputs are "named", and these names are load bearing
+// (eg: they are part of the interface of tensorflow_serving):
+// https://github.com/tensorflow/serving/blob/e0d247b2e4050713194b8fad0be24a0636df7209/tensorflow_serving/apis/predict.proto#L21
+// https://github.com/tensorflow/serving/blob/e0d247b2e4050713194b8fad0be24a0636df7209/tensorflow_serving/apis/predict.proto#L39
+// The name of each input/output is stored in the corresponding tf::Argument in
+// SignatureDefFunctionMetadata::arguments(). Users must ensure the order of
+// TensorHandles passed to the function matches with the order of named
+// arguments. Similarly the name of the outputs is stored in
+// SignatureDefFunctionMetadata::returns().
+class SignatureDefFunction final {
+ public:
+  // Returns FunctionMetadata associated with this ConcreteFunction.
+  const SignatureDefFunctionMetadata* GetFunctionMetadata();
+
+ private:
+  friend class SavedModelAPI;
+  friend class ConcreteFunctionList;
+
+  // TODO(bmzhao): Consider adding a macro for wrapping/unwrapping
+  // when moving out of experimental.
+  static SignatureDefFunction* wrap(TF_SignatureDefFunction* p) {
+    return reinterpret_cast<SignatureDefFunction*>(p);
+  }
+  static TF_SignatureDefFunction* unwrap(SignatureDefFunction* p) {
+    return reinterpret_cast<TF_SignatureDefFunction*>(p);
+  }
+};
+
+inline const SignatureDefFunctionMetadata*
+SignatureDefFunction::GetFunctionMetadata() {
+  return SignatureDefFunctionMetadata::wrap(
+      TF_SignatureDefFunctionGetMetadata(unwrap(this)));
+}
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_

tensorflow/cc/saved_model/experimental/public/signature_def_function_metadata.h

@@ -0,0 +1,47 @@
+/* Copyright 2020 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_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
+#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
+
+#include <memory>
+
+#include "tensorflow/c/experimental/saved_model/public/signature_def_function_metadata.h"
+
+namespace tensorflow {
+namespace experimental {
+namespace cc {
+
+// SignatureDefFunctionMetadata stores additional information on each input
+// and output's names, dtypes, and shape.
+class SignatureDefFunctionMetadata final {
+  // TODO(bmzhao): Add getters here as necessary.
+ private:
+  friend class SignatureDefFunction;
+  static SignatureDefFunctionMetadata* wrap(
+      TF_SignatureDefFunctionMetadata* p) {
+    return reinterpret_cast<SignatureDefFunctionMetadata*>(p);
+  }
+  static TF_SignatureDefFunctionMetadata* unwrap(
+      SignatureDefFunctionMetadata* p) {
+    return reinterpret_cast<TF_SignatureDefFunctionMetadata*>(p);
+  }
+};
+
+}  // namespace cc
+}  // namespace experimental
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_

tensorflow/cc/saved_model/experimental/tests/BUILD

@@ -0,0 +1,28 @@
+# Tests for the C++ header-only SavedModelAPI.
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_cc_test(
+    name = "saved_model_api_test",
+    srcs = [
+        "saved_model_api_test.cc",
+    ],
+    data = [
+        "//tensorflow/cc/saved_model:saved_model_half_plus_two",
+    ],
+    deps = [
+        "//tensorflow/c:tf_status_headers",
+        "//tensorflow/cc/experimental/base/public:runtime",
+        "//tensorflow/cc/experimental/base/public:runtime_builder",
+        "//tensorflow/cc/experimental/base/public:status",
+        "//tensorflow/cc/saved_model/experimental/public:saved_model_api",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)

tensorflow/cc/saved_model/experimental/tests/saved_model_api_test.cc

@@ -0,0 +1,98 @@
+/* Copyright 2020 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/saved_model/experimental/public/saved_model_api.h"
+
+#include <memory>
+#include <string>
+#include <unordered_set>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/cc/experimental/base/public/runtime.h"
+#include "tensorflow/cc/experimental/base/public/runtime_builder.h"
+#include "tensorflow/cc/experimental/base/public/status.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/test.h"
+
+
+namespace {
+
+using tensorflow::experimental::cc::Runtime;
+using tensorflow::experimental::cc::RuntimeBuilder;
+using tensorflow::experimental::cc::SavedModelAPI;
+using tensorflow::experimental::cc::Status;
+
+constexpr char kTestData[] = "cc/saved_model/testdata";
+
+std::string SavedModelPath(absl::string_view saved_model_dir) {
+  return tensorflow::io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),
+                                  kTestData, saved_model_dir);
+}
+
+// This value parameterized test allows us to test both TFRT
+// and non TFRT runtimes.
+// https://github.com/google/googletest/blob/dcc92d0ab6c4ce022162a23566d44f673251eee4/googletest/docs/advanced.md#value-parameterized-tests
+class CPPSavedModelAPITest : public ::testing::TestWithParam<bool> {};
+
+TEST_P(CPPSavedModelAPITest, LoadsSavedModelWithTags) {
+  Status status;
+  RuntimeBuilder builder;
+  bool use_tfrt = GetParam();
+  if (use_tfrt) {
+    GTEST_SKIP();  // TODO(chky) : Enable this once TFRT is open sourced.
+  }
+
+  builder.SetUseTFRT(use_tfrt);
+  std::unique_ptr<Runtime> runtime = builder.Build(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
+  std::unordered_set<std::string> tags = {"serve"};
+  std::unique_ptr<SavedModelAPI> model =
+      SavedModelAPI::Load(model_dir, *runtime, &status, &tags);
+
+  // TODO(bmzhao): Change this to expect TF_OK when loading is implemented.
+  // That unblocks writing other tests that require a TF_SavedModel*,
+  // like loading a ConcreteFunction. This test at least checks that the
+  // C API builds and can be minimally run.
+  EXPECT_EQ(status.code(), TF_UNIMPLEMENTED);
+}
+
+TEST_P(CPPSavedModelAPITest, LoadsSavedModel) {
+  Status status;
+  RuntimeBuilder builder;
+  bool use_tfrt = GetParam();
+  if (use_tfrt) {
+    GTEST_SKIP();  // TODO(chky) : Enable this once TFRT is open sourced.
+  }
+
+  builder.SetUseTFRT(use_tfrt);
+  std::unique_ptr<Runtime> runtime = builder.Build(&status);
+  ASSERT_TRUE(status.ok()) << status.message();
+
+  std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
+  std::unique_ptr<SavedModelAPI> model =
+      SavedModelAPI::Load(model_dir, *runtime, &status);
+
+  EXPECT_EQ(status.code(), TF_OK) << status.message();
+}
+
+INSTANTIATE_TEST_SUITE_P(RuntimeAgnosticCPPSavedModelTests,
+                         CPPSavedModelAPITest, ::testing::Bool());
+
+}  // namespace
+

tensorflow/cc/saved_model/fingerprinting.cc

@@ -0,0 +1,285 @@
+/* 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/saved_model/fingerprinting.h"
+
+#include <cstdint>
+#include <string>
+
+#include "absl/container/btree_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/strings/strip.h"
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h"
+#include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/graph/regularization/simple_delete.h"
+#include "tensorflow/core/graph/regularization/util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system_helper.h"
+#include "tensorflow/core/platform/fingerprint.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/protobuf.h"  // IWYU pragma: keep
+#include "tensorflow/core/protobuf/fingerprint.pb.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
+#include "tensorflow/core/util/tensor_bundle/naming.h"
+// b/291933687, b/291001524
+#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
+#include "tensorflow/cc/saved_model/fingerprinting_utils.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/statusor.h"
+#endif
+// IWYU pragma: no_include "third_party/protobuf/io/coded_stream.h"
+// IWYU pragma: no_include "third_party/protobuf/io/zero_copy_stream_impl_lite.h"
+
+namespace tensorflow::saved_model::fingerprinting {
+
+namespace {
+
+using ::tensorflow::protobuf::Map;
+// NOLINTNEXTLINE: clang-tidy missing-includes false positive
+using ::tensorflow::protobuf::io::CodedOutputStream;
+// NOLINTNEXTLINE: clang-tidy missing-includes false positive
+using ::tensorflow::protobuf::io::StringOutputStream;
+
+// TODO(b/290063184): remove when USM is GA
+uint64_t HashCheckpointIndexFile(absl::string_view model_dir) {
+  std::string meta_filename = MetaFilename(io::JoinPath(
+      model_dir, kSavedModelVariablesDirectory, kSavedModelVariablesFilename));
+  std::string data;
+  absl::Status read_status =
+      ReadFileToString(Env::Default(), meta_filename, &data);
+  if (read_status.ok()) {
+    return tensorflow::Fingerprint64(data);
+  } else {
+    LOG(WARNING) << "Failed to read checkpoint file: " << read_status;
+    return 0;
+  }
+}
+
+uint64_t HashSavedModel(const SavedModel& saved_model) {
+  std::string saved_model_serialized;
+  {
+    // Local scope guarantees coded stream will be trimmed (ensures
+    // serialization determinism).
+    // Unfortunately the saving process itself isn't deterministic, so the
+    // checksum may still change since the saved_model proto may be different.
+    StringOutputStream stream(&saved_model_serialized);
+    CodedOutputStream output(&stream);
+    output.SetSerializationDeterministic(true);
+    saved_model.SerializeToCodedStream(&output);
+  }
+  return tensorflow::Fingerprint64(saved_model_serialized);
+}
+
+uint64_t RegularizeAndHashSignatureDefs(
+    const Map<std::string, SignatureDef>& signature_def_map) {
+  // Sort `signature_def_map`, which is an unordered map from string keys to
+  // SignatureDefs.
+  absl::btree_map<std::string, SignatureDef> sorted_signature_defs;
+  sorted_signature_defs.insert(signature_def_map.begin(),
+                               signature_def_map.end());
+  uint64_t result_hash = 0;
+  for (const auto& item : sorted_signature_defs) {
+    result_hash =
+        FingerprintCat64(result_hash, tensorflow::Fingerprint64(item.first));
+    std::string signature_def_serialized;
+    {
+      StringOutputStream stream(&signature_def_serialized);
+      CodedOutputStream output(&stream);
+      output.SetSerializationDeterministic(true);
+      item.second.SerializeToCodedStream(&output);
+    }
+    result_hash = FingerprintCat64(
+        result_hash, tensorflow::Fingerprint64(signature_def_serialized));
+  }
+  return result_hash;
+}
+
+// The SavedObjectGraph contains two parts: the list of nodes and the map of
+// concrete functions. Regularization treats these two parts separately.
+absl::StatusOr<uint64_t> RegularizeAndHashSavedObjectGraph(
+    const SavedObjectGraph& object_graph_def) {
+  // Sort `concrete_functions`, which is an unordered map from function names to
+  // SavedConcreteFunction, using the suffix UID of the function name. Assumes
+  // that the trackable children are listed in a deterministic order during
+  // serialization.
+  absl::btree_map<int64_t, std::string> uid_to_function_names;
+  for (const auto& [name, concrete_function] :
+       object_graph_def.concrete_functions()) {
+    // All valid function names should end in an UID.
+    TF_ASSIGN_OR_RETURN(int64_t uid, graph_regularization::GetSuffixUID(name));
+    uid_to_function_names.insert({uid, name});
+  }
+  uint64_t result_hash = 0;
+  for (const auto& [uid, function_name] : uid_to_function_names) {
+    // Hash the function name (with the UID stripped).
+    result_hash = FingerprintCat64(result_hash,
+                                   tensorflow::Fingerprint64(absl::StripSuffix(
+                                       function_name, std::to_string(uid))));
+    // Hash the serialized concrete function.
+    std::string concrete_function_serialized;
+    {
+      StringOutputStream stream(&concrete_function_serialized);
+      CodedOutputStream output(&stream);
+      output.SetSerializationDeterministic(true);
+      object_graph_def.concrete_functions()
+          .at(function_name)
+          .SerializeToCodedStream(&output);
+    }
+    result_hash = FingerprintCat64(
+        result_hash, tensorflow::Fingerprint64(concrete_function_serialized));
+  }
+  // TODO(b/241294832): Complete canonicalization of `object_graph_def.nodes`.
+  return result_hash;
+}
+
+void SetFingerprintUUID(FingerprintDef* fingerprint_def) {
+  // Assign a random UUID to the fingerprint. This can happen regardless of
+  // whether the SavedModel was read successfully. It serves as a unique
+  // identifier for the model even in situations where the SavedModel is
+  // non-existent or unreadable.
+  fingerprint_def->set_uuid(CreateRandomUUID());
+}
+
+// Creates a FingerprintDef proto from a SavedModel and the checkpoint meta file
+// (.index) in `export_dir`.
+absl::StatusOr<FingerprintDef> CreateFingerprintDefPb(
+    absl::string_view export_dir, std::string pb_file) {
+  // Version of the code that produced the fingerprint.
+  // Note corresponding version definition in
+  // FingerprintingUtils.cc:CreateFingerprintDefCpb() and in
+  // Fingerprinting.cc:CreateReducedFingerprintDef()
+  const int kFingerprintProducer = 4;
+
+  SavedModel saved_model;
+  TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), pb_file, &saved_model));
+
+  // Create a copy of `metagraph` which will be used and mutated for fingerprint
+  // computation.
+  FingerprintDef fingerprint_def;
+  MetaGraphDef* metagraph = saved_model.mutable_meta_graphs(0);
+  // Set fingerprint field #1.
+  fingerprint_def.set_saved_model_checksum(HashSavedModel(saved_model));
+  // Set fingerprint field #2.
+  graph_regularization::SimpleDelete(*metagraph->mutable_graph_def());
+  fingerprint_def.set_graph_def_program_hash(
+      graph_regularization::ComputeHash(metagraph->graph_def()));
+  // Set fingerprint field #3.
+  fingerprint_def.set_signature_def_hash(
+      RegularizeAndHashSignatureDefs(metagraph->signature_def()));
+  // Set fingerprint field #4.
+  TF_ASSIGN_OR_RETURN(
+      uint64_t object_graph_hash,
+      RegularizeAndHashSavedObjectGraph(metagraph->object_graph_def()));
+  fingerprint_def.set_saved_object_graph_hash(object_graph_hash);
+  // Set fingerprint field #5.
+  fingerprint_def.set_checkpoint_hash(HashCheckpointIndexFile(export_dir));
+  SetFingerprintUUID(&fingerprint_def);
+  // Set version of the fingerprint.
+  VersionDef* version = fingerprint_def.mutable_version();
+  version->set_producer(kFingerprintProducer);
+
+  return fingerprint_def;
+}
+
+absl::StatusOr<FingerprintDef> CreateReducedFingerprintDef() {
+  // Version of the code that produced the fingerprint.
+  // Note corresponding version definition in
+  // FingerprintingUtils.cc:CreateFingerprintDefCpb() and in
+  // Fingerprinting.cc:CreateFingerprintDefPb()
+  const int kFingerprintProducer = 5;
+
+  FingerprintDef fingerprint_def;
+  SetFingerprintUUID(&fingerprint_def);
+  // Set version of the fingerprint.
+  VersionDef* version = fingerprint_def.mutable_version();
+  version->set_producer(kFingerprintProducer);
+
+  return fingerprint_def;
+}
+
+}  // namespace
+
+absl::StatusOr<FingerprintDef> CreateFingerprintDef(
+    absl::string_view export_dir) {
+  std::string prefix = io::JoinPath(export_dir, kSavedModelFilenamePrefix);
+
+#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
+  absl::StatusOr<bool> only_contains_pb =
+      tools::proto_splitter::OnlyContainsPb(prefix);
+  // TODO: b/455882293 - Enable reporting errors if the .[c]pb files are
+  // important but still missing.
+  if (only_contains_pb.ok()) {
+    if (*only_contains_pb) {
+      return CreateFingerprintDefPb(export_dir, absl::StrCat(prefix, ".pb"));
+    }
+    return CreateFingerprintDefCpb(export_dir, absl::StrCat(prefix, ".cpb"));
+  }
+  // At this point we have neither saved_model.pb nor saved_model.cpb.
+  return CreateReducedFingerprintDef();  // Only sets the UUID.
+#else  // The following runs on Windows and Mac.
+  absl::StatusOr<FingerprintDef> fingerprint_def =
+      CreateFingerprintDefPb(export_dir, absl::StrCat(prefix, ".pb"));
+  if (!fingerprint_def.ok()) {
+    return CreateReducedFingerprintDef();
+  }
+  return fingerprint_def;
+#endif
+}
+
+absl::StatusOr<FingerprintDef> ReadSavedModelFingerprint(
+    absl::string_view export_dir) {
+  const std::string fingerprint_pb_path =
+      io::JoinPath(export_dir, kFingerprintFilenamePb);
+  TF_RETURN_IF_ERROR(Env::Default()->FileExists(fingerprint_pb_path));
+
+  FingerprintDef fingerprint_proto;
+  absl::Status result =
+      ReadBinaryProto(Env::Default(), fingerprint_pb_path, &fingerprint_proto);
+  if (!result.ok()) return result;
+
+  return fingerprint_proto;
+}
+
+std::string Singleprint(uint64_t graph_def_program_hash,
+                        uint64_t signature_def_hash,
+                        uint64_t saved_object_graph_hash,
+                        uint64_t checkpoint_hash) {
+  return std::to_string(graph_def_program_hash) + "/" +
+         std::to_string(signature_def_hash) + "/" +
+         std::to_string(saved_object_graph_hash) + "/" +
+         std::to_string(checkpoint_hash);
+}
+
+std::string Singleprint(const FingerprintDef& fingerprint) {
+  return Singleprint(
+      fingerprint.graph_def_program_hash(), fingerprint.signature_def_hash(),
+      fingerprint.saved_object_graph_hash(), fingerprint.checkpoint_hash());
+}
+
+absl::StatusOr<std::string> Singleprint(absl::string_view export_dir) {
+  TF_ASSIGN_OR_RETURN(const FingerprintDef fingerprint_def,
+                      ReadSavedModelFingerprint(export_dir));
+  return Singleprint(fingerprint_def);
+}
+
+}  // namespace tensorflow::saved_model::fingerprinting

tensorflow/cc/saved_model/fingerprinting.h

@@ -0,0 +1,47 @@
+/* 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_SAVED_MODEL_FINGERPRINTING_H_
+#define TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_
+
+#include <string>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/protobuf/fingerprint.pb.h"
+
+namespace tensorflow::saved_model::fingerprinting {
+
+// Creates a FingerprintDef proto from a SavedModel (regular or chunked) and the
+// checkpoint meta file (.index) in `export_dir`.
+absl::StatusOr<FingerprintDef> CreateFingerprintDef(
+    absl::string_view export_dir);
+
+// Loads the `fingerprint.pb` from `export_dir`, returns an error if there is
+// none.
+absl::StatusOr<FingerprintDef> ReadSavedModelFingerprint(
+    absl::string_view export_dir);
+
+// Canonical fingerprinting ID for a SavedModel.
+std::string Singleprint(uint64_t graph_def_program_hash,
+                        uint64_t signature_def_hash,
+                        uint64_t saved_object_graph_hash,
+                        uint64_t checkpoint_hash);
+std::string Singleprint(const FingerprintDef& fingerprint);
+absl::StatusOr<std::string> Singleprint(absl::string_view export_dir);
+
+}  // namespace tensorflow::saved_model::fingerprinting
+
+#endif  // TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_