Function XLA_TEST_P

tensorflow/compiler/xla/tests/cholesky_test.cc:183–218

      public ::testing::WithParamInterface<CholeskyTestCase> {};

XLA_TEST_P(RandomCholeskyTest, Random) {
  XlaBuilder builder(TestName());

  auto test_params = GetParam();
  std::vector<int64> dimensions = {std::get<0>(test_params),
                                   std::get<1>(test_params),
                                   std::get<1>(test_params)};
  bool lower = std::get<2>(test_params);
  Shape shape = ShapeUtil::MakeShape(F32, dimensions);
  TF_ASSERT_OK_AND_ASSIGN(
      auto literal, LiteralUtil::CreateRandomLiteral<F32>(shape, 0.0, 1.0));

  auto input = Parameter(&builder, 0, shape, "input");
  // Form a random positive definite matrix.
  auto matrix =
      BatchDot(input, TransposeInMinorDims(input), PrecisionConfig::HIGHEST);

  auto cholesky = Triangle(Cholesky(matrix, lower), lower);

  // Verify that ||matrix - cholesky * cholesky_t||_2 ~= 0
  XlaOp verification;
  if (lower) {
    verification = BatchDot(cholesky, TransposeInMinorDims(cholesky),
                            PrecisionConfig::HIGHEST);
  } else {
    verification = BatchDot(TransposeInMinorDims(cholesky), cholesky,
                            PrecisionConfig::HIGHEST);
  }
  auto delta = matrix - verification;
  Reduce(delta * delta, ConstantR0<float>(&builder, 0.0),
         CreateScalarAddComputation(F32, &builder), {0, 1, 2});

  TF_ASSERT_OK_AND_ASSIGN(auto input_data, client_->TransferToServer(literal));
  ComputeAndCompareR0<float>(&builder, 0.0, {input_data.get()},
                             ErrorSpec(1e-4, 1e-4));
}

INSTANTIATE_TEST_SUITE_P(RandomCholeskyTestInstance, RandomCholeskyTest,
                         ::testing::Values(CholeskyTestCase{1, 1, true},