xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/conditional_test.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948)

/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <random>
#include <utility>

#include "tensorflow/compiler/xla/client/xla_builder.h"
#include "tensorflow/compiler/xla/client/xla_computation.h"
#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
#include "tensorflow/compiler/xla/tests/literal_test_util.h"
#include "tensorflow/compiler/xla/tests/test_macros.h"

namespace xla {
namespace {

class ConditionalOpTest : public ClientLibraryTestBase {
 protected:
  XlaComputation CreateR0ConstantComputation(float value) {
    XlaBuilder builder("Constant");
    Parameter(&builder, 0, empty_tuple_, "tuple");
    ConstantR0<float>(&builder, value);
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0IdentityComputation() {
    XlaBuilder builder("Identity");
    Parameter(&builder, 0, r0f32_, "x");
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateCeilComputation(const Shape& shape) {
    XlaBuilder builder("Ceil");
    auto param = Parameter(&builder, 0, shape, "param");
    Ceil(param);
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0CeilComputation() {
    return CreateCeilComputation(r0f32_);
  }

  XlaComputation CreateR1CeilComputation() {
    return CreateCeilComputation(r1s2f32_);
  }

  XlaComputation CreateFloorComputation(const Shape& shape) {
    XlaBuilder builder("Floor");
    auto param = Parameter(&builder, 0, shape, "param");
    Floor(param);
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0FloorComputation() {
    return CreateFloorComputation(r0f32_);
  }

  XlaComputation CreateR1FloorComputation() {
    return CreateFloorComputation(r1s2f32_);
  }

  XlaComputation CreateTupleCeilComputation(const std::string& computation_name,
                                            const Shape& tuple_shape) {
    XlaBuilder builder(computation_name);
    auto tuple = Parameter(&builder, 0, tuple_shape, "tuple");
    auto x = GetTupleElement(tuple, 0);
    auto y = GetTupleElement(tuple, 1);
    auto x_ceil = Ceil(x);
    auto y_ceil = Ceil(y);
    Tuple(&builder, {x_ceil, y_ceil});
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0TupleCeilComputation() {
    return CreateTupleCeilComputation("CeilR0", tuple_2_r0f32_);
  }

  XlaComputation CreateR1TupleCeilComputation() {
    return CreateTupleCeilComputation("CeilR1", tuple_2_r1s2f32_);
  }

  XlaComputation CreateTupleFloorComputation(
      const std::string& computation_name, const Shape& tuple_shape) {
    XlaBuilder builder(computation_name);
    auto tuple = Parameter(&builder, 0, tuple_shape, "tuple");
    auto x = GetTupleElement(tuple, 0);
    auto y = GetTupleElement(tuple, 1);
    auto x_floor = Floor(x);
    auto y_floor = Floor(y);
    Tuple(&builder, {x_floor, y_floor});
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0TupleFloorComputation() {
    return CreateTupleFloorComputation("FloorR0", tuple_2_r0f32_);
  }

  XlaComputation CreateR1TupleFloorComputation() {
    return CreateTupleFloorComputation("FloorR1", tuple_2_r1s2f32_);
  }

  XlaComputation CreateTupleAddComputation(const std::string& computation_name,
                                           const Shape& tuple_shape) {
    XlaBuilder builder(computation_name);
    auto tuple = Parameter(&builder, 0, tuple_shape, "tuple");
    auto x = GetTupleElement(tuple, 0);
    auto y = GetTupleElement(tuple, 1);
    Add(x, y);
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0TupleAddComputation() {
    return CreateTupleAddComputation("AddR0", tuple_2_r0f32_);
  }

  XlaComputation CreateR1TupleAddComputation() {
    return CreateTupleAddComputation("AddR1", tuple_2_r1s2f32_);
  }

  XlaComputation CreateTupleSubComputation(const std::string& computation_name,
                                           const Shape& tuple_shape) {
    XlaBuilder builder(computation_name);
    auto tuple = Parameter(&builder, 0, tuple_shape, "tuple");
    auto x = GetTupleElement(tuple, 0);
    auto y = GetTupleElement(tuple, 1);
    Sub(x, y);
    auto build_status = builder.Build();
    EXPECT_IS_OK(build_status.status());
    return std::move(build_status).value();
  }

  XlaComputation CreateR0TupleSubComputation() {
    return CreateTupleSubComputation("SubR0", tuple_2_r0f32_);
  }

  XlaComputation CreateR1TupleSubComputation() {
    return CreateTupleSubComputation("SubR1", tuple_2_r1s2f32_);
  }

  Shape r0f32_ = ShapeUtil::MakeShape(F32, {});
  Shape r1s2f32_ = ShapeUtil::MakeShape(F32, {2});
  Shape tuple_2_r0f32_ = ShapeUtil::MakeTupleShape(
      {ShapeUtil::MakeShape(F32, {}), ShapeUtil::MakeShape(F32, {})});
  Shape tuple_2_r1s2f32_ = ShapeUtil::MakeTupleShape(
      {ShapeUtil::MakeShape(F32, {2}), ShapeUtil::MakeShape(F32, {2})});
  Shape empty_tuple_ = ShapeUtil::MakeTupleShape({});
  ErrorSpec error_spec_{0.001};
};

// Test fixture to run indexed conditional (switch/case) tests with varying
// number of branches.
class CaseOpTest : public ConditionalOpTest,
                   public ::testing::WithParamInterface<int> {};

// Test true and false computations that do not take any parameters.
XLA_TEST_F(ConditionalOpTest, Parameters0) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(true, 0, "pred", &builder, &pred);
  auto operands = Tuple(&builder, {});
  auto true_computation = CreateR0ConstantComputation(56.0f);
  auto false_computation = CreateR0ConstantComputation(12.0f);
  Conditional(pred, operands, true_computation, operands, false_computation);

  ComputeAndCompareR0<float>(&builder, 56.0f, {pred_arg.get()}, error_spec_);
}

// 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_t>(
        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) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.0f);
  auto operand2 = ConstantR0<float>(&builder, 12.0f);
  auto identity = CreateR0IdentityComputation();
  Conditional(pred, operand1, identity, operand2, identity);

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test branch computations that take in 1 parameter.
XLA_TEST_P(CaseOpTest, Parameters1) {
  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_t>(
        bi, 0, "branch_index_arg", &builder, &branch_index);

    auto make_branch = [&builder, this](int i) {
      auto sb = builder.CreateSubBuilder(absl::StrCat("branch_", i));
      Add(ConstantR0<float>(sb.get(), static_cast<float>(i)),
          Parameter(sb.get(), 0, r0f32_, "p0"));
      return sb->BuildAndNoteError();
    };
    std::vector<XlaComputation> branches;
    branches.reserve(num_branches);
    std::vector<const XlaComputation*> branches_p(num_branches);
    std::vector<XlaOp> operands;
    operands.reserve(num_branches);
    std::vector<float> expecteds(num_branches);
    for (int i = 0; i < num_branches; ++i) {
      branches.emplace_back(make_branch(i));
      branches_p[i] = &branches[i];
      auto fi = static_cast<float>(i);
      operands.emplace_back(ConstantR0<float>(&builder, 10 * fi + 7));
      expecteds[i] = 10 * fi + 7 + fi;
    }

    Conditional(branch_index, branches_p, operands);
    float expected = (bi < 0 || bi >= num_branches)
                         ? expecteds[num_branches - 1]
                         : expecteds[bi];
    ComputeAndCompareR0<float>(&builder, expected, {branch_index_arg.get()},
                               error_spec_);
  }
}

// Test conditional with two different computations in the true and false cases
// that take in different arguments.
XLA_TEST_F(ConditionalOpTest, DiffComputationsDiffArgs) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.4f);
  auto operand2 = ConstantR0<float>(&builder, 12.6f);
  Conditional(pred, operand1, CreateR0CeilComputation(), operand2,
              CreateR0FloorComputation());

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test conditional with two different computations in the true and false cases
// that take in the same arguments.
XLA_TEST_F(ConditionalOpTest, DiffComputationsSameArg) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand = ConstantR0<float>(&builder, 12.6f);
  Conditional(pred, operand, CreateR0CeilComputation(), operand,
              CreateR0FloorComputation());

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test conditional with the same computation in the true and false cases but
// take in different arguments.
XLA_TEST_F(ConditionalOpTest, SameComputationDiffArgs) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.4f);
  auto operand2 = ConstantR0<float>(&builder, 12.6f);
  auto floor = CreateR0FloorComputation();
  Conditional(pred, operand1, floor, operand2, floor);

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test conditional with the same computation in the true and false cases that
// take in the same arguments.
XLA_TEST_F(ConditionalOpTest, SameComputationSameArg) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand = ConstantR0<float>(&builder, 12.6f);
  auto floor = CreateR0FloorComputation();
  Conditional(pred, operand, floor, operand, floor);

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test conditional with different instances of the same computation in the true
// and false cases.
XLA_TEST_F(ConditionalOpTest, SameComputationDiffInstances) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.4f);
  auto operand2 = ConstantR0<float>(&builder, 12.6f);
  Conditional(pred, operand1, CreateR0FloorComputation(), operand2,
              CreateR0FloorComputation());

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test the case when a call invokes a computation that contains a conditional.
XLA_TEST_F(ConditionalOpTest, ConditionalWithCall) {
  Shape r0bool = ShapeUtil::MakeShape(PRED, {});
  XlaBuilder inner_builder(TestName() + ".inner_conditional");
  auto pred_cond = Parameter(&inner_builder, 0, r0bool, "param0");
  auto true_operand = Parameter(&inner_builder, 1, r0f32_, "param1");
  auto false_operand = Parameter(&inner_builder, 2, r0f32_, "param2");
  Conditional(pred_cond, true_operand, CreateR0CeilComputation(), false_operand,
              CreateR0FloorComputation());
  auto inner_builder_result = inner_builder.Build().value();

  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.4f);
  auto operand2 = ConstantR0<float>(&builder, 12.6f);
  Call(&builder, inner_builder_result, {pred, operand1, operand2});

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test true and false computations that take in 2 parameters and predicate is
// true.
XLA_TEST_F(ConditionalOpTest, Parameters2TrueBranch) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(true, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.0f);
  auto operand2 = ConstantR0<float>(&builder, 12.0f);
  auto operands = Tuple(&builder, {operand1, operand2});
  Conditional(pred, operands, CreateR0TupleAddComputation(), operands,
              CreateR0TupleSubComputation());

