Function UnpackGrad

tensorflow/core/ops/array_grad_test.cc:77–102

}

std::vector<Tensor> UnpackGrad(const Tensor& x, const Tensor& dy0,
                               const Tensor& dy1, int axis) {
  auto T = DT_FLOAT;
  auto gdef = test::function::GDef(
      {f::NDef("x", "Placeholder", {}, {{"dtype", T}}),
       f::NDef("axis", "Placeholder", {}, {{"dtype", DT_INT32}}),
       f::NDef("dy0", "Placeholder", {}, {{"dtype", T}}),
       f::NDef("dy1", "Placeholder", {}, {{"dtype", T}}),
       f::NDef("dx", "SymbolicGradient", {"x", "dy0", "dy1"},
               {{"f", FDH::FunctionRef("Unpack",
                                       {{"num", 2}, {"T", T}, {"axis", axis}})},
                {"Tin", DataTypeSlice{T, T, T}},
                {"Tout", DataTypeSlice{T}}})});
  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},
                         {"axis:0", test::AsScalar(axis)},
                         {"dy0:0", dy0},
                         {"dy1:0", dy1}},
                        {"dx:0"}, {}, &out));
  CHECK_EQ(out.size(), 1);
  TF_CHECK_OK(sess->Close());
  return out;
}

TEST(ArrayGradTest, UnpackGrad) {
  Tensor x(DT_FLOAT, {2, 2, 3});