| 211 | } |
| 212 | |
| 213 | std::vector<Tensor> SplitGrad(int dim, const Tensor& x, const Tensor& dy0, |
| 214 | const Tensor& dy1) { |
| 215 | auto T = DT_FLOAT; |
| 216 | auto gdef = test::function::GDef( |
| 217 | {f::NDef("dim", "Placeholder", {}, {{"dtype", DT_INT32}}), |
| 218 | f::NDef("x", "Placeholder", {}, {{"dtype", T}}), |
| 219 | f::NDef("dy0", "Placeholder", {}, {{"dtype", T}}), |
| 220 | f::NDef("dy1", "Placeholder", {}, {{"dtype", T}}), |
| 221 | f::NDef("dx", "SymbolicGradient", {"dim", "x", "dy0", "dy1"}, |
| 222 | {{"f", FDH::FunctionRef( |
| 223 | "Split", |
| 224 | {{"split_dim", dim}, {"num_split", 2}, {"T", T}})}, |
| 225 | {"Tin", DataTypeSlice{DT_INT32, T, T, T}}, |
| 226 | {"Tout", DataTypeSlice{DT_INT32, T}}})}); |
| 227 | VLOG(1) << DebugStringWhole(gdef); |
| 228 | auto sess = NewSession(); |
| 229 | TF_CHECK_OK(sess->Create(gdef)); |
| 230 | std::vector<Tensor> out; |
| 231 | TF_CHECK_OK(sess->Run({{"dim", test::AsScalar(dim)}, |
| 232 | {"x:0", x}, |
| 233 | {"dy0:0", dy0}, |
| 234 | {"dy1:0", dy1}}, |
| 235 | {"dx:0", "dx:1"}, {}, &out)); |
| 236 | CHECK_EQ(out.size(), 2); |
| 237 | TF_CHECK_OK(sess->Close()); |
| 238 | return out; |
| 239 | } |
| 240 | |
| 241 | std::vector<Tensor> SplitVGrad(const Tensor& x, const Tensor& size_splits, |
| 242 | int dim, const Tensor& dy0, const Tensor& dy1) { |
no test coverage detected