| 360 | } |
| 361 | |
| 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 | |
| 415 | void HasError(const Status& s, const string& substr) { |