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

    // CHECK-NOT: "tf.VarHandleOp"
    %handle = "tf.VarHandleOp"() {container="", shared_name="var1", device = "/job:worker/replica:0/task:1/device:CPU:0"} : () -> tensor<!tf_type.resource<tensor<0xf32>>>
    %0 = "tf.IfRegion"(%arg0) ({
      %1 = "tf.ReadVariableOp"(%handle) : (tensor<!tf_type.resource<tensor<0xf32>>>) -> (tensor<0xf32>)
      "tf.Yield"(%1) : (tensor<0xf32>) -> ()
     },  {

  %true = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>

  // CHECK: [[Val0:%.*]] = "tf.Mul"(%arg0, %arg1)
  // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0"
  %0 = "tf.IfRegion"(%true) ({
      %true_value = "tf.Mul"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
      "tf.Yield"(%true_value) : (tensor<f32>) -> ()
    }, {

// preserved.
//
// CHECK-LABEL: func @yieldOp
func.func @yieldOp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i1>) -> (tensor<f32>) {
  // CHECK-COUNT-2: tf.Yield
  %0 = "tf.IfRegion"(%arg2) ({
      "tf.Yield"(%arg0) : (tensor<f32>) -> ()
    }, {
      "tf.Yield"(%arg1) : (tensor<f32>) -> ()
    }) { is_stateless = true}: (tensor<i1>) -> tensor<f32>
  func.return %0 : tensor<f32>
}