Function SplitVGrad

tensorflow/core/ops/array_grad_test.cc:241–272

}

std::vector<Tensor> SplitVGrad(const Tensor& x, const Tensor& size_splits,
                               int dim, const Tensor& dy0, const Tensor& dy1) {
  auto T = DT_FLOAT;
  auto Tlen = DT_INT64;
  auto gdef = test::function::GDef(
      {f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
       f::NDef("size_splits", "Placeholder", {}, {{"dtype", Tlen}}),
       f::NDef("dim", "Placeholder", {}, {{"dtype", DT_INT32}}),
       f::NDef("dy0", "Placeholder", {}, {{"dtype", T}}),
       f::NDef("dy1", "Placeholder", {}, {{"dtype", T}}),
       f::NDef("dx", "SymbolicGradient",
               {"x", "size_splits", "dim", "dy0", "dy1"},
               {{"f", FDH::FunctionRef("SplitV", {{"split_dim", dim},
                                                  {"num_split", 2},
                                                  {"T", T},
                                                  {"Tlen", Tlen}})},
                {"Tin", DataTypeSlice{T, Tlen, DT_INT32, T, T}},
                {"Tout", DataTypeSlice{T, Tlen, DT_INT32}}})});
  VLOG(1) << DebugStringWhole(gdef);
  auto sess = NewSession();
  TF_CHECK_OK(sess->Create(gdef));
  std::vector<Tensor> out;
  TF_CHECK_OK(sess->Run({{"x:0", x},
                         {"size_splits:0", size_splits},
                         {"dim", test::AsScalar(dim)},
                         {"dy0:0", dy0},
                         {"dy1:0", dy1}},
                        {"dx:0", "dx:1", "dx:2"}, {}, &out));
  CHECK_EQ(out.size(), 3);
  TF_CHECK_OK(sess->Close());
  return out;
}

TEST(ArrayGradTest, SplitGrad) {
  Tensor x(DT_FLOAT, {2, 4, 5});