-> !tf_res
      // CHECK-NEXT: "tf.VarHandleOp"
      %var_handle = "tf.VarHandleOp"() {container = "c", shared_name = "v0", device = "/TPU:1"}
        : () -> !tf_res
      // CHECK-NEXT: "tf.IfRegion"
      "tf.IfRegion"(%arg1) ({
          func.call @ifregion_then(%id0, %var_handle) : (!tf_res, !tf_res) -> ()
          "tf.Yield"() : () -> ()
        }, {

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

    // CHECK:          %[[GET_NEXT:[0-9]*]] = "tf.IteratorGetNextAsOptional"(%[[ANON_ITER]])
    // CHECK-NEXT:     %[[HAS_VALUE:[0-9]*]] = "tf.OptionalHasValue"(%[[GET_NEXT]])
    // CHECK-NEXT:     %[[IF:.*]] = "tf.IfRegion"(%[[HAS_VALUE]])
    // CHECK-NEXT:       %[[GET_VALUE:[0-9]*]] = "tf.OptionalGetValue"(%[[GET_NEXT]])
    // CHECK-NEXT:       %[[FUNC_CALL:[0-9]*]] = func.call @__reduce_func_1(%[[ARG_5]], %[[GET_VALUE]])

  // CHECK-NOT: _replication_info =
  // CHECK-NOT: _xla_compile_device_type
  // CHECK-NOT: device =
  %1 = "tf_device.cluster"() ({
    %2 = "tf.Add"(%arg1, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
    %3 = "tf.IfRegion"(%arg0) ({
        %4 = "tf.Mul" (%arg1, %2) {device = "y"}: (tensor<f32>, tensor<f32>) -> tensor<f32>
        "tf.Yield"(%4) : (tensor<f32>) -> ()
      }, {
        // CHECK: device = "/device:TPU_REPLICATED_CORE:0"

  // CHECK: %[[IF:.*]]:2 = "tf.IfRegion"(%[[ARG0]])
  %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<4xf32>>>
  %1 = "tf.VarHandleOp"() {container = "c", shared_name = "v2"} : () -> tensor<*x!tf_type.resource<tensor<4xf32>>>
  %2 = "tf_device.cluster"() ({
    %3:2 = "tf.IfRegion"(%arg0) ({
          // CHECK-NEXT: %[[CONST:.*]] = "tf.Const"()

// -----

// Test invalid tf.Yield operation (parent should be IfRegion)
func.func @testInvalidYieldOp(%arg0: f32) -> () {
  // expected-error @+1 {{'tf.Yield' op expects parent op to be one of 'tf.CaseRegion, tf.IfRegion, tf.WhileRegion, tf.GeneratorDatasetRegion'}}
  "tf.Yield"(%arg0) : (f32) -> ()
}

// -----

// Test valid tf.IfRegion operation
// CHECK-LABEL: func @testValidIfRegionOp

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

  // CHECK: "tf.IfRegion"(%[[PRED]]) <{is_stateless = false}> ({
  %case_op = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({
      // CHECK: %[[TA_VAL:.*]] = "tf.ReadVariableOp"(%[[TA_BUFFER]])
      // CHECK: "tf.Slice"(%[[TA_VAL]]
      // CHECK-NOT: tf.TensorArrayReadV3