Function TEST_P

tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc:321–359

}

TEST_P(ArrayElementwiseOpTestParamCount, AddManyValues) {
  const int count = GetParam();
  XlaBuilder builder(TestName());
  std::vector<float> a_values;
  std::vector<float> b_values;
  for (int i = 0; i < count; ++i) {
    a_values.push_back(i / static_cast<float>(count));
    b_values.push_back(2 * i / static_cast<float>(count + 2));
  }

  Literal a_literal = LiteralUtil::CreateR1<float>({a_values});
  std::unique_ptr<GlobalData> a_data =
      client_->TransferToServer(a_literal).ConsumeValueOrDie();
  auto a_constant = ConstantR1<float>(&builder, a_values);
  auto a_param = Parameter(&builder, 0, a_literal.shape(), "a_param");

  Literal b_literal = LiteralUtil::CreateR1<float>({b_values});
  std::unique_ptr<GlobalData> b_data =
      client_->TransferToServer(b_literal).ConsumeValueOrDie();
  auto b_param = Parameter(&builder, 1, a_literal.shape(), "b_param");
  auto b_constant = ConstantR1<float>(&builder, b_values);

  auto sum1 = Add(a_constant, b_param);
  auto sum2 = Add(a_constant, b_constant);
  auto sum3 = Add(a_param, b_param);
  auto sum4 = Add(a_param, b_constant);

  auto sum = Add(sum1, sum2);
  sum = Add(sum, sum3);
  sum = Add(sum, sum4);

  std::vector<float> expected;
  for (int64 i = 0; i < count; ++i) {
    expected.push_back(4 * (a_values[i] + b_values[i]));
  }

  ComputeAndCompareR1<float>(&builder, expected, {a_data.get(), b_data.get()},
                             error_spec_);
}

XLA_TEST_F(ArrayElementwiseOpTest, DeeplyNestedAddWithSlices) {
  XlaBuilder builder(TestName());