1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include <memory>
16
17 #include "tensorflow/c/eager/abstract_context.h"
18 #include "tensorflow/c/eager/c_api.h"
19 #include "tensorflow/c/eager/c_api_unified_experimental.h"
20 #include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
21 #include "tensorflow/c/eager/gradients.h"
22 #include "tensorflow/c/eager/unified_api_testutil.h"
23 #include "tensorflow/c/experimental/ops/math_ops.h"
24 #include "tensorflow/c/tf_status_helper.h"
25 #include "tensorflow/core/platform/errors.h"
26 #include "tensorflow/core/platform/test.h"
27
28 namespace tensorflow {
29 namespace gradients {
30 namespace internal {
31 namespace {
32 using std::vector;
33
34 class CustomGradientTest
35 : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
36 protected:
SetUp()37 void SetUp() override {
38 TF_StatusPtr status(TF_NewStatus());
39 TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
40 Status s = StatusFromTF_Status(status.get());
41 CHECK_EQ(errors::OK, s.code()) << s.error_message();
42 }
43 };
44
45 class PassThroughGradientFunction : public GradientFunction {
46 public:
Compute(AbstractContext * ctx,absl::Span<AbstractTensorHandle * const> grad_outputs,absl::Span<AbstractTensorHandle * > grad_inputs)47 Status Compute(AbstractContext* ctx,
48 absl::Span<AbstractTensorHandle* const> grad_outputs,
49 absl::Span<AbstractTensorHandle*> grad_inputs) override {
50 CHECK_EQ(grad_outputs.size(), 1);
51 CHECK_EQ(grad_inputs.size(), 1);
52 grad_inputs[0] = grad_outputs[0];
53 if (grad_inputs[0]) {
54 grad_inputs[0]->Ref();
55 }
56 return OkStatus();
57 }
58 };
59
60 // Computes:
61 //
62 // @tf.custom_gradient
63 // def f(input):
64 // def grad(grads):
65 // return grads[0]
66 // return tf.exp(input), grad
67 // outputs = [f(inputs[0])]
ExpWithPassThroughGrad(AbstractContext * ctx,absl::Span<AbstractTensorHandle * const> inputs,absl::Span<AbstractTensorHandle * > outputs)68 Status ExpWithPassThroughGrad(AbstractContext* ctx,
69 absl::Span<AbstractTensorHandle* const> inputs,
70 absl::Span<AbstractTensorHandle*> outputs) {
71 Tape tape(/*persistent=*/false);
72 tape.Watch(inputs[0]); // Watch x.
73 AbstractTensorHandle* exp_output;
74 TF_RETURN_IF_ERROR(ops::Exp(ctx, inputs[0], &exp_output, "Exp"));
75 std::unique_ptr<GradientFunction> gradient_function(
76 new PassThroughGradientFunction);
77 tape.RecordOperation(inputs, {exp_output}, gradient_function.release());
78 TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx,
79 /*targets*/ {exp_output},
80 /*sources=*/inputs,
81 /*output_gradients=*/{},
82 /*result=*/outputs));
83 exp_output->Unref();
84 return OkStatus();
85 }
86
TEST_P(CustomGradientTest,ExpWithPassThroughGrad)87 TEST_P(CustomGradientTest, ExpWithPassThroughGrad) {
88 std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
89 TF_NewStatus(), TF_DeleteStatus);
90 AbstractContextPtr ctx;
91 {
92 AbstractContext* ctx_raw = nullptr;
93 Status s =
94 BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
95 ASSERT_EQ(errors::OK, s.code()) << s.error_message();
96 ctx.reset(ctx_raw);
97 }
98
99 AbstractTensorHandlePtr x;
100 {
101 AbstractTensorHandle* x_raw = nullptr;
102 Status s = TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 1.0f, &x_raw);
103 ASSERT_EQ(errors::OK, s.code()) << s.error_message();
104 x.reset(x_raw);
105 }
106
107 // Pseudo-code:
108 //
109 // tape.watch(x)
110 // y = exp(x)
111 // outputs = tape.gradient(y, x)
112 std::vector<AbstractTensorHandle*> outputs(1);
113 Status s = RunModel(ExpWithPassThroughGrad, ctx.get(), {x.get()},
114 absl::MakeSpan(outputs),
115 /*use_function=*/!std::get<2>(GetParam()));
116 ASSERT_EQ(errors::OK, s.code()) << s.error_message();
117
118 TF_Tensor* result_tensor;
119 s = GetValue(outputs[0], &result_tensor);
120 ASSERT_EQ(errors::OK, s.code()) << s.error_message();
121 auto result_value = static_cast<float*>(TF_TensorData(result_tensor));
122 EXPECT_EQ(*result_value, 1.0);
123 outputs[0]->Unref();
124 TF_DeleteTensor(result_tensor);
125 result_tensor = nullptr;
126 }
127
128 #ifdef PLATFORM_GOOGLE
129 INSTANTIATE_TEST_SUITE_P(
130 CustomGradientTest, CustomGradientTest,
131 ::testing::Combine(::testing::Values("graphdef", "mlir"),
132 /*tfrt*/ ::testing::Values(true, false),
133 /*executing_eagerly*/ ::testing::Values(true, false)));
134 #else
135 INSTANTIATE_TEST_SUITE_P(
136 CustomGradientTest, CustomGradientTest,
137 ::testing::Combine(::testing::Values("graphdef", "mlir"),
138 /*tfrt*/ ::testing::Values(false),
139 /*executing_eagerly*/ ::testing::Values(true, false)));
140 #endif
141 } // namespace
142 } // namespace internal
143 } // namespace gradients
144 } // namespace tensorflow
145