chore: import upstream snapshot with attribution

tensorflow/c/experimental/gradients/BUILD

@@ -0,0 +1,254 @@
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_libtpu",
+    "tf_cuda_cc_test",
+)
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "tf_cuda_tests_tags",
+)
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+# Library of gradient functions.
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "array_grad",
+    srcs = ["array_grad.cc"],
+    hdrs = [
+        "array_grad.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:gradients_internal",
+    ],
+)
+
+cc_library(
+    name = "math_grad",
+    srcs = ["math_grad.cc"],
+    hdrs = [
+        "math_grad.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/c/experimental/ops:array_ops",
+        "//tensorflow/c/experimental/ops:math_ops",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core/common_runtime/eager:attr_builder",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "nn_grad",
+    srcs = ["nn_grad.cc"],
+    hdrs = [
+        "nn_grad.h",
+    ],
+    visibility = [
+        "//tensorflow/python/framework/experimental:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/c/eager:immediate_execution_context",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/c/experimental/ops:array_ops",
+        "//tensorflow/c/experimental/ops:math_ops",
+        "//tensorflow/c/experimental/ops:nn_ops",
+        "//tensorflow/core/lib/llvm_rtti",
+        "//tensorflow/core/platform:errors",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "not_differentiable",
+    srcs = ["not_differentiable.cc"],
+    hdrs = [
+        "not_differentiable.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:gradients_internal",
+    ],
+)
+
+cc_library(
+    name = "gradients",
+    hdrs = [
+        "array_grad.h",
+        "math_grad.h",
+        "nn_grad.h",
+        "not_differentiable.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":array_grad",
+        ":math_grad",
+        ":nn_grad",
+        ":not_differentiable",
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:gradients_internal",
+    ],
+)
+
+tf_cuda_cc_test(
+    name = "custom_gradient_test",
+    size = "small",
+    srcs = [
+        "custom_gradient_test.cc",
+    ],
+    args = ["--heap_check="],  # TODO(b/174752220): Remove
+    tags = tf_cuda_tests_tags(),
+    deps = [
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:c_api",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/eager:c_api_unified_internal",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/c/eager:unified_api_testutil",
+        "//tensorflow/c/experimental/ops",
+        "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:errors",
+    ],
+)
+
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "array_grad.h",
+        "math_grad.h",
+        "nn_grad.h",
+        "not_differentiable.h",
+    ],
+    visibility = ["//tensorflow/python:__pkg__"],
+)
+
+cc_library(
+    name = "grad_test_helper",
+    testonly = True,
+    srcs = ["grad_test_helper.cc"],
+    hdrs = ["grad_test_helper.h"],
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/c/eager:gradient_checker",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/c/eager:unified_api_testutil",
+        "//tensorflow/c/experimental/gradients/tape:tape_context",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cuda_cc_test(
+    name = "nn_grad_test",
+    size = "small",
+    srcs = [
+        "nn_grad_test.cc",
+    ],
+    args = ["--heap_check="],  # TODO(b/174752220): Remove
+    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
+    deps = [
+        ":grad_test_helper",
+        ":nn_grad",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/c/eager:c_api_test_util",
+        "//tensorflow/c/eager:c_api_unified_internal",
+        "//tensorflow/c/eager:unified_api_testutil",
+        "//tensorflow/c/experimental/gradients/tape:tape_context",
+        "//tensorflow/c/experimental/ops:nn_ops",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:tensor_float_32_utils",
+    ] + if_libtpu(
+        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
+        if_true = [],
+    ),
+)
+
+tf_cuda_cc_test(
+    name = "math_grad_test",
+    size = "small",
+    srcs = [
+        "math_grad_test.cc",
+    ],
+    args = ["--heap_check="],  # TODO(b/174752220): Remove
+    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
+    deps = [
+        ":grad_test_helper",
+        ":math_grad",
+        "//tensorflow/c:tf_datatype_hdrs",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/eager:c_api_test_util",
+        "//tensorflow/c/eager:c_api_unified_internal",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/c/eager:unified_api_testutil",
+        "//tensorflow/c/experimental/gradients/tape:tape_context",
+        "//tensorflow/c/experimental/ops:math_ops",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:tensor_float_32_utils",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
+    ] + if_libtpu(
+        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
+        if_true = [],
+    ),
+)
+
+tf_cuda_cc_test(
+    name = "array_grad_test",
+    size = "small",
+    srcs = [
+        "array_grad_test.cc",
+    ],
+    args = ["--heap_check="],  # TODO(b/174752220): Remove
+    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
+    deps = [
+        ":array_grad",
+        ":grad_test_helper",
+        "//tensorflow/c:tf_status_helper",
+        "//tensorflow/c/eager:c_api_test_util",
+        "//tensorflow/c/eager:c_api_unified_internal",
+        "//tensorflow/c/eager:unified_api_testutil",
+        "//tensorflow/c/experimental/gradients/tape:tape_context",
+        "//tensorflow/c/experimental/ops:array_ops",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:tensor_float_32_utils",
+    ] + if_libtpu(
+        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
+        if_true = [],
+    ),
+)

tensorflow/c/experimental/gradients/array_grad.cc

@@ -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.
+==============================================================================*/
+#include "tensorflow/c/experimental/gradients/array_grad.h"
+
+#include "tensorflow/c/eager/abstract_context.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace {
+class IdentityNGradientFunction : public GradientFunction {
+ public:
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    for (int i = 0; i < grad_outputs.size(); i++) {
+      auto grad_input = grad_outputs[i];
+      // TODO(srbs): Should we add a copy contructor to AbstractTensorHandle
+      // that takes care of this similar to `Tensor`?
+      if (grad_input) {
+        grad_input->Ref();
+      }
+      grad_inputs[i] = grad_input;
+    }
+    return absl::OkStatus();
+  }
+  ~IdentityNGradientFunction() override {}
+};
+}  // namespace
+
+GradientFunction* IdentityNRegisterer(const ForwardOperation& op) {
+  return new IdentityNGradientFunction;
+}
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/array_grad.h

@@ -0,0 +1,26 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_ARRAY_GRAD_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_ARRAY_GRAD_H_
+
+#include "tensorflow/c/eager/gradients.h"
+
+namespace tensorflow {
+namespace gradients {
+GradientFunction* IdentityNRegisterer(const ForwardOperation& op);
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_ARRAY_GRAD_H_

tensorflow/c/experimental/gradients/array_grad_test.cc

@@ -0,0 +1,133 @@
+/* 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/c/experimental/gradients/array_grad.h"
+
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/experimental/gradients/grad_test_helper.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
+#include "tensorflow/c/experimental/ops/array_ops.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+namespace {
+
+using tensorflow::TF_StatusPtr;
+
+absl::Status IdentityNModel(AbstractContext* ctx,
+                            absl::Span<AbstractTensorHandle* const> inputs,
+                            absl::Span<AbstractTensorHandle*> outputs) {
+  std::vector<AbstractTensorHandle*> temp_outputs(2);
+  TF_RETURN_IF_ERROR(
+      ops::IdentityN(ctx, inputs, absl::MakeSpan(temp_outputs), "IdentityN"));
+  // Although, `ops::IdentityN` returns 2 tensors, the first tensor isn't needed
+  // for computing gradient so we could safely drop it.
+  outputs[0] = temp_outputs[1];
+  temp_outputs[0]->Unref();
+  return absl::OkStatus();
+}
+
+class CppGradients
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+    status_ = StatusFromTF_Status(status.get());
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+    {
+      AbstractContext* ctx_raw = nullptr;
+      status_ =
+          BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+      ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+      immediate_execution_ctx_.reset(ctx_raw);
+    }
+
+    // Computing numerical gradients with TensorFloat-32 is numerically
+    // unstable. Some forward pass tests also fail with TensorFloat-32 due to
+    // low tolerances
+    enable_tensor_float_32_execution(false);
+  }
+
+  AbstractContextPtr immediate_execution_ctx_;
+  GradientRegistry registry_;
+  absl::Status status_;
+
+ public:
+  bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
+  bool UseFunction() const { return std::get<2>(GetParam()); }
+};
+
+TEST_P(CppGradients, TestIdentityNGrad) {
+  // This test is interesting because the current implementation of GradientTape
+  // would return [0, 1] whereas we use build_default_zeros_grads=false here
+  // so we get back [nullptr, 1].
+
+  AbstractTensorHandlePtr x1;
+  {
+    AbstractTensorHandle* x1_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 1.0f, &x1_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x1.reset(x1_raw);
+  }
+
+  AbstractTensorHandlePtr x2;
+  {
+    AbstractTensorHandle* x2_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 1.0f, &x2_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x2.reset(x2_raw);
+  }
+
+  status_ = registry_.Register("IdentityN", IdentityNRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+  auto IdentityNGradModel = BuildGradModel(IdentityNModel, registry_);
+
+  std::vector<AbstractTensorHandle*> outputs(2);
+  status_ =
+      RunModel(IdentityNGradModel, immediate_execution_ctx_.get(),
+               {x1.get(), x2.get()}, absl::MakeSpan(outputs), UseFunction());
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+  EXPECT_EQ(outputs[0], nullptr);
+  ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[1], {1.0f}, /*dims*/ {},
+                                           /*abs_error*/ 0));
+  outputs[1]->Unref();
+}
+
+#ifdef PLATFORM_GOOGLE
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#else
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#endif
+}  // namespace
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/custom_gradient_test.cc

