}
    }

    auto device_attr = op.getAttrOfType<StringAttr>(kDeviceAttr);
    int64_t device_id;
    if (!device_attr) {
      device_id = kAnyDevice;
    } else if (devices.find(device_attr) != devices.end()) {
      device_id = devices[device_attr];
    } else {
      device_id = 1 + devices.size();
      devices[device_attr] = device_id;
    }

  let summary = "Transfer a CPU variable tensor to device.";
  let description = [{
    Transfer a CPU variable tensor to device if the variable has not been
    transferred before.

    Example:
      %device_var = gpurt.maybe_transfer_variable %cpu_var
  }];

  let arguments = (ins TFTensorType);
  let results = (outs TFTensorType);
  let assemblyFormat = "operands attr-dict";
}

#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/session_utils.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/rendezvous.h"
#include "tensorflow/core/framework/resource_var.h"
#include "tensorflow/core/framework/tensor.h"

  }

  SmallVector<func::FuncOp, 25> funcs_to_be_processed;
  // We only process if func has device annotations.
  for (auto func : module.getOps<func::FuncOp>()) {
    auto device_attr = func->getAttrOfType<StringAttr>(kDevice);
    if (device_attr != nullptr) funcs_to_be_processed.push_back(func);
  }

  OpBuilder builder(module);
  // Go head to process those funcs.

#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/lift_variables.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"

#include "tensorflow/compiler/tf2xla/xla_expression.h"
#include "xla/client/xla_builder.h"
#include "xla/client/xla_computation.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace mlir {
namespace mhlo {

class Tf2XlaRewriterTestPeer;

class Tf2XlaRewriter {
 public:

#include "tensorflow/compiler/tf2xla/functionalize_control_flow.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"

#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"

  AttrValue index_attr;
  index_attr.set_i(index);
  (*node_def->mutable_attr())["index"] = index_attr;

  if (auto device_attr =
          func.getArgAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
    *node_def->mutable_device() = device_attr.getValue().str();

  llvm::ArrayRef<mlir::NamedAttribute> func_arg_i_attrs =
      mlir::function_interface_impl::getArgAttrs(func, index);

  AttrValue index_attr;
  index_attr.set_i(index);
  (*node_def->mutable_attr())["index"] = index_attr;

  if (auto device_attr =
          func.getArgAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
    *node_def->mutable_device() = device_attr.getValue().str();

  llvm::ArrayRef<mlir::NamedAttribute> func_arg_i_attrs =
      mlir::function_interface_impl::getArgAttrs(func, index);