Function XLA_TEST_P

tensorflow/compiler/xla/tests/triangular_solve_test.cc:453–485

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

XLA_TEST_P(TriangularSolveParametricTest, Random) {
  TriangularSolveTestSpec spec = GetParam();

  XlaBuilder builder(TestName());

  Array2D<float> avals(spec.m, spec.m);
  avals.FillRandom(1.0);
  for (int i = 0; i < spec.m; ++i) {
    avals(i, i) += 10;
  }

  std::pair<int, int> bdims = spec.left_side ? std::make_pair(spec.m, spec.n)
                                             : std::make_pair(spec.n, spec.m);
  Array2D<float> bvals(bdims.first, bdims.second);
  bvals.FillRandom(1.0);

  XlaOp a, b;
  auto a_data = CreateR2Parameter<float>(avals, 0, "a", &builder, &a);
  auto b_data = CreateR2Parameter<float>(bvals, 1, "b", &builder, &b);
  auto x = TriangularSolve(a, b, spec.left_side, spec.lower,
                           /*unit_diagonal=*/false, spec.transpose_a);
  auto a_tri = Triangle(a, spec.lower);
  a_tri = MaybeTransposeInMinorDims(
      a_tri, spec.transpose_a != TriangularSolveOptions::NO_TRANSPOSE);
  if (spec.left_side) {
    BatchDot(a_tri, x);
  } else {
    BatchDot(x, a_tri);
  }

  ComputeAndCompareR2<float>(&builder, bvals, {a_data.get(), b_data.get()},
                             ErrorSpec(1e-2, 1e-2));
}

std::vector<TriangularSolveTestSpec> TriangularSolveTests() {
  std::vector<TriangularSolveTestSpec> specs;