@@ -0,0 +1,138 @@
+/* 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 <memory>
+
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+namespace {
+using std::vector;
+
+class CustomGradientTest
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+    absl::Status s = StatusFromTF_Status(status.get());
+    CHECK_EQ(errors::OK, s.code()) << s.message();
+  }
+};
+
+class PassThroughGradientFunction : public GradientFunction {
+ public:
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    CHECK_EQ(grad_outputs.size(), 1);
+    CHECK_EQ(grad_inputs.size(), 1);
+    grad_inputs[0] = grad_outputs[0];
+    if (grad_inputs[0]) {
+      grad_inputs[0]->Ref();
+    }
+    return absl::OkStatus();
+  }
+};
+
+// Computes:
+//
+// @tf.custom_gradient
+// def f(input):
+//   def grad(grads):
+//     return grads[0]
+//   return tf.exp(input), grad
+// outputs = [f(inputs[0])]
+absl::Status ExpWithPassThroughGrad(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
+    absl::Span<AbstractTensorHandle*> outputs) {
+  Tape tape(/*persistent=*/false);
+  tape.Watch(inputs[0]);  // Watch x.
+  AbstractTensorHandle* exp_output;
+  TF_RETURN_IF_ERROR(ops::Exp(ctx, inputs[0], &exp_output, "Exp"));
+  std::unique_ptr<GradientFunction> gradient_function(
+      new PassThroughGradientFunction);
+  tape.RecordOperation(inputs, {exp_output}, gradient_function.release());
+  TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx,
+                                          /*targets*/ {exp_output},
+                                          /*sources=*/inputs,
+                                          /*output_gradients=*/{},
+                                          /*result=*/outputs));
+  exp_output->Unref();
+  return absl::OkStatus();
+}
+
+TEST_P(CustomGradientTest, ExpWithPassThroughGrad) {
+  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+      TF_NewStatus(), TF_DeleteStatus);
+  AbstractContextPtr ctx;
+  {
+    AbstractContext* ctx_raw = nullptr;
+    absl::Status s =
+        BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    ctx.reset(ctx_raw);
+  }
+
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    absl::Status s =
+        TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 1.0f, &x_raw);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    x.reset(x_raw);
+  }
+
+  // Pseudo-code:
+  //
+  // tape.watch(x)
+  // y = exp(x)
+  // outputs = tape.gradient(y, x)
+  std::vector<AbstractTensorHandle*> outputs(1);
+  absl::Status s = RunModel(ExpWithPassThroughGrad, ctx.get(), {x.get()},
+                            absl::MakeSpan(outputs),
+                            /*use_function=*/!std::get<2>(GetParam()));
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+
+  TF_Tensor* result_tensor;
+  s = GetValue(outputs[0], &result_tensor);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+  auto result_value = static_cast<float*>(TF_TensorData(result_tensor));
+  EXPECT_EQ(*result_value, 1.0);
+  outputs[0]->Unref();
+  TF_DeleteTensor(result_tensor);
+  result_tensor = nullptr;
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CustomGradientTest, CustomGradientTest,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*executing_eagerly*/ ::testing::Values(true, false)));
+
+}  // namespace
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/grad_test_helper.cc

@@ -0,0 +1,136 @@
+/* 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/c/experimental/gradients/grad_test_helper.h"
+
+#include "tensorflow/c/eager/gradient_checker.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+
+void CompareNumericalAndAutodiffGradients(
+    Model model, Model grad_model, AbstractContext* ctx,
+    absl::Span<AbstractTensorHandle* const> inputs, bool use_function,
+    double abs_error) {
+  auto num_inputs = inputs.size();
+  std::vector<AbstractTensorHandle*> outputs(num_inputs);
+  auto s = RunModel(grad_model, ctx, inputs, absl::MakeSpan(outputs),
+                    /*use_function=*/use_function);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+
+  for (int i = 0; i < num_inputs; ++i) {
+    if (!outputs[i]) continue;
+
+    AbstractTensorHandlePtr numerical_grad;
+    {
+      AbstractTensorHandle* numerical_grad_raw;
+      s = CalcNumericalGrad(ctx, model, inputs,
+                            /*input_index=*/i, use_function,
+                            &numerical_grad_raw);
+      ASSERT_EQ(errors::OK, s.code()) << s.message();
+      numerical_grad.reset(numerical_grad_raw);
+    }
+
+    TF_Tensor* numerical_tensor;
+    s = GetValue(numerical_grad.get(), &numerical_tensor);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    auto num_elem_numerical = TF_TensorElementCount(numerical_tensor);
+
+    TF_Tensor* analytical_tensor;
+    s = GetValue(outputs[i], &analytical_tensor);
+    ASSERT_EQ(errors::OK, s.code()) << s.message();
+    auto num_elem_analytical = TF_TensorElementCount(analytical_tensor);
+
+    ASSERT_EQ(num_elem_numerical, num_elem_analytical);
+
+    float* dnumerical = new float[num_elem_numerical]{0};
+    memcpy(&dnumerical[0], TF_TensorData(numerical_tensor),
+           TF_TensorByteSize(numerical_tensor));
+    float* danalytical = new float[num_elem_analytical]{0};
+    memcpy(&danalytical[0], TF_TensorData(analytical_tensor),
+           TF_TensorByteSize(analytical_tensor));
+
+    for (int j = 0; j < num_elem_numerical; j++) {
+      ASSERT_NEAR(dnumerical[j], danalytical[j], abs_error);
+    }
+    TF_DeleteTensor(analytical_tensor);
+    TF_DeleteTensor(numerical_tensor);
+    delete[] danalytical;
+    delete[] dnumerical;
+    outputs[i]->Unref();
+  }
+}
+
+void CheckTensorValue(AbstractTensorHandle* t, absl::Span<const float> manuals,
+                      absl::Span<const int64_t> dims, double abs_error) {
+  TF_Tensor* analytical_tensor;
+  auto s = GetValue(t, &analytical_tensor);
+  ASSERT_EQ(errors::OK, s.code()) << s.message();
+
+  int64_t num_elem_analytical = 1;
+  auto num_dims_analytical = TF_NumDims(analytical_tensor);
+  ASSERT_EQ(dims.size(), num_dims_analytical);
+  for (int j = 0; j < num_dims_analytical; j++) {
+    auto dim_analytical = TF_Dim(analytical_tensor, j);
+    ASSERT_EQ(dims[j], dim_analytical);
+    num_elem_analytical *= dim_analytical;
+  }
+
+  float* danalytical = new float[num_elem_analytical]{0};
+  memcpy(&danalytical[0], TF_TensorData(analytical_tensor),
+         TF_TensorByteSize(analytical_tensor));
+
+  for (int64_t j = 0; j < num_elem_analytical; j++) {
+    if (abs_error == 0) {
+      ASSERT_EQ(manuals[j], danalytical[j]);
+    } else {
+      ASSERT_NEAR(manuals[j], danalytical[j], abs_error);
+    }
+  }
+
+  TF_DeleteTensor(analytical_tensor);
+  delete[] danalytical;
+}
+
+Model BuildGradModel(Model forward, GradientRegistry registry) {
+  return [forward_model = std::move(forward),
+          grad_registry = std::move(registry)](
+             AbstractContext* ctx,
+             absl::Span<AbstractTensorHandle* const> inputs,
+             absl::Span<AbstractTensorHandle*> outputs) -> absl::Status {
+    Tape tape(/*persistent=*/false);
+    for (size_t i{}; i < inputs.size(); ++i) {
+      tape.Watch(inputs[i]);
+    }
+    std::vector<AbstractTensorHandle*> temp_outputs(1);
+    AbstractContextPtr tape_ctx(new TapeContext(ctx, &tape, grad_registry));
+    TF_RETURN_IF_ERROR(
+        forward_model(tape_ctx.get(), inputs, absl::MakeSpan(temp_outputs)));
+
+    TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx, /*targets=*/temp_outputs,
+                                            /*sources=*/inputs,
+                                            /*output_gradients=*/{}, outputs));
+    for (auto temp_output : temp_outputs) {
+      temp_output->Unref();
+    }
+    return absl::OkStatus();
+  };
+}
+
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/grad_test_helper.h

@@ -0,0 +1,39 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_GRAD_TEST_HELPER_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_GRAD_TEST_HELPER_H_
+
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+
+void CompareNumericalAndAutodiffGradients(
+    Model model, Model grad_model, AbstractContext* ctx,
+    absl::Span<AbstractTensorHandle* const> inputs, bool use_function,
+    double abs_error = 1e-2);
+
+void CheckTensorValue(AbstractTensorHandle* t, absl::Span<const float> manuals,
+                      absl::Span<const int64_t> dims, double abs_error = 1e-2);
+
+Model BuildGradModel(Model forward, GradientRegistry registry);
+
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_GRAD_TEST_HELPER_H_

tensorflow/c/experimental/gradients/math_grad.cc

