// CHECK: [[name:%.*]] = "tf.VarHandleOp"
    // CHECK: "tf.ReadVariableOp"([[name]])
    %cond = builtin.unrealized_conversion_cast to tensor<i1>
    %0 = "tf.IfRegion"(%cond) ({
      "tf.Yield"(%arg0) : (tensor<!tf_type.resource<tensor<10xf32>>>) -> ()
    }, {
      "tf.Yield"(%arg0) : (tensor<!tf_type.resource<tensor<10xf32>>>) -> ()

==============================================================================*/

// This transformation pass transforms functional control flow operations in the
// TensorFlow dialect to their region based counterparts, i.e.,
// tf.If -> tf.IfRegion and tf.While -> tf.WhileRegion

#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project

  // expected-remark@below {{Result #0, ID 0 : 0}}
  %vx = "tf.VarHandleOp"() {container = "cf", shared_name = "vx"} : () -> !tf_res
  func.return %vx, %arg0 : !tf_res, !tf_res
}

// -----
// Test aliasing through IfRegion

!tf_res = tensor<*x!tf_type.resource<tensor<i1>>>

// CHECK-LABEL: func @if_region_aliasing
// expected-remark@below {{Region #0, Arg #0, ID 7 : 1, 4, 6, 7}}

    // CHECK-SAME: device = "/job:worker/replica:0/task:0/device:CPU:0"
    // CHECK-NEXT:   %[[C_OUT:.*]] = "tf.C"(%arg0, %[[CLUSTER_OUT]]#2)
    // CHECK-NOT:    _xla_outside_compilation
    // CHECK         "tf.IfRegion"
    // CHECK:          "tf.D"(%[[C_OUT]], %arg0, %[[CLUSTER_OUT]]#0)
    // CHECK-NOT:      _xla_outside_compilation
    // CHECK:        tf_device.return
    "tf_device.cluster"() ({

their `mlir_module` string attribute.
### `-tf-merge-control-flow`

_Merges IfRegion ops together with a common predicate._

This pass merges IfRegion ops together if they have the same predicate and it
is safe to do so (there are no intermediate dependencies, they are in the
same block, etc).

For example:

```mlir
"tf.IfRegion"(%0) ( {
  %2 = "tf.A"() : () -> (tensor<f32>)
  "tf.Yield"() : () -> ()

LogicalResult SetTokenInputAttrs(ModuleOp module) {
  // Collect all the ops that needs to have token input names attributes. These
  // ops are communication ops and all their parent ops via nesting or function
  // calls. For example, IfRegion op and PartitionedCall op.
  std::vector<Operation*> worklist;
  absl::flat_hash_set<Operation*> ops_with_tokens;
  module.walk([&](Operation* op) {
    if (IsCommunicationOp(op)) {
      ops_with_tokens.insert(op);

  // CHECK-NEXT: %0 = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({
  // CHECK-NEXT:   "tf.Yield"(%arg1) : (tensor<1xi8>) -> ()
  // CHECK-NEXT:   }, {
  // CHECK-NEXT:   "tf.Yield"(%arg2) : (tensor<1xi8>) -> ()
  // CHECK-NEXT:  }) : (tensor<i1>) -> tensor<1xi8>
  // CHECK-NEXT: return %0 : tensor<1xi8>
  %0 = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({

==============================================================================*/

// This transformation pass transforms region bases control flow operations in
// the TensorFlow dialect to their functional counterparts, i.e.,
// tf.IfRegion ->  tf.If and tf.WhileRegion -> tf.While

#include <iterator>
#include <memory>
#include <optional>
#include <string>

#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"

  "tf.StackCloseV2"(%stack) : (tensor<!tf_type.resource>) -> ()
  func.return
}

// -----

// Tests that the pass returns meaningful error message when IfRegion op has
// resource returns.

func.func @main(%arg0: tensor<i1>) -> () {
  %max_size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>

    %0 = "tf.Sum"(%arg0, %cst_0) <{keep_dims = false}> {device = ""} : (tensor<1x4xf32>, tensor<2xi32>) -> tensor<f32>
    %1 = "tf.Greater"(%0, %cst) {device = ""} : (tensor<f32>, tensor<f32>) -> tensor<i1>
    %2:2 = "tf.IfRegion"(%1) <{_else_func_name = "cond_false_80", _then_func_name = "cond_true_70", is_stateless = true}> ({
      %4 = "tf.Identity"(%cst_3) {device = ""} : (tensor<i1>) -> tensor<i1>