Function XLA_TEST_P

tensorflow/compiler/xla/tests/conditional_test.cc:192–219

Test branch computations that do not take any parameters.

// Test branch computations that do not take any parameters.
XLA_TEST_P(CaseOpTest, Parameters0) {
  int num_branches = GetParam();
  for (int bi = -1; bi <= num_branches; ++bi) {
    SCOPED_TRACE(bi);
    XlaBuilder builder(TestName());
    XlaOp branch_index;
    auto branch_index_arg = CreateR0Parameter<int32>(bi, 0, "branch_index_arg",
                                                     &builder, &branch_index);
    auto operand = Tuple(&builder, {});

    std::vector<XlaOp> operands(num_branches, operand);
    std::vector<XlaComputation> branches;
    branches.reserve(num_branches);
    std::vector<const XlaComputation*> branches_p(num_branches);
    for (int i = 0; i < num_branches; ++i) {
      branches.emplace_back(
          CreateR0ConstantComputation(static_cast<float>(i) * 10));
      branches_p[i] = &branches[i];
    }
    Conditional(branch_index, branches_p, operands);

    float expected = 10 * static_cast<float>((bi < 0 || bi >= num_branches)
                                                 ? num_branches - 1
                                                 : bi);
    ComputeAndCompareR0<float>(&builder, expected, {branch_index_arg.get()},
                               error_spec_);
  }
}

// Test true and false computations that take in 1 parameter.
XLA_TEST_F(ConditionalOpTest, Parameters1) {