TF_Int32Tensor:$max_num_elements
  );
}

def TF_IfOp : TF_Op<"If", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
  let summary = "output = cond ? then_branch(input) : else_branch(input)";

  let description = [{
output = cond ? then_branch(input) : else_branch(input)

cond: A Tensor. If the tensor is a scalar of non-boolean type, the
    scalar is converted to a boolean according to the

  // CHECK: %[[GVAR2:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<5x3xf32>>>
  // CHECK: "tf.If"(%[[COND]], %[[VAR]], %[[GVAR1]], %[[GVAR2]])
  %1 = "tf.If"(%cond, %ta#0) {
    then_branch = @then_branch, else_branch = @else_branch, device = "", is_stateless = false}
       : (tensor<i1>, tensor<!tf_type.resource>) -> tensor<!tf_type.resource>

  // CHECK: %[[CLUSTER:.*]]:2 = "tf_device.cluster"()
  %2 = "tf_device.cluster"() ({
    // CHECK: %[[IF:.*]]:2 = "tf.If"(%[[ARG0]], %[[READ0]], %[[READ1]])
    %3:2 = "tf.If"(%arg0, %0, %1) {then_branch = @if_then, else_branch = @if_else,
        is_stateless = false}
      : (tensor<i1>, tensor<*x!tf_type.resource<tensor<4xf32>>>, tensor<*x!tf_type.resource<tensor<4xf32>>>)

    } else if (auto if_op = llvm::dyn_cast<TF::IfOp>(&op)) {
      auto then_branch = if_op.then_function();
      auto else_branch = if_op.else_function();
      // Recursively handle the nested control flow.
      (void)HoistForControlFlow(&then_branch.front(), module, vars_initialized,
                                lifted_partitioned_call_callees);

      /*is_stateless=*/false, if_region.get_thenFuncNameAttr(),
      if_region.get_elseFuncNameAttr());

  // Create empty then branch region.
  auto& then_branch = host_side_if.getThenBranch();
  then_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&then_branch.front());
  builder.create<mlir::TF::YieldOp>(if_region.getLoc(),
                                    /*operands=*/ArrayRef<Value>{});

    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);

    }
    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))};
  }

`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

```mlir

  }

  @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() {