@@ -0,0 +1,517 @@
+/* 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/c/experimental/gradients/math_grad.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/experimental/ops/array_ops.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/eager/attr_builder.h"
+#include "tensorflow/core/framework/types.h"
+
+using std::vector;
+using tensorflow::ops::AddV2;
+using tensorflow::ops::Div;
+using tensorflow::ops::DivNoNan;
+using tensorflow::ops::MatMul;
+using tensorflow::ops::Mul;
+using tensorflow::ops::Neg;
+using tensorflow::ops::OnesLike;
+using tensorflow::ops::SqrtGrad;
+
+namespace tensorflow {
+namespace gradients {
+namespace {
+
+static absl::Status SafeConj(AbstractContext* ctx,
+                             AbstractTensorHandle* const input,
+                             AbstractTensorHandle** output, const char* name) {
+  auto dtype = input->DataType();
+  if (DataTypeIsFloating(BaseType(dtype)) ||
+      DataTypeIsInteger(BaseType(dtype))) {
+    return tensorflow::ops::Identity(ctx, input, output, name);
+  } else if (!DataTypeIsComplex(BaseType(dtype)) &&
+             BaseType(dtype) != DT_VARIANT) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Expected numeric or variant tensor, got dtype ", dtype));
+  }
+  return tensorflow::ops::Conj(ctx, input, output, name);
+}
+
+class AddGradientFunction : public GradientFunction {
+ public:
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    // TODO(b/161805092): Support broadcasting.
+
+    DCHECK(grad_outputs[0]);
+    grad_inputs[0] = grad_outputs[0];
+    grad_inputs[1] = grad_outputs[0];
+
+    grad_inputs[0]->Ref();
+    grad_inputs[1]->Ref();
+    return absl::OkStatus();
+  }
+  ~AddGradientFunction() override = default;
+};
+
+class ExpGradientFunction : public GradientFunction {
+ public:
+  explicit ExpGradientFunction(AbstractTensorHandle* exp) : exp_(exp) {
+    exp->Ref();
+  }
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    AbstractTensorHandle* conj_output;
+    std::string name = "Conj_Exp_Grad";
+    TF_RETURN_IF_ERROR(SafeConj(ctx, exp_.get(), &conj_output, name.c_str()));
+    AbstractTensorHandlePtr conj_output_releaser(conj_output);
+
+    name = "Mul_Exp_Grad";
+    TF_RETURN_IF_ERROR(
+        Mul(ctx, conj_output, grad_outputs[0], &grad_inputs[0], name.c_str()));
+    return absl::OkStatus();
+  }
+  ~ExpGradientFunction() override = default;
+
+ private:
+  AbstractTensorHandlePtr exp_;
+};
+
+class SqrtGradientFunction : public GradientFunction {
+ public:
+  explicit SqrtGradientFunction(AbstractTensorHandle* sqrt) : sqrt_(sqrt) {
+    sqrt->Ref();
+  }
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    std::string name = "Sqrt_Grad";
+    TF_RETURN_IF_ERROR(SqrtGrad(ctx, sqrt_.get(), grad_outputs[0],
+                                &grad_inputs[0], name.c_str()));
+    return absl::OkStatus();
+  }
+  ~SqrtGradientFunction() override = default;
+
+ private:
+  AbstractTensorHandlePtr sqrt_;
+};
+
+class MatMulGradientFunction : public GradientFunction {
+ public:
+  explicit MatMulGradientFunction(vector<AbstractTensorHandle*> f_inputs,
+                                  AttrBuilder f_attrs)
+      : forward_inputs_(f_inputs), forward_attrs_(f_attrs) {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Ref();
+      }
+    }
+  }
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    /* Given upstream grad U and a matmul op A*B, the gradients are:
+     *
+     *    dA = U * B.T
+     *    dB = A.T * U
+     *
+     *    where A.T means `transpose(A)`
+     */
+    AbstractTensorHandle* upstream_grad = grad_outputs[0];
+
+    // Get transpose attrs
+    bool t_a;
+    TF_RETURN_IF_ERROR(forward_attrs_.Get("transpose_a", &t_a));
+
+    bool t_b;
+    TF_RETURN_IF_ERROR(forward_attrs_.Get("transpose_b", &t_b));
+
+    // Conj each input
+    AbstractTensorHandle* conj_output;
+    std::string name = "Conj_A_MatMul_Grad";
+    TF_RETURN_IF_ERROR(
+        SafeConj(ctx, forward_inputs_[0], &conj_output, name.c_str()));
+
+    AbstractTensorHandlePtr A(conj_output);
+
+    name = "Conj_B_MatMul_Grad";
+    TF_RETURN_IF_ERROR(
+        SafeConj(ctx, forward_inputs_[1], &conj_output, name.c_str()));
+
+    AbstractTensorHandlePtr B(conj_output);
+
+    // Calc Grad
+    AbstractTensorHandle* matmul_A_output;
+    AbstractTensorHandle* matmul_B_output;
+    std::string name_grad_A = "MatMul_Grad_A";
+    std::string name_grad_B = "MatMul_Grad_B";
+    if (!t_a && !t_b) {
+      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, B.get(), &matmul_A_output,
+                                /*transpose_a = */ false,
+                                /*transpose_b = */ true,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_A.c_str()));
+
+      TF_RETURN_IF_ERROR(MatMul(ctx, A.get(), upstream_grad, &matmul_B_output,
+                                /*transpose_a = */ true,
+                                /*transpose_b = */ false,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_B.c_str()));
+    } else if (!t_a && t_b) {
+      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, B.get(), &matmul_A_output,
+                                /*transpose_a = */ false,
+                                /*transpose_b = */ false,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_A.c_str()));
+
+      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, A.get(), &matmul_B_output,
+                                /*transpose_a = */ true,
+                                /*transpose_b = */ false,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_B.c_str()));
+
+    } else if (t_a && !t_b) {
+      TF_RETURN_IF_ERROR(MatMul(ctx, B.get(), upstream_grad, &matmul_A_output,
+                                /*transpose_a = */ false,
+                                /*transpose_b = */ true,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_A.c_str()));
+
+      TF_RETURN_IF_ERROR(MatMul(ctx, A.get(), upstream_grad, &matmul_B_output,
+                                /*transpose_a = */ false,
+                                /*transpose_b = */ false,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_B.c_str()));
+    } else {  // t_a && t_b
+      TF_RETURN_IF_ERROR(MatMul(ctx, B.get(), upstream_grad, &matmul_A_output,
+                                /*transpose_a = */ true,
+                                /*transpose_b = */ true,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_A.c_str()));
+
+      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, A.get(), &matmul_B_output,
+                                /*transpose_a = */ true,
+                                /*transpose_b = */ true,
+                                /*grad_a = */ false, /*grad_b = */ false,
+                                name_grad_B.c_str()));
+    }
+
+    // Gradient for A
+    grad_inputs[0] = matmul_A_output;
+
+    // Gradient for B
+    grad_inputs[1] = matmul_B_output;
+    return absl::OkStatus();
+  }
+  ~MatMulGradientFunction() override {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Unref();
+      }
+    }
+  }
+
+ private:
+  // TODO(b/174778737): Only hold needed inputs.
+  vector<AbstractTensorHandle*> forward_inputs_;
+  AttrBuilder forward_attrs_;
+};
+
+class NegGradientFunction : public GradientFunction {
+ public:
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    /* Given upstream grad U and a Neg op Y = -X, the gradients are:
+     *
+     *    dX =  -U
+     *
+     */
+
+    std::string name = "Neg_Grad";
+    TF_RETURN_IF_ERROR(
+        ops::Neg(ctx, grad_outputs[0], &grad_inputs[0], name.c_str()));
+    return absl::OkStatus();
+  }
+  ~NegGradientFunction() override = default;
+};
+
+class SubGradientFunction : public GradientFunction {
+ public:
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    /* Given upstream grad U and a Sub op A-B, the gradients are:
+     *
+     *    dA =  U
+     *    dB = -U
+     *
+     */
+
+    // Grad for A
+    DCHECK(grad_outputs[0]);
+    grad_inputs[0] = grad_outputs[0];
+    grad_inputs[0]->Ref();
+
+    // Grad for B
+    // negate the upstream grad
+    std::string name = "Neg_Sub_Grad_B";
+    TF_RETURN_IF_ERROR(
+        ops::Neg(ctx, grad_outputs[0], &grad_inputs[1], name.c_str()));
+
+    return absl::OkStatus();
+  }
+  ~SubGradientFunction() override = default;
+};
+
+class MulGradientFunction : public GradientFunction {
+ public:
+  explicit MulGradientFunction(vector<AbstractTensorHandle*> f_inputs)
+      : forward_inputs_(f_inputs) {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Ref();
+      }
+    }
+  }
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    /* Given upstream grad U and a mul op A*B, the gradients are:
+     *
+     *    dA = U * B
+     *    dB = A * U
+     *
+     */
+
+    AbstractTensorHandle* upstream_grad = grad_outputs[0];
+
+    // Gradient for A
+    std::string name = "Mul_Grad_A";
+    TF_RETURN_IF_ERROR(Mul(ctx, upstream_grad, forward_inputs_[1],
+                           &grad_inputs[0], name.c_str()));
+
+    // Gradient for B
+    name = "Mul_Grad_B";
+    TF_RETURN_IF_ERROR(Mul(ctx, forward_inputs_[0], upstream_grad,
+                           &grad_inputs[1], name.c_str()));
+    return absl::OkStatus();
+  }
+  ~MulGradientFunction() override {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Unref();
+      }
+    }
+  }
+
+ private:
+  // TODO(b/174778737): Only hold needed inputs.
+  vector<AbstractTensorHandle*> forward_inputs_;
+};
+
+class Log1pGradientFunction : public GradientFunction {
+ public:
+  explicit Log1pGradientFunction(vector<AbstractTensorHandle*> f_inputs)
+      : forward_inputs_(f_inputs) {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Ref();
+      }
+    }
+  }
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    // TODO(vnvo2409): Add control dependency
+    /* Given upstream grad U and a Log1p op: Y = log(1 + X), the gradients are:
+     *
+     *    dX = U / (1 + X)
+     *
+     */
+
+    AbstractTensorHandle* upstream_grad = grad_outputs[0];
+    AbstractTensorHandle* X = forward_inputs_[0];
+
+    AbstractTensorHandle* temp_output;
+
+    // Calculate conjugate of X
+    std::string name = "Conj_Log1p_Grad_X";
+    TF_RETURN_IF_ERROR(SafeConj(ctx, X, &temp_output, name.c_str()));
+
+    AbstractTensorHandlePtr Conj_X(temp_output);
+
+    // Creates Ones
+    name = "OnesLike_Log1p_Grad_X";
+    TF_RETURN_IF_ERROR(OnesLike(ctx, Conj_X.get(), &temp_output, name.c_str()));
+
+    AbstractTensorHandlePtr Ones_X(temp_output);
+
+    name = "Add_Log1p_Grad_X";
+    // Calculate 1 + Conj(X)
+    TF_RETURN_IF_ERROR(
+        AddV2(ctx, Ones_X.get(), Conj_X.get(), &temp_output, name.c_str()));
+
+    AbstractTensorHandlePtr Conj_XP1(temp_output);
+
+    name = "Div_Log1p_Grad_X";
+    // Calculate U / (1 + Conj(X))
+    TF_RETURN_IF_ERROR(
+        Div(ctx, upstream_grad, Conj_XP1.get(), &grad_inputs[0], name.c_str()));
+
+    return absl::OkStatus();
+  }
+  ~Log1pGradientFunction() override {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Unref();
+      }
+    }
+  }
+
+ private:
+  // TODO(b/174778737): Only hold needed inputs.
+  vector<AbstractTensorHandle*> forward_inputs_;
+};
+
+class DivNoNanGradientFunction : public GradientFunction {
+ public:
+  explicit DivNoNanGradientFunction(vector<AbstractTensorHandle*> f_inputs,
+                                    vector<AbstractTensorHandle*> f_outputs)
+      : forward_inputs_(f_inputs), forward_outputs_(f_outputs) {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Ref();
+      }
+    }
+    for (auto output : forward_outputs_) {
+      if (output) {
+        output->Ref();
+      }
+    }
+  }
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    // TODO(vnvo2409): Add shape broadcasting
+    /* Given upstream grad U and a Div op: Z = X/Y, the gradients are:
+     *
+     *    dX = U / Y
+     *    dY = -U*X / Y^2 = (X/Y) * -U / Y = -U*Z / Y
+     *
+     */
+
+    AbstractTensorHandle* upstream_grad = grad_outputs[0];
+    AbstractTensorHandle* Y = forward_inputs_[1];
+    AbstractTensorHandle* Z = forward_outputs_[0];
+
+    // Calculate dX =  U / Y
+    std::string name = "Div_Grad_X";
+    TF_RETURN_IF_ERROR(
+        DivNoNan(ctx, upstream_grad, Y, &grad_inputs[0], name.c_str()));
+
+    AbstractTensorHandle* temp_output;
+    // Calculate dY = -U*Z / Y
+    name = "Neg_Div_Grad_Y";
+    TF_RETURN_IF_ERROR(Neg(ctx, upstream_grad, &temp_output,
+                           name.c_str()));  // -U
+    AbstractTensorHandlePtr MinusU(temp_output);
+
+    name = "Mul_Div_Grad_Y";
+    TF_RETURN_IF_ERROR(Mul(ctx, MinusU.get(), Z, &temp_output,
+                           name.c_str()));  // -U*Z
+    AbstractTensorHandlePtr UZ(temp_output);
+
+    name = "Div_Grad_Y";
+    TF_RETURN_IF_ERROR(DivNoNan(ctx, UZ.get(), Y, &grad_inputs[1],
+                                name.c_str()));  // -U*Z / Y
+
+    return absl::OkStatus();
+  }
+  ~DivNoNanGradientFunction() override {
+    for (auto input : forward_inputs_) {
+      if (input) {
+        input->Unref();
+      }
+    }
+    for (auto output : forward_outputs_) {
+      if (output) {
+        output->Unref();
+      }
+    }
+  }
+
+ private:
+  // TODO(b/174778737): Only hold needed inputs and outputs.
+  vector<AbstractTensorHandle*> forward_inputs_;
+  vector<AbstractTensorHandle*> forward_outputs_;
+};
+
+}  // namespace
+
+GradientFunction* AddRegisterer(const ForwardOperation& op) {
+  return new AddGradientFunction;
+}
+
+GradientFunction* ExpRegisterer(const ForwardOperation& op) {
+  return new ExpGradientFunction(op.outputs[0]);
+}
+
+GradientFunction* MatMulRegisterer(const ForwardOperation& op) {
+  return new MatMulGradientFunction(op.inputs, op.attrs);
+}
+
+GradientFunction* SqrtRegisterer(const ForwardOperation& op) {
+  return new SqrtGradientFunction(op.outputs[0]);
+}
+
+GradientFunction* NegRegisterer(const ForwardOperation& op) {
+  return new NegGradientFunction;
+}
+
+GradientFunction* SubRegisterer(const ForwardOperation& op) {
+  return new SubGradientFunction;
+}
+
+GradientFunction* MulRegisterer(const ForwardOperation& op) {
+  return new MulGradientFunction(op.inputs);
+}
+
+GradientFunction* Log1pRegisterer(const ForwardOperation& op) {
+  return new Log1pGradientFunction(op.inputs);
+}
+
+GradientFunction* DivNoNanRegisterer(const ForwardOperation& op) {
+  return new DivNoNanGradientFunction(op.inputs, op.outputs);
+}
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/math_grad.h

