};

// Check if op directly uses a key in `virtual_devices`.
bool DirectUseOfVirtualDevice(const DictionaryAttr& virtual_devices,
                              Operation* op) {
  StringAttr op_device = op->getAttrOfType<StringAttr>(kDeviceAttr);
  if (!op_device) return false;
  if (virtual_devices.get(op_device.getValue())) return true;
  return false;
}

// Check if op or its ancestor uses a key in `virtual_devices`.

    return failure();
  }

  // Builds TF operation and sets all the attributes.
  std::string node_name = "unnamed";
  if (const StringAttr attr = inst->getAttrOfType<StringAttr>("name")) {
    node_name = std::string(attr.getValue());
  }
  absl::StatusOr<std::unique_ptr<tensorflow::NodeDef>> node_def =
      tensorflow::ConvertTFDialectOpToNodeDef(

// attribute to be passed to the a tfl.average_pool_2d that can fully replace
// this composite (here, padding is done directly by the tfl.average_pool_2d as
// opposed to being extracted into a separate tfl.pad).
StringAttr GetAvgPoolOpPadAttr(Builder& builder, mhlo::CompositeOp op);

// Get dense attr for a matrix that corrects the over counting of divisors when
// casting an average pool with ceil mode on in terms of average pool with it
// off.

      // 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.

    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();
}

  mlir::Builder builder(context);
  StringAttr resource_name_id = builder.getStringAttr(kResourceNameArgAttr);

  SmallSet<StringRef, 4> resource_names;

  for (func::FuncOp func : module.getOps<func::FuncOp>()) {
    for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
      auto resource_arg =
          func.getArgAttrOfType<StringAttr>(i, kResourceNameArgAttr);
      if (!resource_arg) continue;

#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
  // string) are seen as dependent (same resource instance).

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();
}

// Returns the predecessors of `op` when `op`'s predecessors are wrapped by
// islands.

                                            SymbolTable& symbol_table) {
  if (!call_op->hasAttr(kOriginalStablehloEntryFunctionAttrName)) {
    return;
  }

  const auto original_function_name = call_op->getAttrOfType<StringAttr>(
      kOriginalStablehloEntryFunctionAttrName);
  const auto current_function_name = call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);

        /*config=*/builder.getStringAttr(""),
        /*config_proto=*/builder.getStringAttr(""),
        /*executor_type=*/builder.getStringAttr(""));
    call_op->setAttr("device", launch->getAttrOfType<StringAttr>("device"));
    launch.replaceAllUsesWith(call_op);
    launch.erase();
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateConvertLaunchFuncToTFCallPass() {