func.return
}

// CHECK-LABEL: @IfOpTokenAttrs
func.func @IfOpTokenAttrs(%arg0: tensor<i1>, %arg1: tensor<f32>) -> (tensor<f32>) {
  // CHECK: tf.IfRegion
  %0 = "tf.IfRegion"(%arg0) ({
      // CHECK: tf.XlaRecvFromHost
      // CHECK-SAME-DAG: _xla_original_oc_node_name =
      // CHECK-SAME-DAG: _xla_token_input_nodes = ["_xla_token_arg_node"]

  // CHECK: "tf.If"({{.+}}) <{else_branch = @tf.IfRegion1_else, {{.+}} then_branch = @tf.IfRegion1_then}
  %0 = "tf.IfRegion"(%arg0) ({
    // Outer Then
    %cond = "tf.LogicalNot"(%arg0) : (tensor<i1>) -> tensor<i1>
    %asin = "tf.Asin"(%arg1) : (tensor<*xf32>) -> tensor<*xf32>

    // nested IfRegion
    %1 = "tf.IfRegion"(%cond) ({
        %2 = "tf.Abs"(%asin) : (tensor<*xf32>) -> tensor<*xf32>

    // CHECK-NOT:  _xla_outside_compilation
    // CHECK:      "tf.IfRegion"
    // CHECK:        "tf.StringToNumber"
    // CHECK-SAME:   _xla_outside_compilation
    %1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %2 = "tf.IfRegion"(%arg0) ({
      %3 = "tf.StringToNumber"(%arg1) {out_type = f32} : (tensor<!tf_type.string>) -> tensor<f32>

                                  {"tf.Yield"(%if12) : (!tf_variant) -> ()})
        {is_stateless = true} : (tensor<i1>) -> !tf_variant
    %if14 = "tf.IfRegion"(%arg0) ({"tf.Yield"(%if13) : (!tf_variant) -> ()},
                                  {"tf.Yield"(%if13) : (!tf_variant) -> ()})
        {is_stateless = true} : (tensor<i1>) -> !tf_variant

    return %1 : tensor<1x3xf32>
  }
}

// -----

// Check the IfRegion is set to stateful.
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1833 : i32}, tf_saved_model.semantics} {
  // CHECK-LABEL: func.func @serving_default
  // CHECK: "tf.IfRegion"
  // CHECK-SAME: is_stateless = false

  // CHECK: "tf.CalibrationStatisticsSaver"

                                   llvm::ArrayRef<func::FuncOp> functions,
                                   const BacktrackAnalysis& backtrack_analysis);

  // Analyzes tf.CaseRegion/tf.IfRegion ops to compute resource IDs.
  void AnalyzeRegionCaseOrIfOp(Operation* case_or_if_op,
                               const BacktrackAnalysis& backtrack_analysis);

  let summary = "Merges IfRegion ops together with a common predicate.";

  let description = [{
    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>)

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

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

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

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

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