  ComputeAndCompareR0<float>(&builder, 68.0f, {pred_arg.get()}, error_spec_);
}

// Test true and false computations that take in 2 parameters and predicate is
// false.
XLA_TEST_F(ConditionalOpTest, Parameters2FalseBranch) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 56.0f);
  auto operand2 = ConstantR0<float>(&builder, 12.0f);
  auto operands = Tuple(&builder, {operand1, operand2});
  Conditional(pred, operands, CreateR0TupleAddComputation(), operands,
              CreateR0TupleSubComputation());

  ComputeAndCompareR0<float>(&builder, 44.0f, {pred_arg.get()}, error_spec_);
}

// Test true and false computations that take in 2 array parameters and
// predicate is true.
XLA_TEST_F(ConditionalOpTest, Parameters2ArrayTrueBranch) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(true, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR1<float>(&builder, {24.0f, 56.0f});
  auto operand2 = ConstantR1<float>(&builder, {10.0f, 11.0f});
  auto operands = Tuple(&builder, {operand1, operand2});
  Conditional(pred, operands, CreateR1TupleAddComputation(), operands,
              CreateR1TupleSubComputation());

  ComputeAndCompareR1<float>(&builder, {34.0f, 67.0f}, {pred_arg.get()},
                             error_spec_);
}

