%3:2 = "tf.TPUCopyWithDynamicShape"(%1, %2, %cst, %cst) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)
            // CHECK-NOT: tf.TPUAnnotateTensorsWithDynamicShape
            %4:2 = "tf.TPUAnnotateTensorsWithDynamicShape"(%3#0, %3#1) : (tensor<2048xi32>, tensor<2048xi32>) -> (tensor<2048xi32>, tensor<2048xi32>)

	  %3:2 = "tf.TPUCopyWithDynamicShape"(%1, %2, %cst, %cst) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)
	  tf_device.return %3#0, %3#1 : tensor<2048xi32>, tensor<2048xi32>
	}) {device = "TPU_REPLICATED_HOST_0"} : () -> (tensor<2048xi32>, tensor<2048xi32>)
    return %0#0, %0#1: tensor<2048xi32>, tensor<2048xi32>
  }

     %6 = "tf.Identity"(%4) : (tensor<2048xi32>) -> tensor<2048xi32>
    "tf.TPUExecute"(%6, %arg0) {device = "foobar"} : (tensor<2048xi32>, tensor<!tf_type.string>) -> ()
    return
}

// -----

// CHECK-LABEL test_disconnected
func.func @test_disconnected(%arg0: tensor<!tf_type.string>) {
     %0 = builtin.unrealized_conversion_cast to tensor<2048xi32>
     %1 = builtin.unrealized_conversion_cast to tensor<2048xi32>

      tf_executor.yield %1, %1 : tensor<2048xf32>, tensor<2048xf32>
    }
    tf_executor.fetch %outputs#0, %outputs#1 : tensor<2048xf32>, tensor<2048xf32>
  }
  func.return %0#0, %0#1 : tensor<2048xf32>, tensor<2048xf32>
}
// CHECK-LABEL: func @single_op_island_duplicate_result
// CHECK: %[[CONST:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"
// CHECK: tf_executor.fetch %[[CONST]], %[[CONST]]

  %arg7: tensor<2048x640xf32>,
  %arg8: tensor<2048x640xf32>,
  %arg9: tensor<2048xf32>,
  %arg10: tensor<2048xf32>,
  %arg11: tensor<2048xf32>,
  %arg12: tensor<2048xf32>,
  %arg13: tensor<640x2048xf32>,
  %arg14: tensor<640xf32>,
  %arg15: tensor<2048xf32>,
  %arg16: tensor<2048xf32>,
  %arg17: tensor<2048xf32>,
  %arg18: tensor<2048xf32>,
  %arg19: tensor<1x640xf32>,
  %arg20: tensor<1x2048xf32>

      tf_executor.yield %1, %1 : tensor<2048xf32>, tensor<2048xf32>
    }
    tf_executor.fetch %outputs#0, %outputs#1 : tensor<2048xf32>, tensor<2048xf32>
  }
  func.return %0#0, %0#1 : tensor<2048xf32>, tensor<2048xf32>
}

// CHECK: func @tpu_load_embedding_ops_sink_controls
// CHECK: {{%.+}}, [[READ0:%.+]] = tf_executor.island wraps "tf.ReadVariableOp"

func.func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
  %0:2 = tf_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32>
    tf_executor.fetch %outputs, %arg0 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
  }
  return %0#0, %0#1 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
}

    %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<2x8xi32>>>
    %1 = "tf.VarHandleOp"() {container = "c", shared_name = "w"} : () -> tensor<*x!tf_type.resource<tensor<2xi32>>>
    %2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>

    // CHECK:      %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}>
    // CHECK:      %[[RES_READ_VAL:[0-9]*]] = "tf.ReadVariableOp"
    // CHECK-SAME: (tensor<*x!tf_type.resource<tensor<2x8xi32>>>) -> tensor<2x8xi32>
    // CHECK:      "tf.AddV2"(%[[RES_READ_VAL]], %[[ONE]])
    // CHECK-SAME: (tensor<2x8xi32>, tensor<i32>) -> tensor<2x8xi32>
    // CHECK:      "tf.AssignVariableOp"

      func.return %arg0 : tensor<?xi32>
    }

    func.func @main(%arg0: tensor<20xi32>, %arg1: tensor<?xi32>) -> tensor<*xi32> {
      %2 = "tf.TPUPartitionedCall"(%arg0, %arg1) {f = @tpu_partitioned_call_func} : (tensor<20xi32>, tensor<?xi32>) -> tensor<*xi32>
      func.return %2 : tensor<*xi32>
    }
  })";

  CellReader<int64_t> error(kHasTpuPartitionedCallStreamzName);