@@ -0,0 +1,36 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_MATH_GRAD_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_MATH_GRAD_H_
+
+#include "tensorflow/c/eager/gradients.h"
+
+namespace tensorflow {
+namespace gradients {
+
+GradientFunction* AddRegisterer(const ForwardOperation& op);
+GradientFunction* ExpRegisterer(const ForwardOperation& op);
+GradientFunction* MatMulRegisterer(const ForwardOperation& op);
+GradientFunction* SqrtRegisterer(const ForwardOperation& op);
+GradientFunction* NegRegisterer(const ForwardOperation& op);
+GradientFunction* SubRegisterer(const ForwardOperation& op);
+GradientFunction* MulRegisterer(const ForwardOperation& op);
+GradientFunction* Log1pRegisterer(const ForwardOperation& op);
+GradientFunction* DivNoNanRegisterer(const ForwardOperation& op);
+
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_MATH_GRAD_H_

tensorflow/c/experimental/gradients/math_grad_test.cc

@@ -0,0 +1,461 @@
+/* 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/c/experimental/gradients/math_grad.h"
+
+#include <cstdint>
+#include <cstring>
+#include <tuple>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/experimental/gradients/grad_test_helper.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+namespace {
+
+using tensorflow::TF_StatusPtr;
+
+absl::Status AddModel(AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::AddV2(ctx, inputs[0], inputs[1], &outputs[0], "Add");
+}
+
+absl::Status ExpModel(AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Exp(ctx, inputs[0], &outputs[0], "Exp");
+}
+
+absl::Status SqrtModel(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> inputs,
+                       absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Sqrt(ctx, inputs[0], &outputs[0], "Sqrt");
+}
+
+absl::Status NegModel(AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Neg(ctx, inputs[0], &outputs[0], "Neg");
+}
+
+absl::Status SubModel(AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Sub(ctx, inputs[0], inputs[1], &outputs[0], "Sub");
+}
+
+absl::Status MulModel(AbstractContext* ctx,
+                      absl::Span<AbstractTensorHandle* const> inputs,
+                      absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Mul(ctx, inputs[0], inputs[1], &outputs[0], "Mul");
+}
+
+absl::Status Log1pModel(AbstractContext* ctx,
+                        absl::Span<AbstractTensorHandle* const> inputs,
+                        absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Log1p(ctx, inputs[0], &outputs[0], "Log1p");
+}
+
+absl::Status DivNoNanModel(AbstractContext* ctx,
+                           absl::Span<AbstractTensorHandle* const> inputs,
+                           absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::DivNoNan(ctx, inputs[0], inputs[1], &outputs[0], "DivNoNan");
+}
+
+class CppGradients
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+    status_ = StatusFromTF_Status(status.get());
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+    {
+      AbstractContext* ctx_raw = nullptr;
+      status_ =
+          BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+      ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+      immediate_execution_ctx_.reset(ctx_raw);
+    }
+
+    // Computing numerical gradients with TensorFloat-32 is numerically
+    // unstable. Some forward pass tests also fail with TensorFloat-32 due to
+    // low tolerances
+    enable_tensor_float_32_execution(false);
+  }
+
+  AbstractContextPtr immediate_execution_ctx_;
+  GradientRegistry registry_;
+  absl::Status status_;
+
+ public:
+  bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
+  bool UseFunction() const { return std::get<2>(GetParam()); }
+};
+
+TEST_P(CppGradients, TestAddGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandlePtr y;
+  {
+    AbstractTensorHandle* y_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &y_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    y.reset(y_raw);
+  }
+
+  // TODO(srbs): Rename ops::Add to ops::AddV2 and AddRegister to
+  // AddV2Registerer.
+  status_ = registry_.Register("AddV2", AddRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      AddModel, BuildGradModel(AddModel, registry_),
+      immediate_execution_ctx_.get(), {x.get(), y.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestExpGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  status_ = registry_.Register("Exp", ExpRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      ExpModel, BuildGradModel(ExpModel, registry_),
+      immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestMatMulGrad) {
+  // TODO(vnvo2409): Figure out why `gradient_checker` does not work very
+  // well with `MatMul` and remove `TestMatMul*` in
+  // `mnist_gradients_test` when done.
+  GTEST_SKIP();
+
+  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
+  int64_t A_dims[] = {3, 3};
+  AbstractTensorHandlePtr A;
+  {
+    AbstractTensorHandle* A_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), A_vals, A_dims, 2, &A_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    A.reset(A_raw);
+  }
+
+  float B_vals[] = {9.0f, 8.0f, 7.0f, 6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
+  int64_t B_dims[] = {3, 3};
+  AbstractTensorHandlePtr B;
+  {
+    AbstractTensorHandle* B_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), B_vals, B_dims, 2, &B_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    B.reset(B_raw);
+  }
+
+  status_ = registry_.Register("MatMul", MatMulRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  for (bool transpose_a : {false, true}) {
+    for (bool transpose_b : {false, true}) {
+      Model MatMulModel =
+          [transpose_a, transpose_b](
+              AbstractContext* ctx,
+              absl::Span<AbstractTensorHandle* const> inputs,
+              absl::Span<AbstractTensorHandle*> outputs) -> absl::Status {
+        return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0], transpose_a,
+                           transpose_b, /*grad_a=*/false, /*grad_b=*/false,
+                           "MatMul");
+      };
+      ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+          MatMulModel, BuildGradModel(MatMulModel, registry_),
+          immediate_execution_ctx_.get(), {A.get(), B.get()}, UseFunction()));
+    }
+  }
+}
+
+TEST_P(CppGradients, TestMatMulGradManual) {
+  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
+  int64_t A_dims[] = {3, 3};
+  AbstractTensorHandlePtr A;
+  {
+    AbstractTensorHandle* A_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), A_vals, A_dims, 2, &A_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    A.reset(A_raw);
+  }
+
+  float B_vals[] = {9.0f, 8.0f, 7.0f, 6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
+  int64_t B_dims[] = {3, 3};
+  AbstractTensorHandlePtr B;
+  {
+    AbstractTensorHandle* B_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), B_vals, B_dims, 2, &B_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    B.reset(B_raw);
+  }
+
+  status_ = registry_.Register("MatMul", MatMulRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  bool transpose_a_vals[] = {false, false, true, true};
+  bool transpose_b_vals[] = {false, true, false, true};
+  float dA_vals[4][9] = {{24, 15, 6, 24, 15, 6, 24, 15, 6},
+                         {18, 15, 12, 18, 15, 12, 18, 15, 12},
+                         {24, 24, 24, 15, 15, 15, 6, 6, 6},
+                         {18, 18, 18, 15, 15, 15, 12, 12, 12}};
+  float dB_vals[4][9] = {{12, 12, 12, 15, 15, 15, 18, 18, 18},
+                         {12, 15, 18, 12, 15, 18, 12, 15, 18},
+                         {6, 6, 6, 15, 15, 15, 24, 24, 24},
+                         {6, 15, 24, 6, 15, 24, 6, 15, 24}};
+
+  for (int i{}; i < 4; ++i) {
+    bool transpose_a = transpose_a_vals[i];
+    bool transpose_b = transpose_b_vals[i];
+    Model MatMulModel =
+        [transpose_a, transpose_b](
+            AbstractContext* ctx,
+            absl::Span<AbstractTensorHandle* const> inputs,
+            absl::Span<AbstractTensorHandle*> outputs) -> absl::Status {
+      return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0], transpose_a,
+                         transpose_b, /*grad_a=*/false, /*grad_b=*/false,
+                         "MatMul");
+    };
+    Model MatMulGradModel = BuildGradModel(MatMulModel, registry_);
+    std::vector<AbstractTensorHandle*> outputs(2);
+    status_ =
+        RunModel(MatMulGradModel, immediate_execution_ctx_.get(),
+                 {A.get(), B.get()}, absl::MakeSpan(outputs), UseFunction());
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[0], dA_vals[i],
+                                             /*dims*/ {3, 3},
+                                             /*abs_error*/ 0));
+    ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[1], dB_vals[i],
+                                             /*dims*/ {3, 3},
+                                             /*abs_error*/ 0));
+    outputs[0]->Unref();
+    outputs[1]->Unref();
+  }
+}
+
+TEST_P(CppGradients, TestSqrtGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  status_ = registry_.Register("Sqrt", SqrtRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      SqrtModel, BuildGradModel(SqrtModel, registry_),
+      immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestNegGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  status_ = registry_.Register("Neg", NegRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      NegModel, BuildGradModel(NegModel, registry_),
+      immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestSubGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandlePtr y;
+  {
+    AbstractTensorHandle* y_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &y_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    y.reset(y_raw);
+  }
+
+  status_ = registry_.Register("Sub", SubRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      SubModel, BuildGradModel(SubModel, registry_),
+      immediate_execution_ctx_.get(), {x.get(), y.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestMulGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandlePtr y;
+  {
+    AbstractTensorHandle* y_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &y_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    y.reset(y_raw);
+  }
+
+  status_ = registry_.Register("Mul", MulRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      MulModel, BuildGradModel(MulModel, registry_),
+      immediate_execution_ctx_.get(), {x.get(), y.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestLog1pGrad) {
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  status_ = registry_.Register("Log1p", Log1pRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      Log1pModel, BuildGradModel(Log1pModel, registry_),
+      immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestDivNoNanGrad) {
+  status_ = registry_.Register("DivNoNan", DivNoNanRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  auto DivNoNanGradModel = BuildGradModel(DivNoNanModel, registry_);
+
+  AbstractTensorHandlePtr x;
+  {
+    AbstractTensorHandle* x_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &x_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    x.reset(x_raw);
+  }
+
+  AbstractTensorHandlePtr y;
+  {
+    AbstractTensorHandle* y_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 2.0f, &y_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    y.reset(y_raw);
+  }
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      DivNoNanModel, DivNoNanGradModel, immediate_execution_ctx_.get(),
+      {x.get(), y.get()}, UseFunction()));
+
+  // `DivNoNanGradModel` should return {`0`, `0`} when the denominator is `0`.
+  AbstractTensorHandlePtr z;
+  {
+    AbstractTensorHandle* z_raw = nullptr;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 0.0f, &z_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    z.reset(z_raw);
+  }
+  std::vector<AbstractTensorHandle*> outputs(2);
+  status_ =
+      RunModel(DivNoNanGradModel, immediate_execution_ctx_.get(),
+               {x.get(), z.get()}, absl::MakeSpan(outputs), UseFunction());
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+  ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[0], {0.0f}, /*dims*/ {},
+                                           /*abs_error*/ 0));
+  ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[1], {0.0f}, /*dims*/ {},
+                                           /*abs_error*/ 0));
+  outputs[0]->Unref();
+  outputs[1]->Unref();
+}
+
+#ifdef PLATFORM_GOOGLE
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#else
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#endif
+}  // namespace
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/nn_grad.cc

@@ -0,0 +1,171 @@
+/* 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/c/experimental/gradients/nn_grad.h"
+
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/experimental/ops/array_ops.h"
+#include "tensorflow/c/experimental/ops/math_ops.h"
+#include "tensorflow/c/experimental/ops/nn_ops.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+
+using std::vector;
+using tensorflow::ops::BiasAddGrad;
+using tensorflow::ops::ReluGrad;
+
+namespace tensorflow {
+namespace gradients {
+namespace {
+
+class ReluGradientFunction : public GradientFunction {
+ public:
+  explicit ReluGradientFunction(vector<AbstractTensorHandle*> f_outputs)
+      : forward_outputs_(f_outputs) {
+    for (auto output : forward_outputs_) {
+      if (output) {
+        output->Ref();
+      }
+    }
+  }
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    AbstractTensorHandle* upstream_grad = grad_outputs[0];
+    AbstractTensorHandle* activations = forward_outputs_[0];
+
+    // Calculate Grad
+    std::string name = "relu_grad";
+    TF_RETURN_IF_ERROR(ReluGrad(ctx, upstream_grad, activations,
+                                &grad_inputs[0], name.c_str()));
+    return absl::OkStatus();
+  }
+  ~ReluGradientFunction() override {
+    for (auto output : forward_outputs_) {
+      if (output) {
+        output->Unref();
+      }
+    }
+  }
+
+ private:
+  // TODO(b/174778737): Only hold needed outputs.
+  vector<AbstractTensorHandle*> forward_outputs_;
+};
+
+absl::Status BroadcastMul(AbstractContext* ctx, AbstractTensorHandle* vec,
+                          AbstractTensorHandle* mat,
+                          absl::Span<AbstractTensorHandle*> outputs) {
+  if (!isa<ImmediateExecutionContext>(ctx)) {
+    // TODO(b/168850692): Fix this.
+    return absl::UnimplementedError(
+        "BroadcastMul is not supported in tracing mode yet.");
+  }
+  auto imm_ctx = dyn_cast<ImmediateExecutionContext>(ctx);
+  AbstractTensorPtr minus_1(imm_ctx->CreateInt32Scalar(-1));
+  ImmediateTensorHandlePtr dim(imm_ctx->CreateLocalHandle(minus_1.get()));
+  AbstractTensorHandle* expand_dims_outputs;
+  TF_RETURN_IF_ERROR(
+      ops::ExpandDims(ctx, vec, dim.get(), &expand_dims_outputs, "ExpandDims"));
+  TF_RETURN_IF_ERROR(
+      ops::Mul(ctx, expand_dims_outputs, mat, &outputs[0], "Mul"));
+  expand_dims_outputs->Unref();
+  return absl::OkStatus();
+}
+
+class SparseSoftmaxCrossEntropyWithLogitsGradientFunction
+    : public GradientFunction {
+ public:
+  explicit SparseSoftmaxCrossEntropyWithLogitsGradientFunction(
+      vector<AbstractTensorHandle*> f_outputs)
+      : forward_outputs_(f_outputs) {}
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    // Grad for Softmax Input
+    TF_RETURN_IF_ERROR(BroadcastMul(
+        ctx, grad_outputs[0], forward_outputs_[1],
+        grad_inputs.subspan(0, 1)));  // upstream_grad * local softmax grad
+
+    // Grad for labels is null
+    grad_inputs[1] = nullptr;
+    return absl::OkStatus();
+  }
+  ~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override {}
+
+ private:
+  vector<AbstractTensorHandle*> forward_outputs_;
+};
+
+// TODO(vnvo2409): Add python test
+class BiasAddGradientFunction : public GradientFunction {
+ public:
+  explicit BiasAddGradientFunction(AttrBuilder f_attrs)
+      : forward_attrs_(f_attrs) {}
+
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override {
+    /* Given upstream grad U and a BiasAdd: A + bias, the gradients are:
+     *
+     *    dA = U
+     *    dbias = reduceSum(U, dims = channel_dim)
+     */
+
+    AbstractTensorHandle* upstream_grad = grad_outputs[0];
+    DCHECK(upstream_grad);
+
+    // Recover data format from forward pass for gradient.
+    std::string data_format;
+    TF_RETURN_IF_ERROR(forward_attrs_.Get("data_format", &data_format));
+
+    // Grad for A
+    grad_inputs[0] = upstream_grad;
+    grad_inputs[0]->Ref();
+
+    // Grad for bias
+    std::string name = "bias_add_grad";
+    TF_RETURN_IF_ERROR(BiasAddGrad(ctx, upstream_grad, &grad_inputs[1],
+                                   data_format.c_str(), name.c_str()));
+
+    return absl::OkStatus();
+  }
+  ~BiasAddGradientFunction() override {}
+
+ private:
+  AttrBuilder forward_attrs_;
+};
+
+}  // namespace
+
+GradientFunction* ReluRegisterer(const ForwardOperation& op) {
+  return new ReluGradientFunction(op.outputs);
+}
+
+GradientFunction* SparseSoftmaxCrossEntropyWithLogitsRegisterer(
+    const ForwardOperation& op) {
+  return new SparseSoftmaxCrossEntropyWithLogitsGradientFunction(op.outputs);
+}
+
+GradientFunction* BiasAddRegisterer(const ForwardOperation& op) {
+  return new BiasAddGradientFunction(op.attrs);
+}
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/nn_grad.h