// Test branch computations that take in 2 array parameters.
XLA_TEST_P(CaseOpTest, Parameters2Array) {
  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_t>(bi, 0, "pred", &builder, &branch_index);
    auto operand1 = ConstantR1<float>(&builder, {24.0f, 56.0f});
    auto operand2 = ConstantR1<float>(&builder, {10.0f, 11.0f});
    auto operands = Tuple(&builder, {operand1, operand2});
    auto make_branch = [&builder, this](int i) {
      auto sb = builder.CreateSubBuilder(absl::StrCat("branch_", i));
      auto p = Parameter(sb.get(), 0, tuple_2_r1s2f32_, "p0");
      Add(Mul(ConstantR0<float>(sb.get(), static_cast<float>(i)),
              GetTupleElement(p, 0)),
          GetTupleElement(p, 1));
      return sb->BuildAndNoteError();
    };
    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(make_branch(i));
      branches_p[i] = &branches[i];
    }
    Conditional(branch_index, branches_p,
                std::vector<XlaOp>(num_branches, operands));
    auto modified_bi = static_cast<float>(
        (bi < 0 || bi >= num_branches) ? num_branches - 1 : bi);
    ComputeAndCompareR1<float>(
        &builder, {24.0f * modified_bi + 10, 56.0f * modified_bi + 11},
        {branch_index_arg.get()}, error_spec_);
  }
}

