for (func::FuncOp init_func_op :
       tf_saved_model::GetInitializerFunctions(module)) {
    for (auto func : module.getOps<func::FuncOp>()) {
      auto dict_attr =
          func->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function");
      if (!dict_attr) continue;

      OpBuilder builder(func.getContext());
      builder.setInsertionPointToStart(&func.getBlocks().front());

    if (const auto data_format = op->getAttrOfType<StringAttr>("data_format")) {
      is_supported_affine_op =
          data_format.getValue() == "NHWC" || data_format.getValue() == "NDHWC";
    }
  } else if (llvm::isa<TF::BatchMatMulV2Op>(op)) {
    if (const auto adj_y = op->getAttrOfType<BoolAttr>("adj_y")) {
      is_supported_affine_op = !adj_y.getValue();
    }

  // execution.
  if (llvm::isa<mlir::func::ReturnOp, mlir::quantfork::StatisticsOp>(op))
    return std::nullopt;

  if (!HasValidHardwareTarget(op)) return std::nullopt;

  auto device = op->getAttrOfType<mlir::StringAttr>(mlir::TFL::tac::kDevice);
  if (device == nullptr) return std::nullopt;

  llvm::StringRef device_name_str = device.getValue();
  return device_name_str.str();
}

      // This attribute is used for op colocation. Since all ops are located
      // on a single device cluster, this private attribute is no longer
      // needed.
      op->removeAttr(kClassAttr);
      if (auto attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // Preserve device attribute if the op is placed on a replicated core
        // device. Device attribute is used to infer the appropriate sharding
        // within TPUs for this op.

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

namespace mlir {
namespace TF {

std::string GetDeviceAttrAsResourceInstanceStr(mlir::Operation* op) {
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  // Treat missing device attribute like unspecified (= empty string) attribute.
  // Note that different op instances with the same string (including empty

  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)) {
      if (!device.getValue().empty()) WrapOpInLaunch(op, device.getValue());
    }
    return WalkResult::advance();
  });
}

}  // anonymous namespace

}  // namespace

bool IsTfFuncCustomCall(stablehlo::CustomCallOp op) {
  return op.getCallTargetName() == kTfTargetName;
}

DictionaryAttr GetTfBackendConfig(stablehlo::CustomCallOp op) {
  return op->getAttrOfType<DictionaryAttr>(kTfBackendConfigAttrName);
}

FailureOr<SymbolRefAttr> GetTfFuncCustomCallFuncName(
    stablehlo::CustomCallOp op) {
  if (!IsTfFuncCustomCall(op)) {
    return success(nullptr);
  }

  StringRef device;
  StringRef container;
  StringRef shared_name;
};

ResourceKey GetResourceKey(TF::VarHandleOp var_handle_op) {
  ResourceKey resource_key;

  if (auto attr = var_handle_op->getAttrOfType<StringAttr>("device")) {
    resource_key.device = attr.getValue();
  }

  resource_key.container = var_handle_op.getContainer();
  resource_key.shared_name = var_handle_op.getSharedName();

  return resource_key;

    return status;
  }
  return mlir_module;
}

std::string GetDevice(Operation* op) {
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  return device_attr ? device_attr.getValue().str() : "";
}

bool CanBeIgnoredInCluster(Operation* op) {
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  return !device_attr || device_attr.getValue().empty();
}

  llvm::SmallVector<IslandOp, 8> islands_to_outline;
  getOperation().walk([&](IslandOp island_op) {
    auto parent_func = island_op->getParentOfType<func::FuncOp>();
    auto skip_island_outlining =
        parent_func->getAttrOfType<BoolAttr>(mlir::TF::kSkipIslandOutlining);
    if (skip_island_outlining && skip_island_outlining.getValue()) {
      // Island was marked to be skipped.
      return WalkResult::advance();
    }