%1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  // CHECK:      [[COMPUTATION:%.+]] = "tf_device.parallel_execute"()
  // CHECK:      "tf.OpA"([[READ0]])
  %2 = "tf_device.parallel_execute"() ({
    "tf_device.launch"() ({
      "tf.OpA"(%1) : (tensor<i32>) -> ()
      tf_device.return
    }) {device = "TPU_REPLICATED_HOST_0"} : () -> ()
    tf_device.return
  }, {

  "tf_device.parallel_execute"() ({
// expected-error@-1 {{'tf_device.parallel_execute' op number of output types (3) must match the total number of outputs from all regions (2).}}
    %0 = "tf.opA"() : () -> (tensor<*xi1>)
    %1 = "tf.opB"() : () -> (tensor<*xi32>)
    tf_device.return %0, %1 : tensor<*xi1>, tensor<*xi32>
  },
  {
    %2 = "tf.opC"() : () -> (tensor<*xi1>)
    tf_device.return

	case config.PolicyOPASubSys:
		// In case legacy OPA config is being set, we treat it as if the
		// AuthZPlugin is being set.
		subSys = config.PolicyPluginSubSys
		fallthrough
	case config.PolicyPluginSubSys:
		if ppargs, err := polplugin.LookupConfig(s, GetDefaultConnSettings(), xhttp.DrainBody); err != nil {
			return err
		} else if ppargs.URL == nil {
			// Check if legacy opa is configured.

      "tf.OpA"(%0) : (tensor<f32>) -> ()
    }
    ```

    will be transformed into:

    ``mlir
    %0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32>
    tf_device.replicate {n = 2 : i32} {
      "tf.OpA"(%0) : (tensor<f32>) -> ()
    }
    ```

//               devices = {DEVICE_ALIAS_0 = ["/DEVICE:0", "/DEVICE:1"],
//                          DEVICE_ALIAS_1 = ["/DEVICE:2", "/DEVICE:3"]}} {
//     %a = "tf_device.launch"() ({
//       %2 = "tf.opA"(%ri) : (tensor<i1>) -> tensor<i1>
//       tf_device.return %2 : tensor<i1>
//     }) {device = "DEVICE_ALIAS_0"} : () -> tensor<i1>
//     %b = "tf_device.launch"() ({
//       %3 = "tf.opB"(%a) : (tensor<i1>) -> tensor<i1>

func.func @simple_independent_chains(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>) {
  // CHECK: tf_executor.graph {
  tf_executor.graph {
    // CHECK: %[[A_CONTROL:.*]] = tf_executor.island wraps "tf.OpA"() : () -> ()
    %control_A = tf_executor.island wraps "tf.OpA"() : () -> ()
    // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>

func.func @island_not_direct_parent_of_user() -> () {
  tf_executor.graph {
    %island1:2 = tf_executor.island {
      // CHECK: %[[VAL_0:.*]], %[[VAL_1:.*]] = tf_executor.island wraps "tf.OpA"() : () -> tensor<i64>
      %0 = "tf.OpA"() : () -> (tensor<i64>)
      // CHECK: %[[VAL_2:.*]] = tf_executor.island(%[[VAL_1]]) wraps "tf.OpB"() : () -> ()
      "tf.OpB"() : () -> ()
      tf_executor.yield %0 : tensor<i64>
    }

 during TPURewritePass.

 For example, the tf.OpA with the `mini_batch_splits` attribute will be
 moved after _TPUCompileMlir and the first input will use the
 _TPUCompileMlir program output:

 ```mlir
 "tf_device.launch"() ({
  %cst_0 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
  "tf.OpA"(%cst_0) { mini_batch_splits = ""} : (tensor<1x!tf_type.string>) -> ()

    // CHECK-SAME:        _xla_has_host_transfer = true
    // CHECK:         }, {
    // CHECK:           "tf_device.cluster"
    // CHECK:             %[[A:.+]] = "tf.OpA"
    // CHECK:             %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) <{dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<2x2xi64>) -> tensor<1x2xi64>

          %0 = "tf.opA"(%arg0) : (tensor<f32>) -> tensor<f32>
          return %0 : tensor<f32>
        }
      }
    }
    ```

    will be transformed into:

    ```mlir
    module  {
      func @foo(%arg0: tensor<f32>) -> tensor<f32> {
        %0 = tf_executor.graph {
          %outputs, %control = tf_executor.island {
            %1 = "tf.opA"(%arg0) : (tensor<f32>) -> tensor<f32>