INSTANTIATE_TEST_SUITE_P(CaseOpTest_Instantiation, CaseOpTest,
                         ::testing::Values(1, 2, 3, 4, 5));

// Test true and false computations that take in 2 array parameters and
// predicate is false.
XLA_TEST_F(ConditionalOpTest, Parameters2ArrayFalseBranch) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR1<float>(&builder, {24.0f, 56.0f});
  auto operand2 = ConstantR1<float>(&builder, {10.0f, 11.0f});
  auto operands = Tuple(&builder, {operand1, operand2});
  Conditional(pred, operands, CreateR1TupleAddComputation(), operands,
              CreateR1TupleSubComputation());

  ComputeAndCompareR1<float>(&builder, {14.0f, 45.0f}, {pred_arg.get()},
                             error_spec_);
}

// Test true and false computations that return a tuple of scalars.
XLA_TEST_F(ConditionalOpTest, ReturnTupleOfScalars) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operands = Tuple(&builder, {ConstantR0<float>(&builder, 12.2f),
                                   ConstantR0<float>(&builder, 25.6f)});
  Conditional(pred, operands, CreateR0TupleCeilComputation(), operands,
              CreateR0TupleFloorComputation());

  ComputeAndCompareTuple(
      &builder,
      LiteralUtil::MakeTupleFromSlices({LiteralUtil::CreateR0<float>(12.0f),
                                        LiteralUtil::CreateR0<float>(25.0f)}),
      {pred_arg.get()}, error_spec_);
}

// Test true and false computations that return a tuple of arrays.
XLA_TEST_F(ConditionalOpTest, ReturnTupleOfArrays) {
  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(true, 0, "pred", &builder, &pred);
  auto operands =
      Tuple(&builder, {ConstantR1<float>(&builder, {12.2f, 15.8f}),
                       ConstantR1<float>(&builder, {25.6f, 29.2f})});
  Conditional(pred, operands, CreateR1TupleCeilComputation(), operands,
              CreateR1TupleFloorComputation());

  ComputeAndCompareTuple(&builder,
                         LiteralUtil::MakeTupleFromSlices(
                             {LiteralUtil::CreateR1<float>({13.0f, 16.0f}),
                              LiteralUtil::CreateR1<float>({26.0f, 30.0f})}),
                         {pred_arg.get()}, error_spec_);
}

// Test true and false computations that return a tuple of a predicate, a
// scalar, and an array.
XLA_TEST_F(ConditionalOpTest, ReturnTupleofPredicateScalarArray) {
  XlaBuilder true_builder(TestName() + ".true");
  {
    Parameter(&true_builder, 0, empty_tuple_, "tuple");
    auto true_pred = ConstantR0<bool>(&true_builder, true);
    auto true_scalar = ConstantR0<float>(&true_builder, 12.2f);
    auto true_array = ConstantR1<float>(&true_builder, {12.8f, 14.6f});
    Tuple(&true_builder, {true_pred, true_scalar, true_array});
  }
  auto true_builder_result = true_builder.Build();
  EXPECT_IS_OK(true_builder_result.status());

  XlaBuilder false_builder(TestName() + ".false");
  {
    Parameter(&false_builder, 0, empty_tuple_, "tuple");
    auto false_pred = ConstantR0<bool>(&false_builder, false);
    auto false_scalar = ConstantR0<float>(&false_builder, 25.6f);
    auto false_array = ConstantR1<float>(&false_builder, {26.4f, 32.6f});
    Tuple(&false_builder, {false_pred, false_scalar, false_array});
  }
  auto false_builder_result = false_builder.Build();
  EXPECT_IS_OK(false_builder_result.status());

  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(true, 0, "pred", &builder, &pred);
  auto operands = Tuple(&builder, {});
  Conditional(pred, operands, std::move(true_builder_result).value(), operands,
              std::move(false_builder_result).value());

  ComputeAndCompareTuple(&builder,
                         LiteralUtil::MakeTupleFromSlices(
                             {LiteralUtil::CreateR0<bool>(true),
                              LiteralUtil::CreateR0<float>(12.2f),
                              LiteralUtil::CreateR1<float>({12.8f, 14.6f})}),
                         {pred_arg.get()}, error_spec_);
}

