tf_device.return
      }
      tf_executor.yield
    }
    %1 = tf_executor.island {
      tf_device.replicate {n = 2 : i32, devices = {CORE_0 = ["/CPU:0", "/GPU:1"]}} {
        "tf_device.launch"() ({
          "tf.opA"() : () -> ()
          tf_device.return
        }) {device = "/TPU:1"} : () -> ()
        tf_device.return
      }
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return

//      tf_device.return
//    }) {} : () -> (tensor<i1>)
//    tf_executor.yield %4 : tensor<i1>
//  }
//  tf_executor.fetch %3#0 : tensor<i1>
//
// gets lowered to:
//
//  %0 = tf_executor.island {
//    tf_executor.yield
//  }
//  %1:2 = tf_executor.island {
//    %2 = "tf.opA"(%arg0) : (tensor<i1>) -> tensor<i1>
//    tf_executor.yield %2 : tensor<i1>
//  }
//

    %2 = "tf.IfRegion"(%arg0) ({
      %3 = "tf.StringToNumber"(%arg1) {out_type = f32} : (tensor<!tf_type.string>) -> tensor<f32>
      "tf.Yield"(%3) : (tensor<f32>) -> ()
     },  {
      %4 = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
      "tf.Yield"(%4) : (tensor<f32>) -> ()
    }) {is_stateless = true} : (tensor<i1>) -> (tensor<f32>)
    %5 = "tf.Identity"(%2) : (tensor<f32>) -> tensor<f32>

    "tfl.yield"(%1) : (tensor<i1>) -> ()
  },  {
  ^bb0(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>):
    // CHECK: call @WhileOp_body
    // CHECK-SAME: (tensor<*xi32>, tensor<*xf32>, tensor<i32>)
    %1:3 = func.call @WhileOp_body(%arg0, %arg1, %cst) : (tensor<*xi32>, tensor<*xf32>, tensor<i32>) -> (tensor<*xi32>, tensor<*xf32>, tensor<i32>)
    "tfl.yield"(%1#0, %1#1) : (tensor<*xi32>, tensor<*xf32>) -> ()

  // Operation which is expected to consume all the call results.
  Operation* call_consumer = yield;

  // Allow a single ToBoolOp between the call and the yield (valid only
  // when the yield has a single operand)
  if (allow_to_bool && yield.getNumOperands() == 1 && isa<ToBoolOp>(*it)) {
    if (it->getResult(0) != yield.getOperand(0)) return std::nullopt;
    call_consumer = cast<ToBoolOp>(*it);
    it++;

    // 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>>>) -> ()
    }) { is_stateless = false} : (tensor<i1>) -> tensor<!tf_type.resource<tensor<10xf32>>>

  func.return
}

// -----

// Check that a tf_executor.yield parent is a tf_executor.island.
func.func @parent_is_island() {
  "tf.some_op"() ({
    tf_executor.yield
// expected-error@-1 {{'tf_executor.yield' op expects parent op 'tf_executor.island'}}
  }) : () -> ()
  func.return
}

// -----

// Check that an island yield matches the island results.

    %1 = "tf.WhileRegion"(%0) ({
       // Condition region
       // CHECK: ^bb
       // CHECK: "tf.Yield"
       ^bb0(%carg0: tensor<i32>):
          %c0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
          %c1 = "tf.GreaterEqual"(%carg0, %0) {T = i32, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i1>
          "tf.Yield"(%c1) : (tensor<i1>) -> ()
      }, {
       // Body region
       // CHECK: ^bb0

            "tf.Yield"(%u0, %id0, %id0) : (!tf_res, !tf_res, !tf_res) -> ()
          }, {
            // expected-remark@below {{Result #0, ID 3 : 0, 1, 3, 4, 5}}
            %id0 = "tf.Identity"(%vh0) : (!tf_res) -> !tf_res
            "tf.Yield"(%id0, %id0, %arg0) : (!tf_res, !tf_res, !tf_res) -> ()
          }, {
            "tf.Yield"(%vh0, %arg1, %arg1) : (!tf_res, !tf_res, !tf_res) -> ()

  Operation &yield = island.GetBody().back();
  if (!isa<YieldOp>(yield))
    return yield.emitOpError()
           << "invalid tf_executor.island terminator, yield expected";

  // Ensure that the yield terminator operands matches the island results type.
  int result_count = island.getNumResults() - 1;  // -1 for the control token
  const int num_operands = yield.getNumOperands();
  if (num_operands != result_count)