@@ -0,0 +1,29 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_NN_GRAD_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NN_GRAD_H_
+
+#include "tensorflow/c/eager/gradients.h"
+
+namespace tensorflow {
+namespace gradients {
+GradientFunction* ReluRegisterer(const ForwardOperation& op);
+GradientFunction* SparseSoftmaxCrossEntropyWithLogitsRegisterer(
+    const ForwardOperation& op);
+GradientFunction* BiasAddRegisterer(const ForwardOperation& op);
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NN_GRAD_H_

tensorflow/c/experimental/gradients/nn_grad_test.cc

@@ -0,0 +1,228 @@
+/* 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/c/experimental/gradients/nn_grad.h"
+
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/eager/unified_api_testutil.h"
+#include "tensorflow/c/experimental/gradients/grad_test_helper.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
+#include "tensorflow/c/experimental/ops/nn_ops.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace gradients {
+namespace internal {
+namespace {
+
+using tensorflow::TF_StatusPtr;
+
+absl::Status ReluModel(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> inputs,
+                       absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::Relu(ctx, inputs[0], &outputs[0], "Relu");
+}
+
+absl::Status SparseSoftmaxCrossEntropyWithLogitsModel(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
+    absl::Span<AbstractTensorHandle*> outputs) {
+  AbstractTensorHandle* loss;
+  AbstractTensorHandle* backprop;
+  TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
+      ctx, inputs[0], inputs[1], &loss, &backprop,
+      "SparseSoftmaxCrossEntropyWithLogits"));
+  // `gradient_checker` only works with model that returns only 1 tensor.
+  // Although, `ops::SparseSoftmaxCrossEntropyWithLogits` returns 2 tensors, the
+  // second tensor isn't needed for computing gradient so we could safely drop
+  // it.
+  outputs[0] = loss;
+  backprop->Unref();
+  return absl::OkStatus();
+}
+
+absl::Status BiasAddModel(AbstractContext* ctx,
+                          absl::Span<AbstractTensorHandle* const> inputs,
+                          absl::Span<AbstractTensorHandle*> outputs) {
+  return ops::BiasAdd(ctx, inputs[0], inputs[1], &outputs[0], "NHWC",
+                      "BiasAdd");
+}
+
+class CppGradients
+    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
+ protected:
+  void SetUp() override {
+    TF_StatusPtr status(TF_NewStatus());
+    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
+    status_ = StatusFromTF_Status(status.get());
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+    {
+      AbstractContext* ctx_raw = nullptr;
+      status_ =
+          BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+      ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+      immediate_execution_ctx_.reset(ctx_raw);
+    }
+
+    // Computing numerical gradients with TensorFloat-32 is numerically
+    // unstable. Some forward pass tests also fail with TensorFloat-32 due to
+    // low tolerances
+    enable_tensor_float_32_execution(false);
+  }
+
+  AbstractContextPtr immediate_execution_ctx_;
+  GradientRegistry registry_;
+  absl::Status status_;
+
+ public:
+  bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
+  bool UseFunction() const { return std::get<2>(GetParam()); }
+};
+
+TEST_P(CppGradients, TestReluGrad) {
+  status_ = registry_.Register("Relu", ReluRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  auto ReluGradModel = BuildGradModel(ReluModel, registry_);
+
+  float X_vals[] = {1.0f, 2.0f, 3.0f, -5.0f, -4.0f, -3.0f, 2.0f, 10.0f, -1.0f};
+  int64_t X_dims[] = {3, 3};
+  AbstractTensorHandlePtr X;
+  {
+    AbstractTensorHandle* X_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    X.reset(X_raw);
+  }
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      ReluModel, ReluGradModel, immediate_execution_ctx_.get(), {X.get()},
+      UseFunction()));
+
+  // Mathematically, Relu isn't differentiable at `0`. So `gradient_checker`
+  // does not work with it.
+  AbstractTensorHandlePtr Y;
+  {
+    AbstractTensorHandle* Y_raw;
+    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), 0.0f, &Y_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    Y.reset(Y_raw);
+  }
+
+  std::vector<AbstractTensorHandle*> outputs(1);
+  status_ = RunModel(ReluGradModel, immediate_execution_ctx_.get(), {Y.get()},
+                     absl::MakeSpan(outputs), UseFunction());
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+  ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[0], {0.0f}, /*dims*/ {},
+                                           /*abs_error*/ 0));
+  outputs[0]->Unref();
+}
+
+TEST_P(CppGradients, TestSparseSoftmaxCrossEntropyWithLogitsGrad) {
+  if (UseFunction()) {
+    // TODO(b/168850692): Enable this.
+    GTEST_SKIP() << "Can't take gradient of "
+                    "SparseSoftmaxCrossEntropyWithLogits in tracing mode.";
+  }
+
+  // Score
+  float X_vals[] = {1.0f, 2.0f, 3.0f, -5.0f, -4.0f, -3.0f, 2.0f, 0.0f, -1.0f};
+  int64_t X_dims[] = {3, 3};
+  AbstractTensorHandlePtr X;
+  {
+    AbstractTensorHandle* X_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    X.reset(X_raw);
+  }
+  // Label
+  int32_t Y_vals[] = {1, 0, 1};
+  int64_t Y_dims[] = {3};
+  AbstractTensorHandlePtr Y;
+  {
+    AbstractTensorHandle* Y_raw;
+    status_ = TestTensorHandleWithDims<int32_t, TF_INT32>(
+        immediate_execution_ctx_.get(), Y_vals, Y_dims, 1, &Y_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    Y.reset(Y_raw);
+  }
+
+  status_ = registry_.Register("SparseSoftmaxCrossEntropyWithLogits",
+                               SparseSoftmaxCrossEntropyWithLogitsRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      SparseSoftmaxCrossEntropyWithLogitsModel,
+      BuildGradModel(SparseSoftmaxCrossEntropyWithLogitsModel, registry_),
+      immediate_execution_ctx_.get(), {X.get(), Y.get()}, UseFunction()));
+}
+
+TEST_P(CppGradients, TestBiasAddGrad) {
+  if (UseFunction() && UseMlir()) {
+    GTEST_SKIP() << "SetAttrString has not been implemented yet.\n";
+  }
+
+  // A
+  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f};
+  int64_t A_dims[] = {2, 2};
+  AbstractTensorHandlePtr A;
+  {
+    AbstractTensorHandle* A_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), A_vals, A_dims, 2, &A_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    A.reset(A_raw);
+  }
+  // Bias
+  float Bias_vals[] = {2.0f, 3.0f};
+  int64_t Bias_dims[] = {2};
+  AbstractTensorHandlePtr Bias;
+  {
+    AbstractTensorHandle* Bias_raw;
+    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
+        immediate_execution_ctx_.get(), Bias_vals, Bias_dims, 1, &Bias_raw);
+    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+    Bias.reset(Bias_raw);
+  }
+
+  status_ = registry_.Register("BiasAdd", BiasAddRegisterer);
+  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
+
+  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
+      BiasAddModel, BuildGradModel(BiasAddModel, registry_),
+      immediate_execution_ctx_.get(), {A.get(), Bias.get()}, UseFunction()));
+}
+
+#ifdef PLATFORM_GOOGLE
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#else
+INSTANTIATE_TEST_SUITE_P(
+    UnifiedCAPI, CppGradients,
+    ::testing::Combine(::testing::Values("graphdef", "mlir"),
+                       /*tfrt*/ ::testing::Values(false),
+                       /*use_function*/ ::testing::Values(true, false)));
+#endif
+}  // namespace
+}  // namespace internal
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/not_differentiable.cc

