Function TEST

tensorflow/core/ops/array_grad_test.cc:63–75

}

TEST(ArrayGradTest, PackGrad) {
  Tensor x0(DT_FLOAT, {2, 3});
  x0.flat<float>().setZero();
  Tensor x1(DT_FLOAT, {2, 3});
  x1.flat<float>().setZero();
  Tensor dy(DT_FLOAT, {2, 2, 3});
  test::FillIota<float>(&dy, 0);
  auto dx = PackGrad(x0, x1, dy, 0);
  test::ExpectClose(dx[0],
                    test::AsTensor<float>({0., 1., 2., 3., 4., 5.}, {2, 3}));
  test::ExpectClose(dx[1],
                    test::AsTensor<float>({6., 7., 8., 9., 10., 11.}, {2, 3}));
}

std::vector<Tensor> UnpackGrad(const Tensor& x, const Tensor& dy0,
                               const Tensor& dy1, int axis) {