// Test true and false computations that return a nested tuple.
XLA_TEST_F(ConditionalOpTest, ReturnNestedTuple) {
  XlaBuilder true_builder(TestName() + ".true");
  {
    Parameter(&true_builder, 0, empty_tuple_, "tuple");
    auto true_constant1 = ConstantR0<float>(&true_builder, 12.2f);
    auto true_constant2 = ConstantR1<float>(&true_builder, {12.8f, 14.6f});
    auto true_constant3 = ConstantR1<float>(&true_builder, {25.4f, 29.8f});
    auto true_constant4 = ConstantR0<float>(&true_builder, 35.6f);
    Tuple(&true_builder,
          {Tuple(&true_builder, {true_constant1, true_constant2}),
           Tuple(&true_builder, {true_constant3, true_constant4})});
  }
  auto true_builder_result = true_builder.Build();
  EXPECT_IS_OK(true_builder_result.status());

  XlaBuilder false_builder(TestName() + ".false");
  {
    Parameter(&false_builder, 0, empty_tuple_, "tuple");
    auto false_constant1 = ConstantR0<float>(&false_builder, 46.6f);
    auto false_constant2 = ConstantR1<float>(&false_builder, {54.4f, 58.4f});
    auto false_constant3 = ConstantR1<float>(&false_builder, {62.1f, 67.4f});
    auto false_constant4 = ConstantR0<float>(&false_builder, 9.3f);
    Tuple(&false_builder,
          {Tuple(&false_builder, {false_constant1, false_constant2}),
           Tuple(&false_builder, {false_constant3, false_constant4})});
  }
  auto false_builder_result = false_builder.Build();
  EXPECT_IS_OK(false_builder_result.status());

  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operands = Tuple(&builder, {});
  Conditional(pred, operands, std::move(true_builder_result).value(), operands,
              std::move(false_builder_result).value());

  ComputeAndCompareTuple(
      &builder,
      LiteralUtil::MakeTupleFromSlices(
          {LiteralUtil::MakeTupleFromSlices(
               {LiteralUtil::CreateR0<float>(46.6f),
                LiteralUtil::CreateR1<float>({54.4f, 58.4f})}),
           LiteralUtil::MakeTupleFromSlices(
               {LiteralUtil::CreateR1<float>({62.1f, 67.4f}),
                LiteralUtil::CreateR0<float>(9.3f)})}),
      {pred_arg.get()}, error_spec_);
}

// Test conditional that takes in scalar operands in the form of external
// params.
XLA_TEST_F(ConditionalOpTest, ScalarOperandsFromExternalParams) {
  Shape r0bool = ShapeUtil::MakeShape(PRED, {});
  XlaBuilder builder(TestName());

  XlaOp pred, operand1, operand2;
  auto pred_arg = CreateR0Parameter<bool>(true, 0, "pred", &builder, &pred);
  auto operand1_param =
      CreateR0Parameter<float>(56.3f, 1, "operand1", &builder, &operand1);
  auto operand2_param =
      CreateR0Parameter<float>(12.7f, 2, "operand2", &builder, &operand2);
  Conditional(pred, operand1, CreateR0CeilComputation(), operand2,
              CreateR0FloorComputation());

  ComputeAndCompareR0<float>(
      &builder, 57.0f,
      {pred_arg.get(), operand1_param.get(), operand2_param.get()},
      error_spec_);
}

// Test conditional that takes in array operands in the form of external params.
XLA_TEST_F(ConditionalOpTest, ArrayOperandsFromExternalParams) {
  Shape r0bool = ShapeUtil::MakeShape(PRED, {});
  XlaBuilder builder(TestName());

  XlaOp pred, operand1, operand2;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1_param = CreateR1Parameter<float>({24.3f, 56.7f}, 1, "operand1",
                                                 &builder, &operand1);
  auto operand2_param = CreateR1Parameter<float>({10.2f, 11.6f}, 2, "operand2",
                                                 &builder, &operand2);
  Conditional(pred, operand1, CreateR1CeilComputation(), operand2,
              CreateR1FloorComputation());

  ComputeAndCompareR1<float>(
      &builder, {10.0f, 11.0f},
      {pred_arg.get(), operand1_param.get(), operand2_param.get()},
      error_spec_);
}

