apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nvidia-gpu-device-plugin
  namespace: kube-system
  labels:
    k8s-app: nvidia-gpu-device-plugin
    addonmanager.kubernetes.io/mode: EnsureExists
spec:
  selector:
    matchLabels:
      k8s-app: nvidia-gpu-device-plugin
  template:
    metadata:
      labels:
        k8s-app: nvidia-gpu-device-plugin
    spec:
      priorityClassName: system-node-critical
      affinity:

    %1 = "tf.Shape"(%ri) {device = "", T = "tfdtype$DT_FLOAT", out_type = "tfdtype$DT_INT32"} : (tensor<*xf32>) -> tensor<?xi32>
    %2 = "tf.opA"(%1) {device = "TPU_REPLICATED_CORE_0"} : (tensor<?xi32>) -> tensor<*xi32>
    %3 = "tf_device.launch"() ({
      %b = "tf.opB"(%1) : (tensor<?xi32>) -> tensor<*xi32>
      tf_device.return %b : tensor<*xi32>
    }) {device = "TPU_REPLICATED_CORE_0"} : () -> tensor<*xi32>

    // The name of the compilation device (e.g., "XLA_CPU_JIT");
    string compilation_device_name;

    // If 'use_multiple_streams' is true, we create separate streams for
    // compute, host-to-device, and device-to-host communication.
    bool use_multiple_streams = false;

    // If true, the XLA devices with the same device ordinal will share the same
    // compute stream. Otherwise each XLA device will having their own compute

  %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
  %1 = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
  %2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>

// CHECK-SAME: device = "CPU:0"
// CHECK:      %[[C:.*]], {{.*}} = tf_executor.island wraps "tf.opC"(%[[B]])
// CHECK-SAME: device = "CPU:0"
// CHECK:      %[[D:.*]], {{.*}} = tf_executor.island wraps "tf.opD"
// CHECK-NOT:  "tf_device.launch"
// CHECK-NOT:      tf_executor.yield %[[A]], %[[C]], %[[B]], %[[D]]


// Tests empty device string attributes are overwritten.

  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_1dd4fef4-646d-491f-a3a8-bf5334f45813", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
  "tf.LookupTableImportV2"(%0, %cst, %cst_1) {device = ""} : (tensor<!tf_type.resource>, tensor<3x!tf_type.string>, tensor<3xi64>) -> ()

"909", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg7: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "911", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg8: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "913", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg9: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "915", tf.device = "/job:t...

                {n = 2 : i32,
                 devices = {DEVICE_ALIAS_0 = ["/DEVICE:0", "/DEVICE:1"],
                            DEVICE_ALIAS_1 = ["/DEVICE:2", "/DEVICE:3"]}} {
  // Inside the region, %0, %2, %4, and %6 corresponds to
  // "/DEVICE:0"/"/DEVICE:2" and %1, %3, %5, and %7 corresponds to
  // "/DEVICE:1"/"/DEVICE:3", depending on which device alias is used.
  %k = "tf_device.launch"() ( {

// `container`, `name`, and `device` can be set. If an resource handle is tied
// to an instance of an operation (e.g. TensorFlow runtime operation caching),
// `op` can be set instead.
struct ResourceHandle {
  ResourceHandle(StringRef container, StringRef name, StringRef device,
                 Operation* op)
      : container(container), name(name), device(device), op(op) {}

  // CHECK-SAME: "TPU_REPLICATED_CORE_0"
  // CHECK: "tf.TPUCopyWithDynamicShape"
  func.func @valid_copy_op_in_replicated_host(
    %arg0: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"},
    %arg1: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<2048xi32>, tensor<2048xi32>) {
    %cst = "tf.Const"() {value = dense<1024> : tensor<i32>} : () -> tensor<i32>