auto executor_graph =
        rewriter.create<tf_executor::GraphOp>(loc, func_type.getResults());
    rewriter.inlineRegionBefore(graph.getNodes(), executor_graph.getBody(),
                                executor_graph.getBody().end());

    // Add terminator of tf_executor::graph
    rewriter.setInsertionPointToEnd(&executor_graph.getBody().front());
    rewriter.create<tf_executor::FetchOp>(loc);

    // Add terminator of func

func.func @controls_per_replica() {
  tf_executor.graph {
    %1 = tf_executor.ControlTrigger {}
    %2 = tf_executor.ControlTrigger {}
    %3 = tf_executor.island(%1, %2) {
      tf_device.replicate {n = 2 : i32} {
        tf_device.return
      }
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK: %[[CT_0:.*]] = tf_executor.ControlTrigger

  func.func @main() {
    tf_executor.graph {
      %control = tf_executor.island {
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    tf_executor.graph {
      %control = tf_executor.island {
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }

    tf_executor.graph {
      %control = tf_executor.island {
        "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @partitioned}> : (tensor<*x!tf_type.resource>) -> ()
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }

    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Const"() {device = "", value = dense<"value"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
      %outputs_0, %control_1 = tf_executor.island wraps "tf.Const"() {device = "", value = dense<"value"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>

func.func @if_then(%arg0: tensor<i1>) -> tensor<i1> {
  // expected-remark@above {{ID: 5}}
  %graph = tf_executor.graph {
  // expected-remark@above {{ID: 3}}
    %island:2 = tf_executor.island {
    // expected-remark@above {{ID: 1}}
      tf_executor.yield %arg0 : tensor<i1>
      // expected-remark@above {{ID: 0}}
    }
    tf_executor.fetch %island#0 : tensor<i1>
    // expected-remark@above {{ID: 2}}
  }

    %0 = tf_executor.graph {
      %outputs_0, %control_0 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<501000x128xf32>>>

  func.func @argliveinotherislands(%arg0: tensor<?xi32>) -> tensor<?xi32> {
    %0 = tf_executor.graph {
      // CHECK: %[[OTHER_ISLAND_OUTPUT:[a-z0-9]*]], %{{.*}} = tf_executor.island wraps "tf.D"
      %1:2 = tf_executor.island wraps "tf.D"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

      %2:2 = tf_executor.island {
        // CHECK: %[[TPU0_OUTPUT:[0-9]*]] = "tf_device.launch"
        // CHECK: <{device = "tpu0"}>

  // CHECK: tf_executor.graph
  tf_executor.graph {
    %handle0, %handle0_control = tf_executor.island wraps "tf.SomeOp"() : () -> !tf_res
    %handle1, %handle1_control = tf_executor.island wraps "tf.SomeOp"() : () -> !tf_res
    %control_A = tf_executor.island wraps "tf.OpA"() : () -> ()
    %while_out:2, %while_control = tf_executor.island(%control_A) wraps "tf.While"(

  return llvm::isa<tf_executor::EnterOp, tf_executor::ExitOp,
                   tf_executor::IslandOp, tf_executor::MergeOp,
                   tf_executor::SwitchOp>(op);
}

// Assigns all data results to a specified device.
void PopulateDeviceForOpResults(
    Operation& op, llvm::StringRef device,
    llvm::DenseMap<Value, llvm::StringRef>& value_to_device) {
  Operation* op_to_update = &op;