// Test the case where one conditional is nested within another.
XLA_TEST_F(ConditionalOpTest, NestedConditionals) {
  XlaBuilder inner_builder(TestName() + ".inner_conditional");
  {
    Shape r0bool = ShapeUtil::MakeShape(PRED, {});
    Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_});
    auto param0 = Parameter(&inner_builder, 0, tuple_shape, "param0");
    auto pred_cond = GetTupleElement(param0, 0);
    auto true_operand = GetTupleElement(param0, 1);
    auto false_operand = GetTupleElement(param0, 2);
    Conditional(pred_cond, true_operand, CreateR0CeilComputation(),
                false_operand, CreateR0FloorComputation());
  }
  auto inner_builder_result = inner_builder.Build();
  EXPECT_IS_OK(inner_builder_result.status());

  XlaBuilder builder(TestName());
  XlaOp pred1, pred2;
  auto pred1_arg = CreateR0Parameter<bool>(true, 0, "pred1", &builder, &pred1);
  auto pred2_arg = CreateR0Parameter<bool>(false, 1, "pred2", &builder, &pred2);
  auto operand1 = ConstantR0<float>(&builder, 1.1f);
  auto operand2 = ConstantR0<float>(&builder, 12.2f);
  auto operand3 = ConstantR0<float>(&builder, 43.3f);
  auto tuple_operand = Tuple(&builder, {pred2, operand1, operand2});
  Conditional(pred1, tuple_operand, std::move(inner_builder_result).value(),
              operand3, CreateR0IdentityComputation());

  ComputeAndCompareR0<float>(&builder, 12.0f,
                             {pred1_arg.get(), pred2_arg.get()}, error_spec_);
}

XLA_TEST_F(ConditionalOpTest, ConditionalInNestedComputation) {
  XlaBuilder inner_builder(TestName() + ".inner_conditional");
  {
    Shape r0bool = ShapeUtil::MakeShape(PRED, {});
    Shape tuple_shape = ShapeUtil::MakeTupleShape({r0bool, r0f32_, r0f32_});
    auto param0 = Parameter(&inner_builder, 0, tuple_shape, "param0");
    auto pred_cond = GetTupleElement(param0, 0);
    auto true_operand = GetTupleElement(param0, 1);
    auto false_operand = GetTupleElement(param0, 2);
    Conditional(pred_cond, true_operand, CreateR0CeilComputation(),
                false_operand, CreateR0FloorComputation());
  }
  auto inner_builder_result = inner_builder.Build();
  EXPECT_IS_OK(inner_builder_result.status());

  XlaBuilder builder(TestName());
  XlaOp pred;
  auto pred_arg = CreateR0Parameter<bool>(false, 0, "pred", &builder, &pred);
  auto operand1 = ConstantR0<float>(&builder, 1.1f);
  auto operand2 = ConstantR0<float>(&builder, 12.2f);
  auto tuple_operand = Tuple(&builder, {pred, operand1, operand2});
  Call(&builder, std::move(inner_builder_result).value(), {tuple_operand});

  ComputeAndCompareR0<float>(&builder, 12.0f, {pred_arg.get()}, error_spec_);
}

// Test a mismatch in the shape of the true operand and true computation.
XLA_TEST_F(ConditionalOpTest, ShapeMismatch) {
  XlaBuilder builder(TestName());
  auto pred = ConstantR0<bool>(&builder, true);
  auto operand1 = ConstantR0<float>(&builder, 56.0f);
  auto operand2 = ConstantR0<float>(&builder, 12.0f);
  auto operands = Tuple(&builder, {operand1, operand2});
  Conditional(pred, operands, CreateR1TupleAddComputation(), operands,
              CreateR0TupleSubComputation());

  auto result = builder.Build();
  EXPECT_FALSE(result.ok());
  EXPECT_THAT(result.status().error_message(),
              ::testing::HasSubstr("operand 0 must match the shape of the "
                                   "only parameter of branch computation 0"));
}

