getOperation().walk([&](Operation* op) {
      if (auto device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // We assign default device to ops with device attribute that is empty.
        if (device_attr.getValue().empty()) {
          op->setAttr(kDeviceAttr, builder.getStringAttr(default_device_));
        }
      } else if (op->getDialect() == tf) {

    } else if (StringAttr attr = arg.getDefiningOp()->getAttrOfType<StringAttr>(
                   kDeviceAttr)) {
      return op_device == attr;
    }
    // Fold tf.Identity when arg device is not defined.
    return true;
  };

  func_op.walk([&should_fold_op_func](TF::IdentityOp op) {
    StringAttr op_device = op->getAttrOfType<StringAttr>(kDeviceAttr);
    if (should_fold_op_func(op.getOperand(), op_device)) {

#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_cluster_util.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kClassAttr[] = "_class";

#define GEN_PASS_DEF_TPUCLEANUPCLUSTERATTRIBUTESPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

  op->removeAttr(StringAttr::get(context, kDeviceAttr));
  op->moveBefore(return_op);
}

void DeviceAttributeToLaunch::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");

  getOperation().walk([&](Operation* op) {
    if (op->getDialect() != tf_dialect) return WalkResult::advance();
    if (auto device = op->getAttrOfType<StringAttr>(kDeviceAttr)) {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace TF {
namespace {

constexpr const char *kDeviceAttr = "device";

#define GEN_PASS_DEF_CONSTANTOPDEVICEASSIGNMENTPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct ConstantOpDeviceAssignmentPass

// propagated through.
bool IsSupportedExecutorOp(Operation& op) {
  auto ops_have_same_device = [](Operation* lhs, Operation* rhs) {
    auto lhs_device_attr = lhs->getAttrOfType<StringAttr>(kDeviceAttr);
    auto rhs_device_attr = rhs->getAttrOfType<StringAttr>(kDeviceAttr);
    return (!lhs_device_attr && !rhs_device_attr) ||
           (lhs_device_attr && rhs_device_attr &&
            lhs_device_attr.getValue() == rhs_device_attr.getValue());
  };

    }
    auto device_attr = op->getAttr(kDeviceAttr);
    if (!device_attr) {
      op->setAttr(kDeviceAttr, launch.getDeviceAttr());
      return WalkResult::advance();
    }

    if (auto device_str_attr = mlir::dyn_cast<StringAttr>(device_attr)) {
      if (device_str_attr.getValue().empty()) {
        op->setAttr(kDeviceAttr, launch.getDeviceAttr());
        return WalkResult::advance();

#include "tensorflow/core/util/device_name_utils.h"

namespace mlir {
namespace TF {

namespace {

const char kDeviceAttr[] = "device";
const char kDeviceGpu[] = "GPU";

std::optional<std::string> GetOpDevice(mlir::Operation *op) {
  mlir::StringAttr device = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  if (!device || device.getValue().empty()) {
    return std::nullopt;
  }

#define GEN_PASS_DEF_TPUCOLOCATESPLITSPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

constexpr char kDeviceAttr[] = "device";
// Attribute of colocation classes.
constexpr char kClassAttr[] = "_class";

bool HasDevice(Operation* op) {
  auto attr = op->getAttrOfType<StringAttr>(kDeviceAttr);
  if (!attr) return false;
  return !attr.getValue().empty();
}

      }
      if (op->hasAttr(tensorflow::kCompileDeviceTypeAttr)) {
        return mlir::WalkResult::advance();
      }
      auto device_attr = op->getAttrOfType<StringAttr>(tensorflow::kDeviceAttr);
      if (!device_attr) {
        op->setAttr(tensorflow::kCompileDeviceTypeAttr,
                    builder.getStringAttr(tensorflow::kEmptyDevice));
        return mlir::WalkResult::advance();
      }