| 282 | } |
| 283 | |
| 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]; |