XLA_TEST_F(ConditionalOpTest, SwappedInputsInSequentialConditionals) {
  Shape tuple_shape = ShapeUtil::MakeTupleShape({r0f32_, r0f32_});
  XlaComputation swapper;
  {
    XlaBuilder builder(TestName() + ".swapper");
    auto param0 = Parameter(&builder, 0, tuple_shape, "sp0");
    auto x = GetTupleElement(param0, 0);
    auto y = GetTupleElement(param0, 1);
    Tuple(&builder, {y, x});
    swapper = builder.Build().value();
  }
  XlaComputation forwarder;
  {
    XlaBuilder builder(TestName() + ".forwarder");
    auto param0 = Parameter(&builder, 0, tuple_shape, "fp0");
    auto x = GetTupleElement(param0, 0);
    auto y = GetTupleElement(param0, 1);
    Tuple(&builder, {x, y});
    forwarder = builder.Build().value();
  }
  XlaComputation main;
  {
    XlaBuilder builder(TestName() + ".main");
    auto param0 = Parameter(&builder, 0, tuple_shape, "mp0");
    auto x = GetTupleElement(param0, 0);
    auto y = GetTupleElement(param0, 1);
    auto lt_pred = Lt(x, y);
    auto res = Conditional(lt_pred, param0, forwarder, param0, swapper);
    auto ge_pred = Ge(x, y);
    Conditional(ge_pred, res, swapper, res, forwarder);
    main = builder.Build().value();
  }

  auto test_swap = [&](float a, float b) {
    XlaBuilder builder(TestName());
    XlaOp x, y;
    auto x_arg = CreateR0Parameter<float>(a, 0, "x", &builder, &x);
    auto y_arg = CreateR0Parameter<float>(b, 1, "y", &builder, &y);
    auto tuple_operand = Tuple(&builder, {x, y});
    Call(&builder, main, {tuple_operand});

    ComputeAndCompareTuple(
        &builder,
        LiteralUtil::MakeTupleFromSlices(
            {LiteralUtil::CreateR0<float>(a), LiteralUtil::CreateR0<float>(b)}),
        {x_arg.get(), y_arg.get()}, error_spec_);
  };

  test_swap(3.11f, 9.4f);
  test_swap(11.24f, 5.55f);
}

// Test conditional that duplicates tuple elements in the then and else
// computations. This is a regression test for b/112550242.
XLA_TEST_F(ConditionalOpTest, DuplicateElementsConditional) {
  const Shape scalar = ShapeUtil::MakeShape(S32, {});
  const Shape tuple2 = ShapeUtil::MakeTupleShape({scalar, scalar});
  XlaComputation then_comp;
  {
    XlaBuilder builder(TestName() + ".then");
    auto p = Parameter(&builder, 0, tuple2, "then.p");
    auto e0 = GetTupleElement(p, 0);
    auto e1 = GetTupleElement(p, 1);
    Tuple(&builder, {e0, e1, e0});
    then_comp = builder.Build().value();
  }
  XlaComputation else_comp;
  {
    XlaBuilder builder(TestName() + ".else");
    auto p = Parameter(&builder, 0, tuple2, "else.p");
    auto e0 = GetTupleElement(p, 0);
    auto e1 = GetTupleElement(p, 1);
    Tuple(&builder, {e0, e1, e1});
    else_comp = builder.Build().value();
  }

  {
    // Pred is true case.
    std::vector<Literal> args;
    args.push_back(LiteralUtil::MakeTupleFromSlices(
        {LiteralUtil::CreateR0<int32_t>(123),
         LiteralUtil::CreateR0<int32_t>(-42)}));
    args.push_back(LiteralUtil::CreateR0<bool>(true));
    XlaBuilder builder(TestName() + ".main");
    auto p = Parameter(&builder, 0, tuple2, "p0");
    auto p_pred = Parameter(&builder, 1, ShapeUtil::MakeShape(PRED, {}), "p1");
    Conditional(p_pred, p, then_comp, p, else_comp);
    ComputeAndCompare(&builder, args);
  }
  {
    // Pred is false case.
    std::vector<Literal> args;
    args.push_back(LiteralUtil::MakeTupleFromSlices(
        {LiteralUtil::CreateR0<int32_t>(123),
         LiteralUtil::CreateR0<int32_t>(-42)}));
    args.push_back(LiteralUtil::CreateR0<bool>(false));
    XlaBuilder builder(TestName() + ".main");
    auto p = Parameter(&builder, 0, tuple2, "p0");
    auto p_pred = Parameter(&builder, 1, ShapeUtil::MakeShape(PRED, {}), "p1");
    Conditional(p_pred, p, then_comp, p, else_comp);
    ComputeAndCompare(&builder, args);
  }
}

}  // namespace
}  // namespace xla