// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else([[x]], [[y]])
// CHECK: return [[z]]
func.func @main(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
  %z = "tf.If"(%cond, %x, %y) {then_branch = @then, else_branch = @else, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
  return %z: tensor<i32>
}

// -----

// Test lowering AsyncOpKernel

// CHECK-LABEL: func @main

    %1 = "tf.Greater"(%0, %cst_0) {device = ""} : (tensor<f32>, tensor<f32>) -> tensor<i1>
    %2:2 = "tf.If"(%1, %arg0) <{else_branch = @cond_false_80, is_stateless = true, then_branch = @cond_true_70}> {Tcond = i1, Tin = [f32], Tout = [i1, f32], _lower_using_switch_merge = true, _read_only_resource_inputs = [], device = ""} : (tensor<i1>, tensor<1x4xf32>) -> (tensor<i1>, tensor<1x3xf32>)

    auto original_point = rewriter.saveInsertionPoint();
    rewriter.setInsertionPointAfter(op->getParentOfType<func::FuncOp>());
    SymbolTable manager(op->getParentOfType<ModuleOp>());

    // Constructs `then_branch`, which is executed when `if_cond` evaluates to
    // true.
    auto then_branch_op =
        rewriter.create<func::FuncOp>(loc, "cond_true", func_type);

  // CHECK-SAME: device = "noodle"
  %3 = "tf.If"(%1, %2, %arg1) {then_branch = @add, else_branch = @sub, output_shapes = [#tf_type.shape<>], device = "noodle", _underscore_attr = "something", is_stateless = false} : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>

  // CHECK: %2 = "tf.If"

                    %arg1: tensor<4x!quant.uniform<u8:f32, 1.0>>) -> (tensor<4x!quant.uniform<u8:f32, 1.0>>) {
  %0 = "tfl.dequantize"(%arg1) : (tensor<4x!quant.uniform<u8:f32, 1.0>>) -> tensor<4xf32>
  %1 = "tf.If"(%arg0, %0) {then_branch = @testIfThen, else_branch = @testIfElse, is_stateless = false} : (tensor<i1>, tensor<4xf32>) -> tensor<4xf32>

    `tf.IfRegion` and `tf.While` is transformed to `tf.WhileRegion`.

    For example, this functional operation

    ```mlir
      %0 = "tf.If"(%arg0, %arg1) {
        then_branch = @then_branch_func, else_branch = @else_branch_func, is_stateless = false
      } : (tensor<i1>, tensor<*xf32>) -> tensor<*xf32>
    ```

    will be transformed into this region-based operation

    }
    auto else_attr =
        mlir::SymbolRefAttr::get(builder.getContext(), func_names.at(else_idx));

    return llvm::SmallVector<mlir::NamedAttribute, 4>{
        builder.getNamedAttr("then_branch", then_attr),
        builder.getNamedAttr("else_branch", else_attr),
        // TODO(b/139667752): Analyze statelessness correctly
        builder.getNamedAttr("is_stateless", builder.getBoolAttr(false))};
  }

LogicalResult IfOp::verifySymbolUses(SymbolTableCollection& symbol_table) {
  auto branch_name = [](unsigned index) -> std::string {
    return index == 0 ? "'then_branch'" : "'else_branch'";
  };
  return VerifyCaseOrIfOpBranchFunctions(
      symbol_table, *this, {getThenBranchAttr(), getElseBranchAttr()},
      branch_name);
}

  }

  @Test
  @TestMetadata("when_branch.kt")
  public void testWhen_branch() {
    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/when_branch.kt");
  }

  @Test
  @TestMetadata("when_branchUsed.kt")
  public void testWhen_branchUsed() {

  }

  @Test
  @TestMetadata("when_branch.kt")
  public void testWhen_branch() {
    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/when_branch.kt");
  }

  @Test
  @TestMetadata("when_branchUsed.kt")
  public void testWhen_branchUsed() {