void runOnOperation() override;

 private:
  bool ShouldPreserveOp(Operation* op);
  bool ShouldPreserveIsland(IslandOp island);
  void PruneGraph(GraphOp graph);

  llvm::SmallDenseSet<mlir::StringAttr, 4> ops_to_preserve_ids_;
};

// Checks if a tf_executor.Graph can be pruned.
// For TensorFlow V1.0 compatibility: when importing a graph without providing

};

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

    if (!matchPattern(input, m_Constant(&arg_value))) {
      return failure();
    }
    auto attr_name = signature.getArgAttrOfType<StringAttr>(
        operand.index(), kAttrArgumentNameAttr);
    auto attr_type = signature.getArgAttrOfType<StringAttr>(
        operand.index(), kAttrArgumentTypeAttr);
    auto value = ProcessAttributeValue(arg_value, attr_type);

  // propagated to the reduce function call.
  reduce_call->setAttr(kDeviceAttr,
                       reduce_dataset->getAttrOfType<StringAttr>(kDeviceAttr));
  reduce_call->setAttr(
      TF::kCompileDeviceTypeAttr,
      reduce_dataset->getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr));

  SmallVector<Value, 4> if_returns;

  builder.create<TF::YieldOp>(loc,

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