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

  // However, its probably better to fold the IfRegion instead of having the
  // dead branch stay.

  // Inline the region in place of the IfRegion op, and forward the yield
  // inputs to the IfRegion op results. This is possible only if the yield
  // types match the result types.
  auto yield = cast<YieldOp>(region.front().getTerminator());
  auto updated_results = llvm::to_vector<4>(yield.getOperands());

  func.return %graph : tensor<i1>
  // expected-remark@above {{ID: 4}}
  // expected-remark@above {{Sinks: {}}}
}

// -----

// Tests that the pass does not add control dependencies for a stateless
// IfRegion op.

// CHECK-LABEL: func @stateless_ifregion_op
func.func @stateless_ifregion_op(
  // expected-remark@above {{ID: 18}}
  %arg0: tensor<*x!tf_type.resource<tensor<32xf32>>>,
  %arg1: tensor<i1>) {