// CHECK-NEXT: "tf.Asin"
// CHECK-NEXT: "tf.If"({{.+}}) <{else_branch = @tf.IfRegion_else, {{.+}} then_branch = @tf.IfRegion_then}

// CHECK: func private @tf.IfRegion_else
// CHECK-NEXT: "tf.Neg"
// CHECK: func private @tf.IfRegion_then
// CHECK-NEXT: "tf.Abs"

func.func @testNested(%arg0: tensor<i1>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  // CHECK: "tf.If"({{.+}}) <{else_branch = @tf.IfRegion1_else, {{.+}} then_branch = @tf.IfRegion1_then}

  dataset_if->setAttr("_lower_using_switch_merge", builder.getBoolAttr(true));
  // Empty else branch, if there is no more data, do nothing.
  auto& else_branch = dataset_if.getElseBranch();
  else_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&else_branch.front());
  // Return only the state variables from the body arguments.
  SmallVector<Value, 4> else_returns;
  for (int i = 1; i < state_size + 1; i++) {

      if (type_from_body.has_value()) return type_from_body;
    } else if (auto if_op = llvm::dyn_cast<TF::IfOp>(use.getOwner())) {
      auto then_branch = if_op.then_function();
      auto else_branch = if_op.else_function();
      assert(then_branch && else_branch);
      auto type_from_then = GetElementTypeFromAccess(
          then_branch.getArgument(use.getOperandNumber() - 1), module,
          infer_from_op);

  // CHECK: [[cond_res:%.*]]:2 = tfrt.cond [[cond]]
  // CHECK-SAME: @cond_true @cond_false(%arg0, %arg2) : (!tfrt.chain, !tfrt_fallback.tf_tensor)
  %2 = "tf.If"(%arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = true} : (tensor<i1>, tensor<i32>) -> tensor<i32>
  // CHECK: [[out_ch:%.*]] = tfrt.merge.chains [[cond_res]]#0, %arg0 : !tfrt.chain, !tfrt.chain

  // Rewrite the branches.
  Region& then_branch = if_op.getThenBranch();
  Region& else_branch = if_op.getElseBranch();
  if (failed(DecomposeTensorListOpsInternal(
          &then_branch.front(), module, buffer_to_size,
          decomposed_partitioned_call_callees)))
    return failure();
  if (failed(DecomposeTensorListOpsInternal(
          &else_branch.front(), module, buffer_to_size,
          decomposed_partitioned_call_callees)))

namespace tensorflow {

TEST(RearrangeFunctionArgumentForFunctionTest, Basic) {
  FunctionDefLibrary fdl;
  {
    // Function for StatefulPartitionedCall's "f", If's
    // "then_branch"/"else_branch".
    // "arg0" (T=DT_RESOURCE), "arg1" (T=DT_BOOL)
    // "ret0" = "arg1"
    // "ret1" = "arg0"
    tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  if (!s.ok()) {
    *uncompilable_reason = s.message();
    return false;
  }
  return true;
}

// Tests whether 'if_node' is compilable. Every operator in the then_branch and
// else_branch functions must be compilable for 'if_node' to be compilable.
bool RecursiveCompilabilityChecker::IsCompilableIf(
    const Node& if_node, FunctionLibraryRuntime* lib_runtime,
    std::vector<StackFrameView>* stack_trace,

  // CHECK:  "tf._TfrtGetResource"()
  %cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
  // CHECK: tf.If
  %x = "tf.If"(%cond, %handle) {then_branch = @some_func, else_branch = @some_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %r : tensor<i32>

    %1 = "tf.Greater"(%0, %cst_0) {device = ""} : (tensor<f32>, tensor<f32>) -> tensor<i1>

// 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