if (auto split = llvm::dyn_cast<TF::SplitOp>(op)) {
      if (HasDevice(split) || split->getAttrOfType<ArrayAttr>(kClassAttr))
        return WalkResult::advance();
      for (Operation* pred : IslandPredecessors(split)) {
        if (auto colocation_classes =
                pred->getAttrOfType<ArrayAttr>(kClassAttr)) {
          split->setAttr(kClassAttr, colocation_classes);
          return WalkResult::advance();

  void runOnOperation() override {
    Builder builder(&getContext());
    Dialect* tf = getContext().getLoadedDialect<TensorFlowDialect>();
    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_));

    // attribute modulo the number of TPU cores in the system.
    for (auto op : require_device_ordinal_ops) {
      if (op->getAttrOfType<IntegerAttr>(kDeviceOrdinalAttr).getInt() == -1) {
        OpBuilder builder(op);
        int device_ordinal =
            op->getAttrOfType<IntegerAttr>(kReplicaIdAttr).getInt() %
            device_num;
        op->setAttr(kDeviceOrdinalAttr,

namespace mlir {
namespace TF {

using ::tensorflow::kValidDeviceTypes;

LogicalResult HasValidCompilationAndReplicationAttributes(Operation& op) {
  auto replicate_attr = op.getAttrOfType<StringAttr>(kReplicationInfoAttr);
  auto compile_attr = op.getAttrOfType<StringAttr>(kCompileDeviceTypeAttr);
  if (!replicate_attr && !compile_attr) return success();
  if (!replicate_attr || !compile_attr)

      }
    } 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)) {

  double real_multiplier = scale_product / o_spec.getScale();
  output_multipliers->push_back(QuantizeMultiplier(real_multiplier));

  // output ranges
  auto min = rop->getAttrOfType<FloatAttr>("min");
  auto max = rop->getAttrOfType<FloatAttr>("max");
  output_ranges->push_back(CalculateQuantizedRange(
      o_spec.getScale(), o_spec.getZeroPoint(),
      (min ? std::optional<double>(min.getValueAsDouble()) : std::nullopt),

  StringAttr shared_name_attr = op->getAttrOfType<StringAttr>(kSharedNameAttr);
  if (shared_name_attr) {
    auto shared_name = StringRef(shared_name_attr.getValue());
    if (!shared_name.empty()) {
      return shared_name;
    }
  }
  // Attempt to parse "_class" attribute if there is no "shared_name"
  // attribute.
  ArrayAttr classes_attr = op->getAttrOfType<ArrayAttr>(kClassAttr);
  if (!classes_attr) {

  // The sizes of indices, shared_name and container must be equal.
  int32_t indices_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("indices").size();
  int32_t shared_name_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("shared_name").size();
  int32_t container_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("container").size();

  if (!(indices_size == shared_name_size &&

ResourceHandle GetResourceHandle(Operation *op) {
  llvm::StringRef device;
  if (auto attr = op->getAttrOfType<StringAttr>("device")) {
    device = attr.getValue();
  }

  llvm::StringRef container;
  if (auto attr = op->getAttrOfType<StringAttr>("container")) {
    container = attr.getValue();
  }

  llvm::StringRef shared_name;
  if (auto attr = op->getAttrOfType<StringAttr>("shared_name")) {
    shared_name = attr.getValue();
  }

}

mlir::LogicalResult ExtractTfVersions(mlir::ModuleOp module,
                                      VersionDef* versions) {
  versions->Clear();
  auto version_attr =
      module->getAttrOfType<mlir::DictionaryAttr>("tf.versions");
  if (!version_attr) return mlir::failure();

  auto producer =
      mlir::dyn_cast_or_null<mlir::IntegerAttr>(version_attr.get("producer"));
  if (!producer) return mlir::failure();