// `{ tac.device = "GPU", tac.inference_type = "FLOAT"}` to a function
// with the matching attributes. Assumed is that device type "CPU"
// is the only device that is allowed to call other devices. I.e. ancestors of a
// "CPU" `Operation` may only `Operations` without a device or other "CPU"
// `Operations`. Implied is that "CPU" ops may contain subgraphs of different

  for (const auto& device : devices) {
    if (inference_type == QUANTIZED_INT8) {
      all_device_inference_types.push_back({device, QUANTIZED_INT8});
    } else if (inference_type == QUANTIZED_UINT8) {
      all_device_inference_types.push_back({device, QUANTIZED_UINT8});
    }

    // We will alway enable float.
    all_device_inference_types.push_back({device, FLOAT});
  }

      // TODO(bfontain): Check for other attributes.
      replicated_output->setAttr(kDevice, builder.getStringAttr(""));
      TF::TPUReplicatedInputOp input = builder.create<TF::TPUReplicatedInputOp>(
          op->getLoc(), result.getType(), replicated_output.getResults());
      input->setAttr(kDevice, builder.getStringAttr(""));
      mlir::Value new_value = input.getOutput();

          // Set interface_name & target to the call_op as well.
          new_call->setAttr(kInterfaceNameAttr,
                            builder->getStringAttr(interface_name));
          new_call->setAttr(
              kDevice,
              builder->getStringAttr(preferred_inference_device_type.hardware));
          new_call->setAttr(
              kInferenceType,
              builder->getStringAttr(GetInferenceString(

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-resource-device-inference %s | FileCheck %s

!tf_res = tensor<*x!tf_type.resource<tensor<32xf32>>>

// Tests that the pass can correctly propagate device attributes inside the same
// function.

// CHECK-LABEL: func @propagate_in_function
func.func @propagate_in_function(
  %arg0: !tf_res {tf.device = "/TPU:0"},
  %arg1: !tf_res {tf.device = "/TPU:1"}) {
  tf_executor.graph {

// RUN: tac-opt-all-backends -tfl-device-transform-gpu %s -split-input-file -verify-diagnostics | FileCheck %s

func.func @pack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> {
  %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return %0 : tensor<2x1xf32>
}

// CHECK:   func @pack(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> {

  ASSERT_FALSE(arrived);
  TFE_TensorHandle* hdevice =
      TFE_TensorHandleCopyToDevice(hcpu, context, name, status.get());
  ASSERT_TRUE(arrived);
  ASSERT_FALSE(executed);
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> matmul(
      MatMulOp(context, hcpu, hdevice), TFE_DeleteOp);
  TFE_OpSetDevice(matmul.get(), name, status.get());

                                                                               \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("VarHandleOp").Device(DEVICE).HostMemory("resource"), VarHandleOp); \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("_VarHandlesOp").Device(DEVICE).HostMemory("resources"),            \

      %control_31 = tf_executor.island wraps "tf.NoOp"() {device = "/device:CPU:0"} : () -> ()
      %outputs_32, %control_33 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<4> : tensor<i32>} : () -> tensor<i32>
      %control_34 = tf_executor.island wraps "tf.NoOp"() {device = "/device:CPU:0"} : () -> ()

  device_setup_.AddDevicesAndSetUp({DEVICE_GPU});
  Device* device = device_setup_.GetDevice(DEVICE_GPU);
  XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);

  ResourceMgr resource_mgr("");
  OpKernelContext::Params params;
  params.resource_manager = &resource_mgr;
  params.device = device;
  OpKernelContext ctx(&params, 0);

  PjRtDeviceCompiler* pjrt_device_compiler = nullptr;