%1 = "tf.A"() : () -> (tensor<2x!tf_type.string>)
     "tf_device.launch"() ({
       "tf.TPUExecuteAndUpdateVariables"(%arg1, %1) {device_var_reads_indices = [0], device_var_updates_indices = [-1]} : (tensor<*x!tf_type.resource<tensor<4xf32>>>, tensor<2x!tf_type.string>) -> ()
       tf_device.return
    }) {device = "TPU_REPLICATED_CORE_0"} : () -> ()

            // CHECK: "tf.TPUExecuteAndUpdateVariables"(%[[ID]], %[[R1]], %[[COMPILE]]#1)
            "tf_device.launch"() ({
              "tf.TPUExecuteAndUpdateVariables"(%id, %arg31, %compile#1)
                    {device_var_reads_indices = [0, 1], device_var_updates_indices = [0, 1]}
                      : (tensor<*x!tf_type.resource<tensor<f32>>>, tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>, tensor<2x!tf_type.string>) -> ()
              tf_device.return

  llvm::SmallVector<int64_t, 8> device_var_reads_indices;
  llvm::SmallVector<int64_t, 8> device_var_updates_indices;
  for (auto resource : var_access_info.resources_read) {
    auto info_it = var_access_info.per_resource_info.find(resource);
    if (info_it == var_access_info.per_resource_info.end()) continue;
    device_var_reads_indices.push_back(info_it->second.execute_input_index);

    "tf.TPUExecuteAndUpdateVariables"(%4, %arg0) {
        device = "foobar",
        device_var_reads_indices = [],
        device_var_updates_indices = []} : (
                tensor<2048xi32>,
                tensor<2x!tf_type.string>) -> ()
    return
}

// -----

// CHECK-LABEL test_identity

    // CHECK: [[exe:%.*]] = tf_executor.island({{.*}}) wraps "tf.TPUExecuteAndUpdateVariables"
    %exe_control = tf_executor.island wraps "tf.TPUExecuteAndUpdateVariables"(%arg0, %arg0, %key) {
        device_var_reads_indices = [0, 1],
        device_var_updates_indices = [0, 1],
        device = "task:0"
    } : (!tf_res, !tf_res, !tf_str) -> ()

  // CHECK-SAME: <{device = "/job:localhost/replica:0/task:0/device:TPU:0"}>
  // CHECK-NEXT: "tf.TPUExecuteAndUpdateVariables"(%[[ID_0]], %[[ARG_1]], %[[READ_2]], %[[COMPILE]]#1)
  // CHECK-SAME: device_var_reads_indices = [0, 1],
  // CHECK-SAME: device_var_updates_indices = [0, -1]
  %execute:2 = "tf_device.launch"() ({
    %0:2 = "tf.TPUExecute"(%read0, %read1, %read2, %compile#1) {

    %read2, %read2_control  = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) : (!tf_res) -> tensor<f32>
    tf_executor.island wraps "tf.TPUExecuteAndUpdateVariables"(%ret0, %ret1, %key) {
        device_var_reads_indices = [0, 1],
        device_var_updates_indices = [0, 1]} : (!tf_res, !tf_res, !tf_str) -> ()
    tf_executor.fetch %ret0, %ret1: !tf_res, !tf_res
  }
  func.return %graph#0, %graph#1 : !tf_res, !tf_res
}

      %3 = "tf.AssignVariableOp"(%arg0, %2)
    ```

    will be transformed into

    ```mlir
      %2 = "tf.TPUExecuteAndUpdateVariables"(%arg0, %arg1, %compile)
        { device_var_reads_indices = [0, 1],
          device_var_updates_indices = [0, -1] }
    ````

    The transformation happens only for on-device variables. The above

  // expected-error@below {{requires 'device_var_reads_indices' to be the same size as number of resource handles in 'args' (1), but got 2}}
  "tf.TPUExecuteAndUpdateVariables"(%arg0, %arg1) {device_var_reads_indices = [0, 1], device_var_updates_indices = [0]} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<3x!tf_type.string>) -> ()
  func.return
}

  %3 = "tf.AssignVariableOp"(%arg0, %2)
```

will be transformed into

```mlir
  %2 = "tf.TPUExecuteAndUpdateVariables"(%arg0, %arg1, %compile)
    { device_var_reads_indices = [0, 1],
      device_var_updates_indices = [0, -1] }
````

The transformation happens only for on-device variables. The above
transformation requires `%arg0`, `%arg1` to have the same device assignment