Reduction grad
| 172 | |
| 173 | // Reduction grad |
| 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 | |
| 228 | Tensor ReduceSum(const Tensor& x, gtl::ArraySlice<int32> axes) { |
| 229 | int num_axes = axes.length(); |