1 /* Copyright 2016 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 
16 #include <memory>
17 #include <vector>
18 
19 #include "tensorflow/core/framework/function_testlib.h"
20 #include "tensorflow/core/framework/op_kernel.h"
21 #include "tensorflow/core/framework/tensor_testutil.h"
22 #include "tensorflow/core/lib/gtl/array_slice.h"
23 #include "tensorflow/core/lib/strings/str_util.h"
24 #include "tensorflow/core/platform/test.h"
25 #include "tensorflow/core/public/session.h"
26 
27 namespace tensorflow {
28 namespace {
29 
30 namespace f = test::function;
31 using FDH = FunctionDefHelper;
32 
NewSession()33 std::unique_ptr<Session> NewSession() {
34   SessionOptions opts;
35   (*opts.config.mutable_device_count())["CPU"] = 1;
36   return std::unique_ptr<Session>(NewSession(opts));
37 }
38 
39 class MathGradTest : public ::testing::Test {
40  protected:
41   // Unary
42   // dst is the output dtype of op_node.
Unary(const FDH::Node & op_node,const Tensor & x,const DataType dst,Tensor * y)43   Status Unary(const FDH::Node& op_node, const Tensor& x, const DataType dst,
44                Tensor* y) {
45     const DataType src = x.dtype();
46     auto adef = [](const string& name,
47                    const DataType type) {  // E.g., x:float, dy:double
48       return strings::StrCat(name, ":", DataTypeString(type));
49     };
50     // Sum(op(x)), sum all output of op(x).
51     auto test = FDH::Define("Test", {adef("x", src)}, {adef("l", dst)}, {},
52                             {
53                                 op_node,
54                                 FDH::Const("zero", 0),
55                                 FDH::Const("one", 1),
56                                 {{"r"}, "Rank", {"x"}, {{"T", src}}},
57                                 {{"indices"}, "Range", {"zero", "r", "one"}},
58                                 {{"l"}, "Sum", {"y", "indices"}, {{"T", dst}}},
59                             });
60 
61     // TestGrad = Test'(x)
62     auto grad = FDH::Define(
63         "TestGrad", {adef("x", src)}, {adef("dx", src)}, {},
64         {
65             FDH::Const("one", 1),
66             {{"dy"}, "Cast", {"one"}, {{"DstT", dst}, {"SrcT", DT_INT32}}},
67             {{"grad"},
68              "SymbolicGradient",
69              {"x", "dy"},
70              {
71                  {"f", FDH::FunctionRef("Test")},
72                  {"Tin", DataTypeSlice{src, dst}},
73                  {"Tout", DataTypeSlice{src}},
74              }},
75             {{"dx"}, "Identity", {"grad"}, {{"T", src}}},
76         });
77     // Each test case will feed in "x:0" and expects to get "dx:0".
78     auto gdef = test::function::GDef(
79         {
80             f::NDef("x", "Placeholder", {}, {{"dtype", src}}),
81             f::NDef("dx", "TestGrad", {"x"}, {}),
82         },
83         {test, grad});
84 
85     auto sess = NewSession();
86     TF_CHECK_OK(sess->Create(gdef));
87     std::vector<Tensor> outputs;
88     auto s = sess->Run({{"x:0", x}}, {"dx:0"}, {}, &outputs);
89     if (s.ok()) {
90       CHECK_EQ(outputs.size(), 1);
91       *y = outputs[0];
92     }
93     TF_CHECK_OK(sess->Close());
94     return s;
95   }
96 
Unary(const string & op,const Tensor & x,Tensor * y)97   Status Unary(const string& op, const Tensor& x, Tensor* y) {
98     const FDH::Node op_node = {{"y"}, op, {"x"}, {{"T", x.dtype()}}};
99     return Unary(op_node, x, x.dtype(), y);
100   }
101 
102   // Unary op expecting OK.
SymGrad(const string & op,const Tensor & x)103   Tensor SymGrad(const string& op, const Tensor& x) {
104     Tensor ret;
105     TF_CHECK_OK(Unary(op, x, &ret));
106     return ret;
107   }
108 
SymCastGrad(const Tensor & x,const DataType dst)109   Tensor SymCastGrad(const Tensor& x, const DataType dst) {
110     Tensor ret;
111     const FDH::Node op_node = {
112         {"y"}, "Cast", {"x"}, {{"SrcT", x.dtype()}, {"DstT", dst}}};
113     TF_CHECK_OK(Unary(op_node, x, dst, &ret));
114     return ret;
115   }
116 
117   // Binary
SymGrad(const string & op,const Tensor & x,const Tensor & y,Tensor * dx,Tensor * dy)118   void SymGrad(const string& op, const Tensor& x, const Tensor& y, Tensor* dx,
119                Tensor* dy) {
120     const DataType T = x.dtype();
121     auto adef = [T](const string& name) {  // E.g., x:float, dy:double
122       return strings::StrCat(name, ":", DataTypeString(T));
123     };
124     // Sum(op(x)), sum all output of op(x).
125     auto test = FDH::Define("Test", {adef("x"), adef("y")}, {adef("l")}, {},
126                             {
127                                 {{"z"}, op, {"x", "y"}, {{"T", T}}},
128                                 FDH::Const("zero", 0),
129                                 FDH::Const("one", 1),
130                                 {{"r"}, "Rank", {"z"}, {{"T", T}}},
131                                 {{"indices"}, "Range", {"zero", "r", "one"}},
132                                 {{"l"}, "Sum", {"z", "indices"}, {{"T", T}}},
133                             });
134 
135     // TestGrad = Test'(x, y)
136     auto grad = FDH::Define(
137         "TestGrad", {adef("x"), adef("y")}, {adef("dx"), adef("dy")}, {},
138         {
139             FDH::Const("one", 1),
140             {{"dz"}, "Cast", {"one"}, {{"DstT", T}, {"SrcT", DT_INT32}}},
141             {{"grad0", "grad1"},
142              "SymbolicGradient",
143              {"x", "y", "dz"},
144              {
145                  {"f", FDH::FunctionRef("Test")},
146                  {"Tin", DataTypeSlice{T, T, T}},
147                  {"Tout", DataTypeSlice{T, T}},
148              }},
149             {{"dx"}, "Identity", {"grad0"}, {{"T", T}}},
150             {{"dy"}, "Identity", {"grad1"}, {{"T", T}}},
151         });
152     // Each test case will feed in "x:0" and "y:0" and expects to get "d0" and
153     // "d:0".
154     auto gdef = test::function::GDef(
155         {
156             f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
157             f::NDef("y", "Placeholder", {}, {{"dtype", T}}),
158             f::NDef("d", "TestGrad", {"x", "y"}, {}),
159         },
160         {test, grad});
161 
162     auto sess = NewSession();
163     TF_CHECK_OK(sess->Create(gdef));
164     std::vector<Tensor> outputs;
165     TF_CHECK_OK(
166         sess->Run({{"x:0", x}, {"y:0", y}}, {"d:0", "d:1"}, {}, &outputs));
167     CHECK_EQ(outputs.size(), 2);
168     TF_CHECK_OK(sess->Close());
169     *dx = outputs[0];
170     *dy = outputs[1];
171   }
172 
173   // Reduction grad
ReductionGrad(const string & op,const Tensor & x,const Tensor & idx,Tensor * dx,Tensor * di)174   void ReductionGrad(const string& op, const Tensor& x, const Tensor& idx,
175                      Tensor* dx, Tensor* di) {
176     const DataType T = x.dtype();
177     auto adef = [T](const string& name) {  // E.g., x:float, dy:double
178       return strings::StrCat(name, ":", DataTypeString(T));
179     };
180     // Sum(op(x, idx)), sum all output of op(x, idx).
181     auto test = FDH::Define("Test", {adef("x"), "i:int32"}, {adef("l")}, {},
182                             {
183                                 {{"y"}, op, {"x", "i"}, {{"T", T}}},
184                                 FDH::Const("zero", 0),
185                                 FDH::Const("one", 1),
186                                 {{"r"}, "Rank", {"y"}, {{"T", T}}},
187                                 {{"indices"}, "Range", {"zero", "r", "one"}},
188                                 {{"l"}, "Sum", {"y", "indices"}, {{"T", T}}},
189                             });
190 
191     // TestGrad = Test'(x)
192     auto grad = FDH::Define(
193         "TestGrad", {adef("x"), "i:int32"}, {adef("dx"), "di:int32"}, {},
194         {
195             FDH::Const("one", 1),
196             {{"dy"}, "Cast", {"one"}, {{"DstT", T}, {"SrcT", DT_INT32}}},
197             {{"grad0", "grad1"},
198              "SymbolicGradient",
199              {"x", "i", "dy"},
200              {
201                  {"f", FDH::FunctionRef("Test")},
202                  {"Tin", DataTypeSlice{T, DT_INT32, T}},
203                  {"Tout", DataTypeSlice{T, DT_INT32}},
204              }},
205             {{"dx"}, "Identity", {"grad0"}, {{"T", T}}},
206             {{"di"}, "Identity", {"grad1"}, {{"T", DT_INT32}}},
207         });
208     // Each test case will feed in "x:0" and expects to get "dx:0".
209     auto gdef = test::function::GDef(
210         {
211             f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
212             f::NDef("i", "Placeholder", {}, {{"dtype", DT_INT32}}),
213             f::NDef("d", "TestGrad", {"x", "i"}, {}),
214         },
215         {test, grad});
216 
217     auto sess = NewSession();
218     TF_CHECK_OK(sess->Create(gdef));
219     std::vector<Tensor> outputs;
220     TF_CHECK_OK(
221         sess->Run({{"x:0", x}, {"i:0", idx}}, {"d:0", "d:1"}, {}, &outputs));
222     CHECK_EQ(outputs.size(), 2);
223     TF_CHECK_OK(sess->Close());
224     *dx = outputs[0];
225     *di = outputs[1];
226   }
227 
ReduceSum(const Tensor & x,gtl::ArraySlice<int32> axes)228   Tensor ReduceSum(const Tensor& x, gtl::ArraySlice<int32> axes) {
229     int num_axes = axes.length();
230     Tensor y(DT_INT32, TensorShape({num_axes}));
231     for (size_t i = 0; i < axes.size(); ++i) {
232       y.flat<int32>()(i) = axes[i];
233     }
234     auto T = x.dtype();
235     auto gdef = test::function::GDef(
236         {
237             f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
238             f::NDef("y", "Const", {}, {{"dtype", DT_INT32}, {"value", y}}),
239             f::NDef("z", "Sum", {"x", "y"}, {{"T", T}}),
240         },
241         {});
242     auto sess = NewSession();
243     TF_CHECK_OK(sess->Create(gdef));
244     std::vector<Tensor> outputs;
245     TF_CHECK_OK(sess->Run({{"x:0", x}}, {"z:0"}, {}, &outputs));
246     CHECK_EQ(outputs.size(), 1);
247     TF_CHECK_OK(sess->Close());
248     return outputs[0];
249   }
250 
MatMulCommon(const string & opname,const string & attr_adj_x,const string & attr_adj_y,const Tensor & x,bool ax,const Tensor & y,bool ay)251   Tensor MatMulCommon(const string& opname, const string& attr_adj_x,
252                       const string& attr_adj_y, const Tensor& x, bool ax,
253                       const Tensor& y, bool ay) {
254     auto T = x.dtype();
255     auto gdef = test::function::GDef(
256         {
257             f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
258             f::NDef("y", "Placeholder", {}, {{"dtype", T}}),
259             f::NDef("z", opname, {"x", "y"},
260                     {{"T", T}, {attr_adj_x, ax}, {attr_adj_y, ay}}),
261         },
262         {});
263     auto sess = NewSession();
264     TF_CHECK_OK(sess->Create(gdef));
265     std::vector<Tensor> outputs;
266     TF_CHECK_OK(sess->Run({{"x:0", x}, {"y:0", y}}, {"z:0"}, {}, &outputs));
267     CHECK_EQ(outputs.size(), 1);
268     TF_CHECK_OK(sess->Close());
269     return outputs[0];
270   }
271 
MatMul(const Tensor & x,bool ax,const Tensor & y,bool ay)272   Tensor MatMul(const Tensor& x, bool ax, const Tensor& y, bool ay) {
273     return MatMulCommon("MatMul", "transpose_a", "transpose_b", x, ax, y, ay);
274   }
275 
BatchMatMul(const Tensor & x,bool ax,const Tensor & y,bool ay)276   Tensor BatchMatMul(const Tensor& x, bool ax, const Tensor& y, bool ay) {
277     return MatMulCommon("BatchMatMul", "adj_x", "adj_y", x, ax, y, ay);
278   }
279 
BatchMatMulV2(const Tensor & x,bool ax,const Tensor & y,bool ay)280   Tensor BatchMatMulV2(const Tensor& x, bool ax, const Tensor& y, bool ay) {
281     return MatMulCommon("BatchMatMulV2", "adj_x", "adj_y", x, ax, y, ay);
282   }
283 
MatMulGradCommon(const string & opname,const string & attr_adj_x,const string & attr_adj_y,const Tensor & x,bool ax,const Tensor & y,bool ay,Tensor * dx,Tensor * dy)284   void MatMulGradCommon(const string& opname, const string& attr_adj_x,
285                         const string& attr_adj_y, const Tensor& x, bool ax,
286                         const Tensor& y, bool ay, Tensor* dx, Tensor* dy) {
287     const DataType T = x.dtype();
288     auto adef = [T](const string& name) {  // E.g., x:float, dy:double
289       return strings::StrCat(name, ":", DataTypeString(T));
290     };
291     // Sum(op(x)), sum all output of op(x).
292     auto test =
293         FDH::Define("Test", {adef("x"), adef("y")}, {adef("l")}, {},
294                     {
295                         {{"z"},
296                          opname,
297                          {"x", "y"},
298                          {{"T", T}, {attr_adj_x, ax}, {attr_adj_y, ay}}},
299                         FDH::Const("zero", 0),
300                         FDH::Const("one", 1),
301                         {{"r"}, "Rank", {"z"}, {{"T", T}}},
302                         {{"indices"}, "Range", {"zero", "r", "one"}},
303                         {{"l"}, "Sum", {"z", "indices"}, {{"T", T}}},
304                     });
305 
306     // TestGrad = Test'(x, y)
307     auto grad = FDH::Define(
308         "TestGrad", {adef("x"), adef("y")}, {adef("dx"), adef("dy")}, {},
309         {
310             FDH::Const("one", 1),
311             {{"dz"}, "Cast", {"one"}, {{"DstT", T}, {"SrcT", DT_INT32}}},
312             {{"grad0", "grad1"},
313              "SymbolicGradient",
314              {"x", "y", "dz"},
315              {
316                  {"f", FDH::FunctionRef("Test")},
317                  {"Tin", DataTypeSlice{T, T, T}},
318                  {"Tout", DataTypeSlice{T, T}},
319              }},
320             {{"dx"}, "Identity", {"grad0"}, {{"T", T}}},
321             {{"dy"}, "Identity", {"grad1"}, {{"T", T}}},
322         });
323     // Each test case will feed in "x:0" and "y:0" and expects to get "d0" and
324     // "d:0".
325     auto gdef = test::function::GDef(
326         {
327             f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
328             f::NDef("y", "Placeholder", {}, {{"dtype", T}}),
329             f::NDef("d", "TestGrad", {"x", "y"}, {}),
330         },
331         {test, grad});
332 
333     auto sess = NewSession();
334     TF_CHECK_OK(sess->Create(gdef));
335     std::vector<Tensor> outputs;
336     TF_CHECK_OK(
337         sess->Run({{"x:0", x}, {"y:0", y}}, {"d:0", "d:1"}, {}, &outputs));
338     CHECK_EQ(outputs.size(), 2);
339     TF_CHECK_OK(sess->Close());
340     *dx = outputs[0];
341     *dy = outputs[1];
342   }
343 
MatMulGrad(const Tensor & x,bool ax,const Tensor & y,bool ay,Tensor * dx,Tensor * dy)344   void MatMulGrad(const Tensor& x, bool ax, const Tensor& y, bool ay,
345                   Tensor* dx, Tensor* dy) {
346     return MatMulGradCommon("MatMul", "transpose_a", "transpose_b", x, ax, y,
347                             ay, dx, dy);
348   }
349 
BatchMatMulGrad(const Tensor & x,bool ax,const Tensor & y,bool ay,Tensor * dx,Tensor * dy)350   void BatchMatMulGrad(const Tensor& x, bool ax, const Tensor& y, bool ay,
351                        Tensor* dx, Tensor* dy) {
352     return MatMulGradCommon("BatchMatMul", "adj_x", "adj_y", x, ax, y, ay, dx,
353                             dy);
354   }
355 
BatchMatMulV2Grad(const Tensor & x,bool ax,const Tensor & y,bool ay,Tensor * dx,Tensor * dy)356   void BatchMatMulV2Grad(const Tensor& x, bool ax, const Tensor& y, bool ay,
357                          Tensor* dx, Tensor* dy) {
358     return MatMulGradCommon("BatchMatMulV2", "adj_x", "adj_y", x, ax, y, ay, dx,
359                             dy);
360   }
361 
SelectGrad(const Tensor & c,const Tensor & x,const Tensor & y,Tensor * dc,Tensor * dx,Tensor * dy)362   void SelectGrad(const Tensor& c, const Tensor& x, const Tensor& y, Tensor* dc,
363                   Tensor* dx, Tensor* dy) {
364     auto T = DT_FLOAT;
365     // Sum(Select(c, x, y))
366     auto test =
367         FDH::Define("Test", {"c:bool", "x:float", "y:float"}, {"l:float"}, {},
368                     {
369                         {{"z"}, "Select", {"c", "x", "y"}, {{"T", T}}},
370                         FDH::Const("zero", 0),
371                         FDH::Const("one", 1),
372                         {{"r"}, "Rank", {"z"}, {{"T", T}}},
373                         {{"indices"}, "Range", {"zero", "r", "one"}},
374                         {{"l"}, "Sum", {"z", "indices"}, {{"T", T}}},
375                     });
376 
377     // TestGrad(x, y) = Test'(c, x, y)
378     auto grad = FDH::Define("TestGrad", {"c:bool", "x:float", "y:float"},
379                             {"dc:bool", "dx:float", "dy:float"}, {},
380                             {FDH::Const("dz", 1.f),
381                              {{"grad0", "grad1", "grad2"},
382                               "SymbolicGradient",
383                               {"c", "x", "y", "dz"},
384                               {
385                                   {"f", FDH::FunctionRef("Test")},
386                                   {"Tin", DataTypeSlice{DT_BOOL, T, T, T}},
387                                   {"Tout", DataTypeSlice{DT_BOOL, T, T}},
388                               }},
389                              {{"dc"}, "Identity", {"grad0"}, {{"T", DT_BOOL}}},
390                              {{"dx"}, "Identity", {"grad1"}, {{"T", T}}},
391                              {{"dy"}, "Identity", {"grad2"}, {{"T", T}}}});
392     // Each test case will feed in "x:0" and expects to get "dx:0".
393     auto gdef = test::function::GDef(
394         {
395             f::NDef("c", "Placeholder", {}, {{"dtype", DT_BOOL}}),
396             f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
397             f::NDef("y", "Placeholder", {}, {{"dtype", T}}),
398             f::NDef("d", "TestGrad", {"c", "x", "y"}, {}),
399         },
400         {test, grad});
401 
402     auto sess = NewSession();
403     TF_CHECK_OK(sess->Create(gdef));
404     std::vector<Tensor> outputs;
405     TF_CHECK_OK(sess->Run({{"c:0", c}, {"x:0", x}, {"y:0", y}},
406                           {"d:0", "d:1", "d:2"}, {}, &outputs));
407     CHECK_EQ(outputs.size(), 3);
408     TF_CHECK_OK(sess->Close());
409     *dc = outputs[0];
410     *dx = outputs[1];
411     *dy = outputs[2];
412   }
413 };
414 
HasError(const Status & s,const string & substr)415 void HasError(const Status& s, const string& substr) {
416   EXPECT_TRUE(absl::StrContains(s.ToString(), substr))
417       << s << ", expected substring " << substr;
418 }
419 
420 REGISTER_OP("TestOpWithNoGrad")
421     .Input("x: T")
422     .Output("y: T")
423     .Attr("T: {float, double}")
424     .Doc(R"doc(
425 Test op with no grad registered.
426 
427 x: input
428 y: output
429 )doc");
430 
431 class TestOp : public OpKernel {
432  public:
TestOp(OpKernelConstruction * ctx)433   explicit TestOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
Compute(OpKernelContext * ctx)434   void Compute(OpKernelContext* ctx) override { ctx->set_output(0, Tensor()); }
435 };
436 REGISTER_KERNEL_BUILDER(Name("TestOpWithNoGrad").Device(DEVICE_CPU), TestOp);
437 
TEST_F(MathGradTest,Error_Reporting)438 TEST_F(MathGradTest, Error_Reporting) {
439   auto x = test::AsTensor<float>({-3.f});
440   auto dx = test::AsTensor<float>({3.f});
441   Tensor donotcare;
442   HasError(Unary("TestOpWithNoGrad", x, &donotcare),
443            "No gradient defined for op: TestOpWithNoGrad");
444 }
445 
TEST_F(MathGradTest,Abs)446 TEST_F(MathGradTest, Abs) {
447   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
448                                  TensorShape({2, 3}));
449   auto g = [](float x) { return x < 0 ? -1.f : 1.f; };
450   auto dx = test::AsTensor<float>(
451       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
452   auto ans = SymGrad("Abs", x);
453   test::ExpectClose(ans, dx);
454 }
455 
TEST_F(MathGradTest,Neg)456 TEST_F(MathGradTest, Neg) {
457   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
458                                  TensorShape({2, 3}));
459   auto g = [](float x) { return -1.f; };
460   auto dx = test::AsTensor<float>(
461       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
462   auto ans = SymGrad("Neg", x);
463   test::ExpectClose(ans, dx);
464 }
465 
TEST_F(MathGradTest,Reciprocal)466 TEST_F(MathGradTest, Reciprocal) {
467   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
468                                  TensorShape({2, 3}));
469   auto g = [](float x) { return -1.f / (x * x); };
470   auto dx = test::AsTensor<float>(
471       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
472   auto ans = SymGrad("Reciprocal", x);
473   test::ExpectClose(ans, dx);
474 }
475 
TEST_F(MathGradTest,Square)476 TEST_F(MathGradTest, Square) {
477   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
478                                  TensorShape({2, 3}));
479   auto g = [](float x) { return 2 * x; };
480   auto dx = test::AsTensor<float>(
481       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
482   auto ans = SymGrad("Square", x);
483   test::ExpectClose(ans, dx);
484 }
485 
TEST_F(MathGradTest,Sqrt)486 TEST_F(MathGradTest, Sqrt) {
487   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
488                                  TensorShape({2, 3}));
489   auto g = [](float x) { return 0.5f / std::sqrt(x); };
490   auto dx = test::AsTensor<float>(
491       {g(1.f), g(2.f), g(3.f), g(4.f), g(5.f), g(6.f)}, TensorShape({2, 3}));
492   auto ans = SymGrad("Sqrt", x);
493   test::ExpectClose(ans, dx);
494 }
495 
TEST_F(MathGradTest,Rsqrt)496 TEST_F(MathGradTest, Rsqrt) {
497   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
498                                  TensorShape({2, 3}));
499   auto g = [](float x) { return -0.5f / (x * std::sqrt(x)); };
500   auto dx = test::AsTensor<float>(
501       {g(1.f), g(2.f), g(3.f), g(4.f), g(5.f), g(6.f)}, TensorShape({2, 3}));
502   auto ans = SymGrad("Rsqrt", x);
503   test::ExpectClose(ans, dx);
504 }
505 
TEST_F(MathGradTest,Exp)506 TEST_F(MathGradTest, Exp) {
507   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
508                                  TensorShape({2, 3}));
509   auto g = [](float x) { return std::exp(x); };
510   auto dx = test::AsTensor<float>(
511       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
512   auto ans = SymGrad("Exp", x);
513   test::ExpectClose(ans, dx);
514 }
515 
TEST_F(MathGradTest,Expm1)516 TEST_F(MathGradTest, Expm1) {
517   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
518                                  TensorShape({2, 3}));
519   auto g = [](float x) { return std::exp(x); };
520   auto dx = test::AsTensor<float>(
521       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
522   auto ans = SymGrad("Expm1", x);
523   test::ExpectClose(ans, dx);
524 }
525 
TEST_F(MathGradTest,Log)526 TEST_F(MathGradTest, Log) {
527   auto x = test::AsTensor<float>({0.1f, 1.f, 2.f, 3.f, 4.f, 10.f},
528                                  TensorShape({2, 3}));
529   auto g = [](float x) { return 1 / x; };
530   auto dx = test::AsTensor<float>(
531       {g(.1f), g(1.f), g(2.f), g(3.f), g(4.f), g(10.f)}, TensorShape({2, 3}));
532   auto ans = SymGrad("Log", x);
533   test::ExpectClose(ans, dx);
534 }
535 
TEST_F(MathGradTest,Log1p)536 TEST_F(MathGradTest, Log1p) {
537   auto x = test::AsTensor<float>({0.1f, 1.f, 2.f, 3.f, 4.f, 10.f},
538                                  TensorShape({2, 3}));
539   auto g = [](float x) { return 1 / (1 + x); };
540   auto dx = test::AsTensor<float>(
541       {g(.1f), g(1.f), g(2.f), g(3.f), g(4.f), g(10.f)}, TensorShape({2, 3}));
542   auto ans = SymGrad("Log1p", x);
543   test::ExpectClose(ans, dx);
544 }
545 
TEST_F(MathGradTest,Sinh)546 TEST_F(MathGradTest, Sinh) {
547   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
548                                  TensorShape({2, 3}));
549   auto g = [](float x) { return std::cosh(x); };
550   auto dx = test::AsTensor<float>(
551       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
552   auto ans = SymGrad("Sinh", x);
553   test::ExpectClose(ans, dx);
554 }
555 
TEST_F(MathGradTest,Cosh)556 TEST_F(MathGradTest, Cosh) {
557   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
558                                  TensorShape({2, 3}));
559   auto g = [](float x) { return std::sinh(x); };
560   auto dx = test::AsTensor<float>(
561       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
562   auto ans = SymGrad("Cosh", x);
563   test::ExpectClose(ans, dx);
564 }
565 
TEST_F(MathGradTest,Tanh)566 TEST_F(MathGradTest, Tanh) {
567   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
568                                  TensorShape({2, 3}));
569   auto g = [](float x) {
570     auto y = std::tanh(x);
571     return 1 - y * y;
572   };
573   auto dx = test::AsTensor<float>(
574       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
575   auto ans = SymGrad("Tanh", x);
576   test::ExpectClose(ans, dx);
577 }
578 
TEST_F(MathGradTest,Asinh)579 TEST_F(MathGradTest, Asinh) {
580   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
581                                  TensorShape({2, 3}));
582   auto g = [](float x) {
583     auto y = std::asinh(x);
584     return std::cosh(y);
585   };
586   auto dx = test::AsTensor<float>(
587       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
588   auto ans = SymGrad("Asinh", x);
589   test::ExpectClose(ans, dx);
590 }
591 
TEST_F(MathGradTest,Acosh)592 TEST_F(MathGradTest, Acosh) {
593   auto x = test::AsTensor<float>({6.f, 5.f, 4.f, 1.f, 2.f, 3.f},
594                                  TensorShape({2, 3}));
595   auto g = [](float x) {
596     auto y = std::acosh(x);
597     return std::sinh(y);
598   };
599   auto dx = test::AsTensor<float>(
600       {g(6.f), g(5.f), g(4.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
601   auto ans = SymGrad("Acosh", x);
602   test::ExpectClose(ans, dx);
603 }
604 
TEST_F(MathGradTest,Atanh)605 TEST_F(MathGradTest, Atanh) {
606   auto x = test::AsTensor<float>({-0.3f, -0.2f, -0.1f, 0.1f, 0.2f, 0.3f},
607                                  TensorShape({2, 3}));
608   auto g = [](float x) { return 1.f / (1.f - x * x); };
609   auto dx = test::AsTensor<float>(
610       {g(-0.3f), g(-0.2f), g(-0.1f), g(0.1f), g(0.2f), g(0.3f)},
611       TensorShape({2, 3}));
612   auto ans = SymGrad("Atanh", x);
613   test::ExpectClose(ans, dx);
614 }
615 
TEST_F(MathGradTest,Sigmoid)616 TEST_F(MathGradTest, Sigmoid) {
617   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
618                                  TensorShape({2, 3}));
619   auto g = [](float x) {
620     auto y = 1.f / (1.f + std::exp(-x));
621     return y * (1 - y);
622   };
623   auto dx = test::AsTensor<float>(
624       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
625   auto ans = SymGrad("Sigmoid", x);
626   test::ExpectClose(ans, dx);
627 }
628 
TEST_F(MathGradTest,Sign)629 TEST_F(MathGradTest, Sign) {
630   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
631                                  TensorShape({2, 3}));
632   auto g = [](float x) { return 0.f; };
633   auto dx = test::AsTensor<float>(
634       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
635   auto ans = SymGrad("Sign", x);
636   test::ExpectClose(ans, dx);
637 }
638 
TEST_F(MathGradTest,Sin)639 TEST_F(MathGradTest, Sin) {
640   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
641                                  TensorShape({2, 3}));
642   auto g = [](float x) { return std::cos(x); };
643   auto dx = test::AsTensor<float>(
644       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
645   auto ans = SymGrad("Sin", x);
646   test::ExpectClose(ans, dx);
647 }
648 
TEST_F(MathGradTest,Cos)649 TEST_F(MathGradTest, Cos) {
650   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
651                                  TensorShape({2, 3}));
652   auto g = [](float x) { return -std::sin(x); };
653   auto dx = test::AsTensor<float>(
654       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
655   auto ans = SymGrad("Cos", x);
656   test::ExpectClose(ans, dx);
657 }
658 
TEST_F(MathGradTest,Cast)659 TEST_F(MathGradTest, Cast) {
660   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
661                                  TensorShape({2, 3}));
662   auto g = [](float x) { return 1.f; };
663   auto dx = test::AsTensor<float>(
664       {g(-3.f), g(-2.f), g(-1.f), g(1.f), g(2.f), g(3.f)}, TensorShape({2, 3}));
665   Tensor ans = SymCastGrad(x, DT_INT32);
666   test::ExpectClose(ans, dx);
667 }
668 
669 // TODO(zhifengc)
670 // TEST_F(MathGradSComplexTest, Real) {}
671 // TEST_F(MathGradSComplexTest, Imag) {}
672 // TEST_F(MathGradSComplexTest, Angle) {}
673 // TEST_F(MathGradSComplexTest, Conj) {}
674 // TEST_F(MathGradTernary, Select) {}
675 
TEST_F(MathGradTest,Add)676 TEST_F(MathGradTest, Add) {
677   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
678                                  TensorShape({2, 3}));
679   auto y = test::AsTensor<float>({-10.f, 10.f}, TensorShape({2, 1}));
680   auto ans_dx = test::AsTensor<float>({1.f, 1.f, 1.f, 1.f, 1.f, 1.f},
681                                       TensorShape({2, 3}));
682   auto ans_dy = test::AsTensor<float>({3.f, 3.f}, TensorShape({2, 1}));
683   Tensor dx;
684   Tensor dy;
685   {
686     SymGrad("Add", x, y, &dx, &dy);
687     test::ExpectClose(ans_dx, dx);
688     test::ExpectClose(ans_dy, dy);
689   }
690   {  // Swap x and y
691     SymGrad("Add", y, x, &dy, &dx);
692     test::ExpectClose(ans_dx, dx);
693     test::ExpectClose(ans_dy, dy);
694   }
695 }
696 
TEST_F(MathGradTest,Sub)697 TEST_F(MathGradTest, Sub) {
698   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
699                                  TensorShape({2, 3}));
700   auto y = test::AsTensor<float>({-10.f, 10.f}, TensorShape({2, 1}));
701   Tensor dx;
702   Tensor dy;
703   {
704     SymGrad("Sub", x, y, &dx, &dy);
705     auto ans_dx = test::AsTensor<float>({1.f, 1.f, 1.f, 1.f, 1.f, 1.f},
706                                         TensorShape({2, 3}));
707     auto ans_dy = test::AsTensor<float>({-3.f, -3.f}, TensorShape({2, 1}));
708     test::ExpectClose(ans_dx, dx);
709     test::ExpectClose(ans_dy, dy);
710   }
711   {  // Swap x and y
712     SymGrad("Sub", y, x, &dy, &dx);
713     auto ans_dx = test::AsTensor<float>({-1.f, -1.f, -1.f, -1.f, -1.f, -1.f},
714                                         TensorShape({2, 3}));
715     auto ans_dy = test::AsTensor<float>({3.f, 3.f}, TensorShape({2, 1}));
716     test::ExpectClose(ans_dx, dx);
717     test::ExpectClose(ans_dy, dy);
718   }
719 }
720 
TEST_F(MathGradTest,Mul)721 TEST_F(MathGradTest, Mul) {
722   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
723                                  TensorShape({2, 3}));
724   auto y = test::AsTensor<float>({-10.f, 10.f}, TensorShape({2, 1}));
725   auto ans_dx = test::AsTensor<float>({-10.f, -10.f, -10.f, 10.f, 10.f, 10.f},
726                                       TensorShape({2, 3}));
727   auto ans_dy = test::AsTensor<float>({-3.f + (-2.f) + (-1.f), 1.f + 2.f + 3.f},
728                                       TensorShape({2, 1}));
729   Tensor dx;
730   Tensor dy;
731   {
732     SymGrad("Mul", x, y, &dx, &dy);
733     test::ExpectClose(ans_dx, dx);
734     test::ExpectClose(ans_dy, dy);
735   }
736   {  // Swap x and y
737     SymGrad("Mul", y, x, &dy, &dx);
738     test::ExpectClose(ans_dx, dx);
739     test::ExpectClose(ans_dy, dy);
740   }
741 }
742 
TEST_F(MathGradTest,Div)743 TEST_F(MathGradTest, Div) {
744   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
745                                  TensorShape({2, 3}));
746   auto y = test::AsTensor<float>({-10.f, 10.f}, TensorShape({2, 1}));
747   Tensor dx;
748   Tensor dy;
749   {
750     SymGrad("Div", x, y, &dx, &dy);
751     {
752       auto g = [](float x, float y) { return 1.f / y; };
753       test::ExpectClose(dx, test::AsTensor<float>(
754                                 {g(-3.f, -10.f), g(-2.f, -10.f), g(-1.f, -10.f),
755                                  g(1.f, 10.f), g(2.f, 10.f), g(3.f, 10.f)},
756                                 TensorShape({2, 3})));
757     }
758     {
759       auto g = [](float x, float y) { return -x / (y * y); };
760       test::ExpectClose(dy,
761                         test::AsTensor<float>(
762                             {g(-3.f, -10.f) + g(-2.f, -10.f) + g(-1.f, -10.f),
763                              g(1.f, 10.f) + g(2.f, 10.f) + g(3.f, 10.f)},
764                             TensorShape({2, 1})));
765     }
766   }
767   {  // Swap x and y
768     SymGrad("Div", y, x, &dy, &dx);
769     {
770       auto g = [](float x, float y) { return 1.f / y; };
771       test::ExpectClose(dy,
772                         test::AsTensor<float>(
773                             {g(-10.f, -3.f) + g(-10.f, -2.f) + g(-10.f, -1.f),
774                              g(10.f, 1.f) + g(10.f, 2.f) + g(10.f, 3.f)},
775                             TensorShape({2, 1})));
776     }
777     {
778       auto g = [](float x, float y) { return -x / (y * y); };
779       test::ExpectClose(dx, test::AsTensor<float>(
780                                 {g(-10.f, -3.f), g(-10.f, -2.f), g(-10.f, -1.f),
781                                  g(10.f, 1.f), g(10.f, 2.f), g(10.f, 3.f)},
782                                 TensorShape({2, 3})));
783     }
784   }
785 }
786 
TEST_F(MathGradTest,DivNoNan)787 TEST_F(MathGradTest, DivNoNan) {
788   auto x = test::AsTensor<float>(
789       {0.f, -3.f, -2.f, -1.f, 0.f, 1.f, 2.f, 3.f, 0.f}, TensorShape({3, 3}));
790   auto y = test::AsTensor<float>({-10.f, 0.f, 10.f}, TensorShape({3, 1}));
791   Tensor dx;
792   Tensor dy;
793   {
794     SymGrad("DivNoNan", x, y, &dx, &dy);
795     {
796       auto g = [](float x, float y) {
797         if (y == 0.f) {
798           return 0.f;
799         } else {
800           return 1.f / y;
801         }
802       };
803       test::ExpectClose(dx, test::AsTensor<float>(
804                                 {g(0.f, -10.f), g(-3.f, -10.f), g(-2.f, -10.f),
805                                  g(-1.f, 0.f), g(0.f, 0.f), g(1.f, 0.f),
806                                  g(2.f, 10.f), g(3.f, 10.f), g(0.f, 10.f)},
807                                 TensorShape({3, 3})));
808     }
809     {
810       auto g = [](float x, float y) {
811         if (y == 0.f) {
812           return 0.f;
813         } else {
814           return -x / (y * y);
815         }
816       };
817       test::ExpectClose(dy,
818                         test::AsTensor<float>(
819                             {g(0.f, -10.f) + g(-3.f, -10.f) + g(-2.f, -10.f),
820                              g(-1.f, 0.f) + g(0.f, 0.f) + g(1.f, 0.f),
821                              g(2.f, 10.f) + g(3.f, 10.f) + g(0.f, 10.f)},
822                             TensorShape({3, 1})));
823     }
824   }
825   {  // Swap x and y.
826     SymGrad("DivNoNan", y, x, &dy, &dx);
827     {
828       auto g = [](float x, float y) {
829         if (y == 0.f) {
830           return 0.f;
831         } else {
832           return 1.f / y;
833         }
834       };
835       test::ExpectClose(dy,
836                         test::AsTensor<float>(
837                             {g(-10.f, 0.f) + g(-10.f, -3.f) + g(-10.f, -2.f),
838                              g(0.f, -1.f) + g(0.f, 0.f) + g(0.f, 1.f),
839                              g(10.f, 2.f) + g(10.f, 3.f) + g(10.f, 0.f)},
840                             TensorShape({3, 1})));
841     }
842     {
843       auto g = [](float x, float y) {
844         if (y == 0.f) {
845           return 0.f;
846         } else {
847           return -x / (y * y);
848         }
849       };
850       test::ExpectClose(dx, test::AsTensor<float>(
851                                 {g(-10.f, 0.f), g(-10.f, -3.f), g(-10.f, -2.f),
852                                  g(0.f, -1.f), g(0.f, 0.f), g(0.f, 1.f),
853                                  g(10.f, 2.f), g(10.f, 3.f), g(10.f, 0.f)},
854                                 TensorShape({3, 3})));
855     }
856   }
857 }
858 
TEST_F(MathGradTest,Pow)859 TEST_F(MathGradTest, Pow) {
860   auto x = test::AsTensor<float>({0.f, 1.f, 2.f, 3.f, 4.f, 5.f},
861                                  TensorShape({2, 3}));
862   auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
863   Tensor dx;
864   Tensor dy;
865   auto g = [](float x, float y) { return y * std::pow(x, y - 1); };
866   auto h = [](float x, float y) {
867     return std::pow(x, y) * (x ? std::log(x) : 0);
868   };
869   {
870     SymGrad("Pow", x, y, &dx, &dy);
871     test::ExpectClose(
872         dx, test::AsTensor<float>({g(0.f, .5f), g(1.f, .5f), g(2.f, .5f),
873                                    g(3.f, 2.f), g(4.f, 2.f), g(5.f, 2.f)},
874                                   TensorShape({2, 3})));
875     test::ExpectClose(
876         dy, test::AsTensor<float>({h(0.f, .5f) + h(1.f, .5f) + h(2.f, .5f),
877                                    h(3.f, 2.f) + h(4.f, 2.f) + h(5.f, 2.f)},
878                                   TensorShape({2, 1})));
879   }
880   {  // Swap x and y
881     SymGrad("Pow", y, x, &dy, &dx);
882     test::ExpectClose(
883         dy, test::AsTensor<float>({g(.5f, 0.f) + g(.5f, 1.f) + g(.5f, 2.f),
884                                    g(2.f, 3.f) + g(2.f, 4.f) + g(2.f, 5.f)},
885                                   TensorShape({2, 1})));
886     test::ExpectClose(
887         dx, test::AsTensor<float>({h(.5f, 0.f), h(.5f, 1.f), h(.5f, 2.f),
888                                    h(2.f, 3.f), h(2.f, 4.f), h(2.f, 5.f)},
889                                   TensorShape({2, 3})));
890   }
891 }
892 
TEST_F(MathGradTest,ComplexPow)893 TEST_F(MathGradTest, ComplexPow) {
894   auto x = test::AsTensor<complex64>({0.f, 2.f, -2.f}, TensorShape({3}));
895   auto y = test::AsTensor<complex64>({2.f, 2.f, 2.f}, TensorShape({3}));
896   Tensor dx;
897   Tensor dy;
898   auto g = [](complex64 x, complex64 y) { return y * std::pow(x, y - 1.f); };
899   auto h = [](complex64 x, complex64 y) {
900     return std::pow(x, y) * (x != complex64(0) ? std::log(x) : 0);
901   };
902   SymGrad("Pow", x, y, &dx, &dy);
903 
904   // This case failed on Kokoro MacOS:
905   // dx[2] = (-4,6.0398321011234657e-07),
906   // test::AsTensor[2] = (-4,-3.4969110629390343e-07).
907   // dx[2] on linux is close to test::AsTensor[2].
908   // This error hasn't shown up before because
909   // ExpectClose used to check just the magnitude of a complex number, i.e.,
910   // std::abs(complex) = sqrt(real^2 + imag^2).
911   // Now ExpectClose checks the value of each component separately.
912   // Workaround: I set a big tolerance to make the case pass for now.
913   // TODO(penporn): Fix this or file a bug. This is not a precision issue.
914   // Even the most significant digit (or the sign) doesn't match.
915   test::ExpectClose(
916       dx,
917       test::AsTensor<complex64>({g(0.f, 2.f), g(2.f, 2.f), g(-2.f, 2.f)},
918                                 TensorShape({3})),
919       1e-6f);
920 
921   // This case failed on Kokoro MacOS:
922   // dx[2] = (2.7725925445556641,12.56636905670166),
923   // test::AsTensor[2] = (2.7725865840911865,12.566371917724609)
924   // dx[2] on linux is close to test::AsTensor[2].
925   // Default atol = rtol = 5.96046e-07.
926   // Real: diff = 5.96046e-06 > threshold = 2.248633e-06 <- failed
927   // Complex: diff = 2.86102e-06 <= threshold = 8.08618e-06 <- passed
928   // Again, this error hasn't shown up before because ExpectClose used to
929   // check just the magnitude of the complex number. Now it checks each
930   // component separately.
931   // Workaround: Set a larger tolerance for now.
932   // TODO(penporn): See if this is a precision issue or a bug.
933   test::ExpectClose(
934       dy,
935       test::AsTensor<complex64>({h(0.f, 2.f), h(2.f, 2.f), h(-2.f, 2.f)},
936                                 TensorShape({3})),
937       4.5e-6f);
938 }
939 
TEST_F(MathGradTest,Xlogy)940 TEST_F(MathGradTest, Xlogy) {
941   auto x = test::AsTensor<float>({0.f, 0.f, 2.f, 3.f, 4.f, 5.f},
942                                  TensorShape({2, 3}));
943   auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
944   Tensor dx;
945   Tensor dy;
946   auto g = [](float x, float y) -> float { return x == 0. ? 0. : std::log(y); };
947   auto h = [](float x, float y) -> float { return x == 0. ? 0. : x / y; };
948   SymGrad("Xlogy", x, y, &dx, &dy);
949   test::ExpectClose(
950       dx, test::AsTensor<float>({g(0.f, .5f), g(0.f, 0.f), g(2.f, .5f),
951                                  g(3.f, 2.f), g(4.f, 2.f), g(5.f, 2.f)},
952                                 TensorShape({2, 3})));
953   test::ExpectClose(
954       dy, test::AsTensor<float>({h(0.f, .5f) + h(0.f, 0.f) + h(2.f, .5f),
955                                  h(3.f, 2.f) + h(4.f, 2.f) + h(5.f, 2.f)},
956                                 TensorShape({2, 1})));
957 }
958 
TEST_F(MathGradTest,Xlog1py)959 TEST_F(MathGradTest, Xlog1py) {
960   auto x = test::AsTensor<float>({0.f, 0.f, 2.f, 3.f, 4.f, 5.f},
961                                  TensorShape({2, 3}));
962   auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
963   Tensor dx;
964   Tensor dy;
965   auto g = [](float x, float y) -> float {
966     return x == 0. ? 0. : std::log1p(y);
967   };
968   auto h = [](float x, float y) -> float {
969     return x == 0. ? 0. : x / (y + 1.);
970   };
971   SymGrad("Xlog1py", x, y, &dx, &dy);
972   test::ExpectClose(
973       dx, test::AsTensor<float>({g(0.f, .5f), g(0.f, 0.f), g(2.f, .5f),
974                                  g(3.f, 2.f), g(4.f, 2.f), g(5.f, 2.f)},
975                                 TensorShape({2, 3})));
976   test::ExpectClose(
977       dy, test::AsTensor<float>({h(0.f, .5f) + h(0.f, 0.f) + h(2.f, .5f),
978                                  h(3.f, 2.f) + h(4.f, 2.f) + h(5.f, 2.f)},
979                                 TensorShape({2, 1})));
980 }
981 
TEST_F(MathGradTest,Xdivy)982 TEST_F(MathGradTest, Xdivy) {
983   auto x = test::AsTensor<float>({0.f, 0.f, 2.f, 3.f, 4.f, 5.f},
984                                  TensorShape({2, 3}));
985   auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
986   Tensor dx;
987   Tensor dy;
988   auto g = [](float x, float y) -> float { return x == 0. ? 0. : 1 / y; };
989   auto h = [](float x, float y) -> float {
990     return x == 0. ? 0. : -x / (y * y);
991   };
992   SymGrad("Xdivy", x, y, &dx, &dy);
993   test::ExpectClose(
994       dx, test::AsTensor<float>({g(0.f, .5f), g(0.f, 0.f), g(2.f, .5f),
995                                  g(3.f, 2.f), g(4.f, 2.f), g(5.f, 2.f)},
996                                 TensorShape({2, 3})));
997   test::ExpectClose(
998       dy, test::AsTensor<float>({h(0.f, .5f) + h(0.f, 0.f) + h(2.f, .5f),
999                                  h(3.f, 2.f) + h(4.f, 2.f) + h(5.f, 2.f)},
1000                                 TensorShape({2, 1})));
1001 }
1002 
TEST_F(MathGradTest,SquaredDifference)1003 TEST_F(MathGradTest, SquaredDifference) {
1004   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1005                                  TensorShape({2, 3}));
1006   auto y = test::AsTensor<float>({.5f, 2.f}, TensorShape({2, 1}));
1007   Tensor dx;
1008   Tensor dy;
1009   auto g = [](float x, float y) -> float { return 2. * (x - y); };
1010   auto h = [](float x, float y) -> float { return 2. * (y - x); };
1011   SymGrad("SquaredDifference", x, y, &dx, &dy);
1012   test::ExpectClose(
1013       dx, test::AsTensor<float>({g(-3.f, .5f), g(-2.f, .5f), g(-1.f, .5f),
1014                                  g(1.f, 2.f), g(2.f, 2.f), g(3.f, 2.f)},
1015                                 TensorShape({2, 3})));
1016   test::ExpectClose(
1017       dy, test::AsTensor<float>({h(-3.f, .5f) + h(-2.f, .5f) + h(-1.f, .5f),
1018                                  h(1.f, 2.f) + h(2.f, 2.f) + h(3.f, 2.f)},
1019                                 TensorShape({2, 1})));
1020 }
1021 
TEST_F(MathGradTest,Maximum)1022 TEST_F(MathGradTest, Maximum) {
1023   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1024                                  TensorShape({2, 3}));
1025   auto y = test::AsTensor<float>({-1.5f, 1.5f}, TensorShape({2, 1}));
1026   Tensor dx;
1027   Tensor dy;
1028   {
1029     SymGrad("Maximum", x, y, &dx, &dy);
1030     {
1031       auto g = [](float x, float y) { return x >= y ? 1.f : 0.f; };
1032       test::ExpectClose(dx, test::AsTensor<float>(
1033                                 {g(-3.f, -1.5f), g(-2.f, -1.5f), g(-1.f, -1.5f),
1034                                  g(1.f, 1.5f), g(2.f, 1.5f), g(3.f, 1.5f)},
1035                                 TensorShape({2, 3})));
1036     }
1037     {
1038       auto g = [](float x, float y) { return x < y ? 1.f : 0.f; };
1039       test::ExpectClose(dy,
1040                         test::AsTensor<float>(
1041                             {g(-3.f, -1.5f) + g(-2.f, -1.5f) + g(-1.f, -1.5f),
1042                              g(1.f, 1.5f) + g(2.f, 1.5f) + g(3.f, 1.5f)},
1043                             TensorShape({2, 1})));
1044     }
1045   }
1046   {  // Swap x and y
1047     SymGrad("Maximum", y, x, &dy, &dx);
1048     {
1049       auto g = [](float x, float y) { return x >= y ? 1.f : 0.f; };
1050       test::ExpectClose(dy,
1051                         test::AsTensor<float>(
1052                             {g(-1.5f, -3.f) + g(-1.5f, -2.f) + g(-1.5f, -1.f),
1053                              g(1.5f, 1.f) + g(1.5f, 2.f) + g(1.5f, 3.f)},
1054                             TensorShape({2, 1})));
1055     }
1056     {
1057       auto g = [](float x, float y) { return x < y ? 1.f : 0.f; };
1058       test::ExpectClose(dx, test::AsTensor<float>(
1059                                 {g(-1.5f, -3.f), g(-1.5f, -2.f), g(-1.5f, -1.f),
1060                                  g(1.5f, 1.f), g(1.5f, 2.f), g(1.5f, 3.f)},
1061                                 TensorShape({2, 3})));
1062     }
1063   }
1064 }
1065 
TEST_F(MathGradTest,Minimum)1066 TEST_F(MathGradTest, Minimum) {
1067   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1068                                  TensorShape({2, 3}));
1069   auto y = test::AsTensor<float>({-1.5f, 1.5f}, TensorShape({2, 1}));
1070   Tensor dx;
1071   Tensor dy;
1072   {
1073     SymGrad("Minimum", x, y, &dx, &dy);
1074     {
1075       auto g = [](float x, float y) { return x <= y ? 1.f : 0.f; };
1076       test::ExpectClose(dx, test::AsTensor<float>(
1077                                 {g(-3.f, -1.5f), g(-2.f, -1.5f), g(-1.f, -1.5f),
1078                                  g(1.f, 1.5f), g(2.f, 1.5f), g(3.f, 1.5f)},
1079                                 TensorShape({2, 3})));
1080     }
1081     {
1082       auto g = [](float x, float y) { return x > y ? 1.f : 0.f; };
1083       test::ExpectClose(dy,
1084                         test::AsTensor<float>(
1085                             {g(-3.f, -1.5f) + g(-2.f, -1.5f) + g(-1.f, -1.5f),
1086                              g(1.f, 1.5f) + g(2.f, 1.5f) + g(3.f, 1.5f)},
1087                             TensorShape({2, 1})));
1088     }
1089   }
1090   {  // Swap x and y
1091     SymGrad("Minimum", y, x, &dy, &dx);
1092     {
1093       auto g = [](float x, float y) { return x <= y ? 1.f : 0.f; };
1094       test::ExpectClose(dy,
1095                         test::AsTensor<float>(
1096                             {g(-1.5f, -3.f) + g(-1.5f, -2.f) + g(-1.5f, -1.f),
1097                              g(1.5f, 1.f) + g(1.5f, 2.f) + g(1.5f, 3.f)},
1098                             TensorShape({2, 1})));
1099     }
1100     {
1101       auto g = [](float x, float y) { return x > y ? 1.f : 0.f; };
1102       test::ExpectClose(dx, test::AsTensor<float>(
1103                                 {g(-1.5f, -3.f), g(-1.5f, -2.f), g(-1.5f, -1.f),
1104                                  g(1.5f, 1.f), g(1.5f, 2.f), g(1.5f, 3.f)},
1105                                 TensorShape({2, 3})));
1106     }
1107   }
1108 }
1109 
TEST_F(MathGradTest,Select)1110 TEST_F(MathGradTest, Select) {
1111   auto c = test::AsTensor<bool>({true, false, false, true, true, false},
1112                                 TensorShape({2, 3}));
1113   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1114                                  TensorShape({2, 3}));
1115   auto y = test::AsTensor<float>({3.f, 2.f, 1.f, 1.f, 2.f, 3.f},
1116                                  TensorShape({2, 3}));
1117   Tensor dc;
1118   Tensor dx;
1119   Tensor dy;
1120   {
1121     SelectGrad(c, x, y, &dc, &dx, &dy);
1122     test::ExpectTensorEqual<bool>(
1123         dc, test::AsTensor<bool>({false, false, false, false, false, false},
1124                                  TensorShape({2, 3})));
1125     test::ExpectTensorEqual<float>(
1126         dx, test::AsTensor<float>({1.f, 0.f, 0.f, 1.f, 1.f, 0.f},
1127                                   TensorShape({2, 3})));
1128     test::ExpectTensorEqual<float>(
1129         dy, test::AsTensor<float>({0.f, 1.f, 1.f, 0.f, 0.f, 1.f},
1130                                   TensorShape({2, 3})));
1131   }
1132 }
1133 
TEST_F(MathGradTest,MatMul_00)1134 TEST_F(MathGradTest, MatMul_00) {
1135   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1136                                  TensorShape({2, 3}));
1137   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({3, 1}));
1138   Tensor dx;
1139   Tensor dy;
1140   MatMulGrad(x, false, y, false, &dx, &dy);
1141   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({2, 1}));
1142   test::ExpectClose(dx, MatMul(dz, false, y, true));
1143   test::ExpectClose(dy, MatMul(x, true, dz, false));
1144 }
1145 
TEST_F(MathGradTest,MatMul_01)1146 TEST_F(MathGradTest, MatMul_01) {
1147   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1148                                  TensorShape({2, 3}));
1149   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 3}));
1150   Tensor dx;
1151   Tensor dy;
1152   MatMulGrad(x, false, y, true, &dx, &dy);
1153   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({2, 1}));
1154   test::ExpectClose(dx, MatMul(dz, false, y, false));
1155   test::ExpectClose(dy, MatMul(dz, true, x, false));
1156 }
1157 
TEST_F(MathGradTest,MatMul_10)1158 TEST_F(MathGradTest, MatMul_10) {
1159   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1160                                  TensorShape({3, 2}));
1161   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({3, 1}));
1162   Tensor dx;
1163   Tensor dy;
1164   MatMulGrad(x, true, y, false, &dx, &dy);
1165   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({2, 1}));
1166   test::ExpectClose(dx, MatMul(y, false, dz, true));
1167   test::ExpectClose(dy, MatMul(x, false, dz, false));
1168 }
1169 
TEST_F(MathGradTest,MatMul_11)1170 TEST_F(MathGradTest, MatMul_11) {
1171   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1172                                  TensorShape({3, 2}));
1173   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 3}));
1174   Tensor dx;
1175   Tensor dy;
1176   MatMulGrad(x, true, y, true, &dx, &dy);
1177   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({2, 1}));
1178   test::ExpectClose(dx, MatMul(y, true, dz, true));
1179   test::ExpectClose(dy, MatMul(dz, true, x, true));
1180 }
1181 
TEST_F(MathGradTest,BatchMatMul_00)1182 TEST_F(MathGradTest, BatchMatMul_00) {
1183   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1184                                  TensorShape({1, 2, 3}));
1185   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 3, 1}));
1186   Tensor dx;
1187   Tensor dy;
1188   BatchMatMulGrad(x, false, y, false, &dx, &dy);
1189   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1190   test::ExpectClose(dx, BatchMatMul(dz, false, y, true));
1191   test::ExpectClose(dy, BatchMatMul(x, true, dz, false));
1192 }
1193 
TEST_F(MathGradTest,BatchMatMul_01)1194 TEST_F(MathGradTest, BatchMatMul_01) {
1195   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1196                                  TensorShape({1, 2, 3}));
1197   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 1, 3}));
1198   Tensor dx;
1199   Tensor dy;
1200   BatchMatMulGrad(x, false, y, true, &dx, &dy);
1201   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1202   test::ExpectClose(dx, BatchMatMul(dz, false, y, false));
1203   test::ExpectClose(dy, BatchMatMul(dz, true, x, false));
1204 }
1205 
TEST_F(MathGradTest,BatchMatMul_10)1206 TEST_F(MathGradTest, BatchMatMul_10) {
1207   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1208                                  TensorShape({1, 3, 2}));
1209   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 3, 1}));
1210   Tensor dx;
1211   Tensor dy;
1212   BatchMatMulGrad(x, true, y, false, &dx, &dy);
1213   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1214   test::ExpectClose(dx, BatchMatMul(y, false, dz, true));
1215   test::ExpectClose(dy, BatchMatMul(x, false, dz, false));
1216 }
1217 
TEST_F(MathGradTest,BatchMatMul_11)1218 TEST_F(MathGradTest, BatchMatMul_11) {
1219   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1220                                  TensorShape({1, 3, 2}));
1221   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 1, 3}));
1222   Tensor dx;
1223   Tensor dy;
1224   BatchMatMulGrad(x, true, y, true, &dx, &dy);
1225   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1226   test::ExpectClose(dx, BatchMatMul(y, true, dz, true));
1227   test::ExpectClose(dy, BatchMatMul(dz, true, x, true));
1228 }
1229 
TEST_F(MathGradTest,BatchMatMulV2_00)1230 TEST_F(MathGradTest, BatchMatMulV2_00) {
1231   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1232                                  TensorShape({1, 2, 3}));
1233   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 3, 1}));
1234   Tensor dx;
1235   Tensor dy;
1236   BatchMatMulV2Grad(x, false, y, false, &dx, &dy);
1237   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1238   test::ExpectClose(dx, BatchMatMulV2(dz, false, y, true));
1239   test::ExpectClose(dy, BatchMatMulV2(x, true, dz, false));
1240 }
1241 
TEST_F(MathGradTest,BatchMatMulV2_01)1242 TEST_F(MathGradTest, BatchMatMulV2_01) {
1243   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1244                                  TensorShape({1, 2, 3}));
1245   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 1, 3}));
1246   Tensor dx;
1247   Tensor dy;
1248   BatchMatMulV2Grad(x, false, y, true, &dx, &dy);
1249   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1250   test::ExpectClose(dx, BatchMatMulV2(dz, false, y, false));
1251   test::ExpectClose(dy, BatchMatMulV2(dz, true, x, false));
1252 }
1253 
TEST_F(MathGradTest,BatchMatMulV2_10)1254 TEST_F(MathGradTest, BatchMatMulV2_10) {
1255   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1256                                  TensorShape({1, 3, 2}));
1257   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 3, 1}));
1258   Tensor dx;
1259   Tensor dy;
1260   BatchMatMulV2Grad(x, true, y, false, &dx, &dy);
1261   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1262   test::ExpectClose(dx, BatchMatMulV2(y, false, dz, true));
1263   test::ExpectClose(dy, BatchMatMulV2(x, false, dz, false));
1264 }
1265 
TEST_F(MathGradTest,BatchMatMulV2_11)1266 TEST_F(MathGradTest, BatchMatMulV2_11) {
1267   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1268                                  TensorShape({1, 3, 2}));
1269   auto y = test::AsTensor<float>({-1.f, .5f, 2.f}, TensorShape({1, 1, 3}));
1270   Tensor dx;
1271   Tensor dy;
1272   BatchMatMulV2Grad(x, true, y, true, &dx, &dy);
1273   auto dz = test::AsTensor<float>({1.f, 1.f}, TensorShape({1, 2, 1}));
1274   test::ExpectClose(dx, BatchMatMulV2(y, true, dz, true));
1275   test::ExpectClose(dy, BatchMatMulV2(dz, true, x, true));
1276 }
1277 
TEST_F(MathGradTest,BatchMatMulV2_LhsBroadcasts)1278 TEST_F(MathGradTest, BatchMatMulV2_LhsBroadcasts) {
1279   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1280                                  TensorShape({2, 3}));
1281   auto y = test::AsTensor<float>(
1282       {1.f, 2.4, 3.f, -1.f, .5f, 2.f, 3.f, 1.f, -1.f, 2.f, -.1f, 0},
1283       TensorShape({2, 3, 2}));
1284   Tensor dx;
1285   Tensor dy;
1286   BatchMatMulV2Grad(x, false, y, false, &dx, &dy);
1287   EXPECT_TRUE(dx.shape().IsSameSize(x.shape()));
1288   EXPECT_TRUE(dy.shape().IsSameSize(y.shape()));
1289   auto dz = test::AsTensor<float>({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f},
1290                                   TensorShape({2, 2, 2}));
1291   Tensor ans_dx;
1292   CHECK(ans_dx.CopyFrom(ReduceSum(BatchMatMulV2(dz, false, y, true), {0}),
1293                         dx.shape()));
1294   Tensor ans_dy = BatchMatMulV2(x, true, dz, false);
1295   test::ExpectClose(dx, ans_dx);
1296   test::ExpectClose(dy, ans_dy);
1297 }
1298 
TEST_F(MathGradTest,BatchMatMulV2_RhsBroadcasts)1299 TEST_F(MathGradTest, BatchMatMulV2_RhsBroadcasts) {
1300   auto x = test::AsTensor<float>(
1301       {1.f, 2.4, 3.f, -1.f, .5f, 2.f, 3.f, 1.f, -1.f, 2.f, -.1f, 0},
1302       TensorShape({2, 2, 3}));
1303   auto y = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1304                                  TensorShape({3, 2}));
1305   Tensor dx;
1306   Tensor dy;
1307   BatchMatMulV2Grad(x, false, y, false, &dx, &dy);
1308   auto dz = test::AsTensor<float>({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f},
1309                                   TensorShape({2, 2, 2}));
1310   Tensor ans_dx = BatchMatMulV2(dz, false, y, true);
1311   Tensor ans_dy;
1312   CHECK(ans_dy.CopyFrom(ReduceSum(BatchMatMulV2(x, true, dz, false), {0}),
1313                         dy.shape()));
1314   test::ExpectClose(dx, ans_dx);
1315   test::ExpectClose(dy, ans_dy);
1316 }
1317 
TEST_F(MathGradTest,BatchMatMulV2_BothLhsAndRhsBroadcast)1318 TEST_F(MathGradTest, BatchMatMulV2_BothLhsAndRhsBroadcast) {
1319   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1320                                  TensorShape({2, 1, 1, 3}));
1321   auto y = test::AsTensor<float>({3.f, 1.f, -1.f, 2.f, -.1f, 0},
1322                                  TensorShape({1, 2, 3, 1}));
1323   Tensor dx;
1324   Tensor dy;
1325   BatchMatMulV2Grad(x, false, y, false, &dx, &dy);
1326   EXPECT_TRUE(dx.shape().IsSameSize(x.shape()));
1327   EXPECT_TRUE(dy.shape().IsSameSize(y.shape()));
1328   auto dz =
1329       test::AsTensor<float>({1.f, 1.f, 1.f, 1.f}, TensorShape({2, 2, 1, 1}));
1330   Tensor ans_dx;
1331   Tensor ans_dy;
1332   CHECK(ans_dx.CopyFrom(ReduceSum(BatchMatMulV2(dz, false, y, true), {1}),
1333                         dx.shape()));
1334   CHECK(ans_dy.CopyFrom(ReduceSum(BatchMatMulV2(x, true, dz, false), {0}),
1335                         dy.shape()));
1336   test::ExpectClose(dx, ans_dx);
1337   test::ExpectClose(dy, ans_dy);
1338 }
1339 
TEST_F(MathGradTest,BatchMatMulV2_BroadcastWhileAdjointed)1340 TEST_F(MathGradTest, BatchMatMulV2_BroadcastWhileAdjointed) {
1341   auto x = test::AsTensor<float>({1.f, 2.f, 3.f, 4.f, 5.f, 6.f},
1342                                  TensorShape({2, 1, 3, 1}));
1343   auto y = test::AsTensor<float>({3.f, 1.f, -1.f, 2.f, -.1f, 0},
1344                                  TensorShape({1, 2, 1, 3}));
1345   Tensor dx;
1346   Tensor dy;
1347   BatchMatMulV2Grad(x, true, y, true, &dx, &dy);
1348   EXPECT_TRUE(dx.shape().IsSameSize(x.shape()));
1349   EXPECT_TRUE(dy.shape().IsSameSize(y.shape()));
1350 
1351   auto dz =
1352       test::AsTensor<float>({1.f, 1.f, 1.f, 1.f}, TensorShape({2, 2, 1, 1}));
1353   Tensor ans_dx;
1354   Tensor ans_dy;
1355   CHECK(ans_dx.CopyFrom(ReduceSum(BatchMatMulV2(y, true, dz, true), {1}),
1356                         dx.shape()));
1357   CHECK(ans_dy.CopyFrom(ReduceSum(BatchMatMulV2(dz, true, x, true), {0}),
1358                         dy.shape()));
1359   test::ExpectClose(dx, ans_dx);
1360   test::ExpectClose(dy, ans_dy);
1361 }
1362 
TEST_F(MathGradTest,Sum_dim0)1363 TEST_F(MathGradTest, Sum_dim0) {
1364   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1365                                  TensorShape({2, 3}));
1366   auto i = test::AsTensor<int32>({0}, TensorShape({}));
1367   Tensor dx;
1368   Tensor di;
1369   ReductionGrad("Sum", x, i, &dx, &di);
1370   test::ExpectTensorEqual<float>(
1371       dx, test::AsTensor<float>({1.f, 1.f, 1.f, 1.f, 1.f, 1.f},
1372                                 TensorShape({2, 3})));
1373   test::ExpectTensorEqual<int32>(di,
1374                                  test::AsTensor<int32>({0}, TensorShape({})));
1375 }
1376 
TEST_F(MathGradTest,Sum_dim1)1377 TEST_F(MathGradTest, Sum_dim1) {
1378   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1379                                  TensorShape({2, 3}));
1380   auto i = test::AsTensor<int32>({1}, TensorShape({}));
1381   Tensor dx;
1382   Tensor di;
1383   ReductionGrad("Sum", x, i, &dx, &di);
1384   test::ExpectTensorEqual<float>(
1385       dx, test::AsTensor<float>({1.f, 1.f, 1.f, 1.f, 1.f, 1.f},
1386                                 TensorShape({2, 3})));
1387   test::ExpectTensorEqual<int32>(di,
1388                                  test::AsTensor<int32>({0}, TensorShape({})));
1389 }
1390 
TEST_F(MathGradTest,Mean_dim0)1391 TEST_F(MathGradTest, Mean_dim0) {
1392   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1393                                  TensorShape({2, 3}));
1394   auto i = test::AsTensor<int32>({0}, TensorShape({}));
1395   Tensor dx;
1396   Tensor di;
1397   ReductionGrad("Mean", x, i, &dx, &di);
1398   test::ExpectTensorEqual<float>(
1399       dx, test::AsTensor<float>(
1400               {1.f / 2, 1.f / 2, 1.f / 2, 1.f / 2, 1.f / 2, 1.f / 2},
1401               TensorShape({2, 3})));
1402   test::ExpectTensorEqual<int32>(di,
1403                                  test::AsTensor<int32>({0}, TensorShape({})));
1404 }
1405 
TEST_F(MathGradTest,Mean_dim1)1406 TEST_F(MathGradTest, Mean_dim1) {
1407   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1408                                  TensorShape({2, 3}));
1409   auto i = test::AsTensor<int32>({1}, TensorShape({}));
1410   Tensor dx;
1411   Tensor di;
1412   ReductionGrad("Mean", x, i, &dx, &di);
1413   test::ExpectTensorEqual<float>(
1414       dx, test::AsTensor<float>(
1415               {1.f / 3, 1.f / 3, 1.f / 3, 1.f / 3, 1.f / 3, 1.f / 3},
1416               TensorShape({2, 3})));
1417   test::ExpectTensorEqual<int32>(di,
1418                                  test::AsTensor<int32>({0}, TensorShape({})));
1419 }
1420 
TEST_F(MathGradTest,Mean_dim0_dim1)1421 TEST_F(MathGradTest, Mean_dim0_dim1) {
1422   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1423                                  TensorShape({2, 3}));
1424   auto i = test::AsTensor<int32>({0, 1}, TensorShape({2}));
1425   Tensor dx;
1426   Tensor di;
1427   ReductionGrad("Mean", x, i, &dx, &di);
1428   test::ExpectTensorEqual<float>(
1429       dx, test::AsTensor<float>(
1430               {1.f / 6, 1.f / 6, 1.f / 6, 1.f / 6, 1.f / 6, 1.f / 6},
1431               TensorShape({2, 3})));
1432   test::ExpectTensorEqual<int32>(
1433       di, test::AsTensor<int32>({0, 0}, TensorShape({2})));
1434 }
1435 
TEST_F(MathGradTest,Min_dim0)1436 TEST_F(MathGradTest, Min_dim0) {
1437   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1438                                  TensorShape({2, 3}));
1439   auto i = test::AsTensor<int32>({0}, TensorShape({}));
1440   Tensor dx;
1441   Tensor di;
1442   ReductionGrad("Min", x, i, &dx, &di);
1443   test::ExpectTensorEqual<float>(
1444       dx, test::AsTensor<float>({1.f, 1.f, 1.f, 0.f, 0.f, 0.f},
1445                                 TensorShape({2, 3})));
1446   test::ExpectTensorEqual<int32>(di,
1447                                  test::AsTensor<int32>({0}, TensorShape({})));
1448 }
1449 
TEST_F(MathGradTest,Min_dim1)1450 TEST_F(MathGradTest, Min_dim1) {
1451   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1452                                  TensorShape({2, 3}));
1453   auto i = test::AsTensor<int32>({1}, TensorShape({}));
1454   Tensor dx;
1455   Tensor di;
1456   ReductionGrad("Min", x, i, &dx, &di);
1457   test::ExpectTensorEqual<float>(
1458       dx, test::AsTensor<float>({1.f, 0.f, 0.f, 1.f, 0.f, 0.f},
1459                                 TensorShape({2, 3})));
1460   test::ExpectTensorEqual<int32>(di,
1461                                  test::AsTensor<int32>({0}, TensorShape({})));
1462 }
1463 
TEST_F(MathGradTest,Min_dim0_dim1)1464 TEST_F(MathGradTest, Min_dim0_dim1) {
1465   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1466                                  TensorShape({2, 3}));
1467   auto i = test::AsTensor<int32>({0, 1}, TensorShape({2}));
1468   Tensor dx;
1469   Tensor di;
1470   ReductionGrad("Min", x, i, &dx, &di);
1471   test::ExpectTensorEqual<float>(
1472       dx, test::AsTensor<float>({1.f, 0.f, 0.f, 0.f, 0.f, 0.f},
1473                                 TensorShape({2, 3})));
1474   test::ExpectTensorEqual<int32>(
1475       di, test::AsTensor<int32>({0, 0}, TensorShape({2})));
1476 }
1477 
TEST_F(MathGradTest,Min_dim0_dim1_Dups)1478 TEST_F(MathGradTest, Min_dim0_dim1_Dups) {
1479   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, -3.f},
1480                                  TensorShape({2, 3}));
1481   auto i = test::AsTensor<int32>({0, 1}, TensorShape({2}));
1482   Tensor dx;
1483   Tensor di;
1484   ReductionGrad("Min", x, i, &dx, &di);
1485   test::ExpectTensorEqual<float>(
1486       dx, test::AsTensor<float>({.5f, 0.f, 0.f, 0.f, 0.f, .5f},
1487                                 TensorShape({2, 3})));
1488   test::ExpectTensorEqual<int32>(
1489       di, test::AsTensor<int32>({0, 0}, TensorShape({2})));
1490 }
1491 
TEST_F(MathGradTest,Max_dim0)1492 TEST_F(MathGradTest, Max_dim0) {
1493   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1494                                  TensorShape({2, 3}));
1495   auto i = test::AsTensor<int32>({0}, TensorShape({}));
1496   Tensor dx;
1497   Tensor di;
1498   ReductionGrad("Max", x, i, &dx, &di);
1499   LOG(INFO) << dx.SummarizeValue(6);
1500   test::ExpectTensorEqual<float>(
1501       dx, test::AsTensor<float>({0.f, 0.f, 0.f, 1.f, 1.f, 1.f},
1502                                 TensorShape({2, 3})));
1503   test::ExpectTensorEqual<int32>(di,
1504                                  test::AsTensor<int32>({0}, TensorShape({})));
1505 }
1506 
TEST_F(MathGradTest,Max_dim1)1507 TEST_F(MathGradTest, Max_dim1) {
1508   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1509                                  TensorShape({2, 3}));
1510   auto i = test::AsTensor<int32>({1}, TensorShape({}));
1511   Tensor dx;
1512   Tensor di;
1513   ReductionGrad("Max", x, i, &dx, &di);
1514   test::ExpectTensorEqual<float>(
1515       dx, test::AsTensor<float>({0.f, 0.f, 1.f, 0.f, 0.f, 1.f},
1516                                 TensorShape({2, 3})));
1517   test::ExpectTensorEqual<int32>(di,
1518                                  test::AsTensor<int32>({0}, TensorShape({})));
1519 }
1520 
TEST_F(MathGradTest,Max_dim0_dim1)1521 TEST_F(MathGradTest, Max_dim0_dim1) {
1522   auto x = test::AsTensor<float>({-3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1523                                  TensorShape({2, 3}));
1524   auto i = test::AsTensor<int32>({0, 1}, TensorShape({2}));
1525   Tensor dx;
1526   Tensor di;
1527   ReductionGrad("Max", x, i, &dx, &di);
1528   test::ExpectTensorEqual<float>(
1529       dx, test::AsTensor<float>({0.f, 0.f, 0.f, 0.f, 0.f, 1.f},
1530                                 TensorShape({2, 3})));
1531   test::ExpectTensorEqual<int32>(
1532       di, test::AsTensor<int32>({0, 0}, TensorShape({2})));
1533 }
1534 
TEST_F(MathGradTest,Max_dim0_dim1_Dups)1535 TEST_F(MathGradTest, Max_dim0_dim1_Dups) {
1536   auto x = test::AsTensor<float>({3.f, -2.f, -1.f, 1.f, 2.f, 3.f},
1537                                  TensorShape({2, 3}));
1538   auto i = test::AsTensor<int32>({0, 1}, TensorShape({2}));
1539   Tensor dx;
1540   Tensor di;
1541   ReductionGrad("Max", x, i, &dx, &di);
1542   test::ExpectTensorEqual<float>(
1543       dx, test::AsTensor<float>({.5f, 0.f, 0.f, 0.f, 0.f, .5f},
1544                                 TensorShape({2, 3})));
1545   test::ExpectTensorEqual<int32>(
1546       di, test::AsTensor<int32>({0, 0}, TensorShape({2})));
1547 }
1548 
1549 }  // namespace
1550 }  // namespace tensorflow
1551