@@ -0,0 +1,35 @@
+/* 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/c/experimental/gradients/not_differentiable.h"
+
+namespace tensorflow {
+namespace gradients {
+absl::Status NotDifferentiableGradientFunction::Compute(
+    AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> grad_outputs,
+    absl::Span<AbstractTensorHandle*> grad_inputs) {
+  for (int i = 0; i < grad_inputs.size(); i++) {
+    grad_inputs[i] = nullptr;
+  }
+  return absl::OkStatus();
+}
+
+absl::Status RegisterNotDifferentiable(GradientRegistry* registry,
+                                       const std::string& op) {
+  return registry->Register(op, [](const ForwardOperation& op) {
+    return new NotDifferentiableGradientFunction;
+  });
+}
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/not_differentiable.h

@@ -0,0 +1,35 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_NOT_DIFFERENTIABLE_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NOT_DIFFERENTIABLE_H_
+
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/gradients.h"
+
+namespace tensorflow {
+namespace gradients {
+// Ignores `grad_outputs` and sets all entries in grad_inputs to nullptr.
+class NotDifferentiableGradientFunction : public GradientFunction {
+  absl::Status Compute(AbstractContext* ctx,
+                       absl::Span<AbstractTensorHandle* const> grad_outputs,
+                       absl::Span<AbstractTensorHandle*> grad_inputs) override;
+};
+// Shorthand for registry->Register(op, new NotDifferentiableGradientFunction)
+absl::Status RegisterNotDifferentiable(GradientRegistry* registry,
+                                       const std::string& op);
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NOT_DIFFERENTIABLE_H_

tensorflow/c/experimental/gradients/tape/BUILD

@@ -0,0 +1,93 @@
+# A tape built on top of unified execution APIs.
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "tape_context",
+    srcs = ["tape_context.cc"],
+    hdrs = [
+        "tape_context.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":tape_operation",
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:abstract_function",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+cc_library(
+    name = "tape_operation",
+    srcs = ["tape_operation.cc"],
+    hdrs = [
+        "tape_operation.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        "//tensorflow/c:tensor_interface",
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:abstract_operation",
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:stringpiece",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "tape",
+    hdrs = [
+        "tape_context.h",
+        "tape_operation.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":tape_context",
+        ":tape_operation",
+        "//tensorflow/c:tensor_interface",
+        "//tensorflow/c/eager:abstract_context",
+        "//tensorflow/c/eager:abstract_function",
+        "//tensorflow/c/eager:abstract_operation",
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:gradients_internal",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "tape_context.h",
+        "tape_operation.h",
+    ],
+    visibility = [
+        "//tensorflow:internal",
+    ],
+)

tensorflow/c/experimental/gradients/tape/tape_context.cc

@@ -0,0 +1,52 @@
+/* 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/c/experimental/gradients/tape/tape_context.h"
+
+#include "absl/status/status.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_operation.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace gradients {
+TapeContext::TapeContext(AbstractContext* c, Tape* tape,
+                         const GradientRegistry& registry)
+    : AbstractContext(kTape), parent_ctx_(c), tape_(tape), registry_(registry) {
+  // TODO(srbs): Make AbstractContext ref counted.
+  // parent_ctx_->Ref();
+}
+void TapeContext::Release() {
+  // TODO(srbs): Change to Unref()
+  delete this;
+}
+TapeContext::~TapeContext() {
+  // TODO(srbs): Make AbstractContext ref counted.
+  // parent_ctx_->Unref();
+}
+TapeOperation* TapeContext::CreateOperation() {
+  return new TapeOperation(parent_ctx_->CreateOperation(), tape_, registry_);
+}
+absl::Status TapeContext::RegisterFunction(AbstractFunction* f) {
+  return parent_ctx_->RegisterFunction(f);
+}
+absl::Status TapeContext::RemoveFunction(const std::string& func) {
+  return parent_ctx_->RemoveFunction(func);
+}
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/tape/tape_context.h

@@ -0,0 +1,49 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_CONTEXT_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_CONTEXT_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_operation.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace gradients {
+class TapeContext : public AbstractContext {
+ public:
+  explicit TapeContext(AbstractContext*, Tape*, const GradientRegistry&);
+  void Release() override;
+  TapeOperation* CreateOperation() override;
+  absl::Status RegisterFunction(AbstractFunction*) override;
+  absl::Status RemoveFunction(const std::string& func) override;
+  // For LLVM style RTTI.
+  static bool classof(const AbstractContext* ptr) {
+    return ptr->getKind() == kTape;
+  }
+  ~TapeContext() override;
+
+ private:
+  AbstractContext* parent_ctx_;  // Not owned.
+  Tape* tape_;
+  const GradientRegistry& registry_;
+};
+}  // namespace gradients
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_CONTEXT_H_

tensorflow/c/experimental/gradients/tape/tape_operation.cc

@@ -0,0 +1,246 @@
+/* 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/c/experimental/gradients/tape/tape_operation.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace gradients {
+TapeOperation::TapeOperation(AbstractOperation* parent_op, Tape* tape,
+                             const GradientRegistry& registry)
+    : AbstractOperation(kTape),
+      parent_op_(parent_op),
+      tape_(tape),
+      registry_(registry) {
+  // TODO(b/172003047): Consider making AbstractOperation RefCounted.
+  // parent_op_->Ref();
+}
+void TapeOperation::Release() {
+  // TODO(srbs): Change to Unref().
+  delete this;
+}
+TapeOperation::~TapeOperation() {
+  // TODO(b/172003047): Consider making AbstractOperation RefCounted.
+  // parent_op->Unref();
+}
+absl::Status TapeOperation::Reset(const char* op, const char* raw_device_name) {
+  forward_op_.op_name = op;
+  forward_op_.attrs.Reset(op);
+  forward_op_.inputs.clear();
+  forward_op_.outputs.clear();
+  return parent_op_->Reset(op, raw_device_name);
+}
+const std::string& TapeOperation::Name() const { return parent_op_->Name(); }
+const std::string& TapeOperation::DeviceName() const {
+  return parent_op_->DeviceName();
+}
+absl::Status TapeOperation::SetDeviceName(const char* name) {
+  return parent_op_->SetDeviceName(name);
+}
+absl::Status TapeOperation::AddInput(AbstractTensorHandle* input) {
+  TF_RETURN_IF_ERROR(parent_op_->AddInput(input));
+  forward_op_.inputs.push_back(input);
+  return absl::OkStatus();
+}
+absl::Status TapeOperation::AddInputList(
+    absl::Span<AbstractTensorHandle* const> inputs) {
+  TF_RETURN_IF_ERROR(parent_op_->AddInputList(inputs));
+  for (auto input : inputs) {
+    forward_op_.inputs.push_back(input);
+  }
+  return absl::OkStatus();
+}
+absl::Status TapeOperation::SetAttrString(const char* attr_name,
+                                          const char* data, size_t length) {
+  forward_op_.attrs.Set(attr_name, absl::string_view(data, length));
+  return parent_op_->SetAttrString(attr_name, data, length);
+}
+absl::Status TapeOperation::SetAttrInt(const char* attr_name, int64_t value) {
+  forward_op_.attrs.Set(attr_name, static_cast<int64_t>(value));
+  return parent_op_->SetAttrInt(attr_name, value);
+}
+absl::Status TapeOperation::SetAttrFloat(const char* attr_name, float value) {
+  forward_op_.attrs.Set(attr_name, value);
+  return parent_op_->SetAttrFloat(attr_name, value);
+}
+absl::Status TapeOperation::SetAttrBool(const char* attr_name, bool value) {
+  forward_op_.attrs.Set(attr_name, value);
+  return parent_op_->SetAttrBool(attr_name, value);
+}
+absl::Status TapeOperation::SetAttrType(const char* attr_name, DataType value) {
+  forward_op_.attrs.Set(attr_name, value);
+  return parent_op_->SetAttrType(attr_name, value);
+}
+absl::Status TapeOperation::SetAttrShape(const char* attr_name,
+                                         const int64_t* dims,
+                                         const int num_dims) {
+  if (num_dims > TensorShape::MaxDimensions()) {
+    return errors::InvalidArgument("Value specified for `", attr_name, "` has ",
+                                   num_dims,
+                                   " dimensions which is over the limit of ",
+                                   TensorShape::MaxDimensions(), ".");
+  }
+  TensorShapeProto proto;
+  if (num_dims < 0) {
+    proto.set_unknown_rank(true);
+  } else {
+    for (int d = 0; d < num_dims; ++d) {
+      proto.add_dim()->set_size(dims[d]);
+    }
+  }
+
+  forward_op_.attrs.Set(attr_name, proto);
+  return parent_op_->SetAttrShape(attr_name, dims, num_dims);
+}
+absl::Status TapeOperation::SetAttrFunction(const char* attr_name,
+                                            const AbstractOperation* value) {
+  return tensorflow::errors::Unimplemented(
+      "SetAttrFunction has not been implemented yet.");
+}
+absl::Status TapeOperation::SetAttrFunctionName(const char* attr_name,
+                                                const char* value,
+                                                size_t length) {
+  return tensorflow::errors::Unimplemented(
+      "SetAttrFunctionName has not been implemented "
+      "yet.");
+}
+absl::Status TapeOperation::SetAttrTensor(const char* attr_name,
+                                          AbstractTensorInterface* tensor) {
+  return tensorflow::errors::Unimplemented(
+      "SetAttrTensor has not been implemented yet.");
+}
+absl::Status TapeOperation::SetAttrStringList(const char* attr_name,
+                                              const void* const* values,
+                                              const size_t* lengths,
+                                              int num_values) {
+  std::vector<absl::string_view> v(num_values);
+  for (int i = 0; i < num_values; ++i) {
+    v[i] = absl::string_view(static_cast<const char*>(values[i]), lengths[i]);
+  }
+  forward_op_.attrs.Set(attr_name, v);
+  return parent_op_->SetAttrStringList(attr_name, values, lengths, num_values);
+}
+absl::Status TapeOperation::SetAttrFloatList(const char* attr_name,
+                                             const float* values,
+                                             int num_values) {
+  forward_op_.attrs.Set(attr_name,
+                        gtl::ArraySlice<const float>(values, num_values));
+  return parent_op_->SetAttrFloatList(attr_name, values, num_values);
+}
+absl::Status TapeOperation::SetAttrIntList(const char* attr_name,
+                                           const int64_t* values,
+                                           int num_values) {
+  forward_op_.attrs.Set(
+      attr_name, gtl::ArraySlice<const int64_t>(
+                     reinterpret_cast<const int64_t*>(values), num_values));
+  return parent_op_->SetAttrIntList(attr_name, values, num_values);
+}
+absl::Status TapeOperation::SetAttrTypeList(const char* attr_name,
+                                            const DataType* values,
+                                            int num_values) {
+  forward_op_.attrs.Set(attr_name,
+                        gtl::ArraySlice<const DataType>(values, num_values));
+  return parent_op_->SetAttrTypeList(attr_name, values, num_values);
+}
+absl::Status TapeOperation::SetAttrBoolList(const char* attr_name,
+                                            const unsigned char* values,
+                                            int num_values) {
+  std::unique_ptr<bool[]> b(new bool[num_values]);
+  for (int i = 0; i < num_values; ++i) {
+    b[i] = values[i];
+  }
+  forward_op_.attrs.Set(attr_name,
+                        gtl::ArraySlice<const bool>(b.get(), num_values));
+  return parent_op_->SetAttrBoolList(attr_name, values, num_values);
+}
+absl::Status TapeOperation::SetAttrShapeList(const char* attr_name,
+                                             const int64_t** dims,
+                                             const int* num_dims,
+                                             int num_values) {
+  std::unique_ptr<TensorShapeProto[]> proto(new TensorShapeProto[num_values]);
+  for (int i = 0; i < num_values; ++i) {
+    const auto num_dims_i = num_dims[i];
+
+    if (num_dims_i > TensorShape::MaxDimensions()) {
+      return errors::InvalidArgument(
+          strings::StrCat("Value specified for `", attr_name, "` has ",
+                          num_dims_i, " dimensions which is over the limit of ",
+                          TensorShape::MaxDimensions(), "."));
+    }
+    if (num_dims_i < 0) {
+      proto[i].set_unknown_rank(true);
+    } else {
+      const int64_t* dims_i = dims[i];
+      auto proto_i = &proto[i];
+      for (int d = 0; d < num_dims_i; ++d) {
+        proto_i->add_dim()->set_size(dims_i[d]);
+      }
+    }
+  }
+  forward_op_.attrs.Set(
+      attr_name, absl::Span<const TensorShapeProto>(proto.get(), num_values));
+  return parent_op_->SetAttrShapeList(attr_name, dims, num_dims, num_values);
+}
+absl::Status TapeOperation::SetAttrFunctionList(
+    const char* attr_name, absl::Span<const AbstractOperation*> values) {
+  return tensorflow::errors::Unimplemented(
+      "SetAttrFunctionList has not been "
+      "implemented yet.");
+}
+AbstractOperation* TapeOperation::GetBackingOperation() { return parent_op_; }
+absl::Status TapeOperation::Execute(absl::Span<AbstractTensorHandle*> retvals,
+                                    int* num_retvals) {
+  TF_RETURN_IF_ERROR(parent_op_->Execute(retvals, num_retvals));
+  for (int i = 0; i < *num_retvals; i++) {
+    // TODO(srbs): Manage refcount of ForwardOperation's inputs/outputs.
+    forward_op_.outputs.push_back(retvals[i]);
+  }
+  // TODO(b/166669239): This is needed to support AttrBuilder::Get for string
+  // attributes. Number type attrs and DataType attrs work fine without this.
+  // Consider getting rid of this and making the behavior between number types
+  // and string consistent.
+  forward_op_.attrs.BuildNodeDef();
+  // TODO(b/170307493): Populate skip_input_indices here.
+  std::unique_ptr<GradientFunction> backward_fn;
+  TF_RETURN_IF_ERROR(registry_.Lookup(forward_op_, &backward_fn));
+  tape_->RecordOperation(forward_op_.inputs, forward_op_.outputs,
+                         backward_fn.release(), parent_op_->Name());
+  return absl::OkStatus();
+}
+
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/c/experimental/gradients/tape/tape_operation.h

@@ -0,0 +1,94 @@
+/* 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_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_OPERATION_H_
+#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_OPERATION_H_
+
+#include <cstddef>
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace gradients {
+class TapeOperation : public AbstractOperation {
+ public:
+  explicit TapeOperation(AbstractOperation*, Tape*, const GradientRegistry&);
+  void Release() override;
+  absl::Status Reset(const char* op, const char* raw_device_name) override;
+  const std::string& Name() const override;
+  const std::string& DeviceName() const override;
+  absl::Status SetDeviceName(const char* name) override;
+  absl::Status AddInput(AbstractTensorHandle* input) override;
+  absl::Status AddInputList(
+      absl::Span<AbstractTensorHandle* const> inputs) override;
+  absl::Status Execute(absl::Span<AbstractTensorHandle*> retvals,
+                       int* num_retvals) override;
+  absl::Status SetAttrString(const char* attr_name, const char* data,
+                             size_t length) override;
+  absl::Status SetAttrInt(const char* attr_name, int64_t value) override;
+  absl::Status SetAttrFloat(const char* attr_name, float value) override;
+  absl::Status SetAttrBool(const char* attr_name, bool value) override;
+  absl::Status SetAttrType(const char* attr_name, DataType value) override;
+  absl::Status SetAttrShape(const char* attr_name, const int64_t* dims,
+                            const int num_dims) override;
+  absl::Status SetAttrFunction(const char* attr_name,
+                               const AbstractOperation* value) override;
+  absl::Status SetAttrFunctionName(const char* attr_name, const char* value,
+                                   size_t length) override;
+  absl::Status SetAttrTensor(const char* attr_name,
+                             AbstractTensorInterface* tensor) override;
+  absl::Status SetAttrStringList(const char* attr_name,
+                                 const void* const* values,
+                                 const size_t* lengths,
+                                 int num_values) override;
+  absl::Status SetAttrFloatList(const char* attr_name, const float* values,
+                                int num_values) override;
+  absl::Status SetAttrIntList(const char* attr_name, const int64_t* values,
+                              int num_values) override;
+  absl::Status SetAttrTypeList(const char* attr_name, const DataType* values,
+                               int num_values) override;
+  absl::Status SetAttrBoolList(const char* attr_name,
+                               const unsigned char* values,
+                               int num_values) override;
+  absl::Status SetAttrShapeList(const char* attr_name, const int64_t** dims,
+                                const int* num_dims, int num_values) override;
+  absl::Status SetAttrFunctionList(
+      const char* attr_name,
+      absl::Span<const AbstractOperation*> values) override;
+  AbstractOperation* GetBackingOperation();
+  // For LLVM style RTTI.
+  static bool classof(const AbstractOperation* ptr) {
+    return ptr->getKind() == kTape;
+  }
+  ~TapeOperation() override;
+
+ private:
+  AbstractOperation* parent_op_;
+  ForwardOperation forward_op_;
+  Tape* tape_;
+  const GradientRegistry& registry_;
+};
+
+}  // namespace gradients
+}  // namespace tensorflow
+#endif  // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_OPERATION_H_