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,

    if (clone.getName() == entry_func.getName()) {
      // We can simply change name of TPU program's main function because there
      // should be no other reference to it.
      clone.setName(StringAttr::get(clone.getContext(), "main"));
      clone.setPublic();
    } else {
      clone.setPrivate();
    }
    symbol_table.insert(clone);
  }

  *serialized_func_module =

// initializer function. Assumes that there exists such an attribute.
std::string GetInitializerType(func::FuncOp init_func_op) {
  return init_func_op
      ->getAttrOfType<StringAttr>(kTfSavedModelInitializerTypeAttr)
      .str();
}

// An initializer function should satisfy the follwing conditions:
// * Its GraphOp should only have control outputs.

  auto input_attr = tf_attrs.get("inputs");
  if (!input_attr) {
    return errors::InvalidArgument("no inputs attribute found");
  }
  auto input_names = mlir::cast<mlir::StringAttr>(input_attr).getValue();
  input_names.split(function_input_names, ",", /*MaxSplit=*/-1,
                    /*KeepEmpty=*/false);
  const int function_input_names_size = function_input_names.size();

      return func.getArgAttrOfType<mlir::StringAttr>(
          resource_arg.getArgNumber(), kFuncDeviceAttr);
    } else {
      return mlir::Attribute{};
    }
  }
}

bool IsCompositeDevice(mlir::Attribute attr) {
  const auto str_attr = attr.dyn_cast_or_null<mlir::StringAttr>();
  return str_attr &&

  size_t numRead;
  mlir::Attribute attr = mlir::parseAttribute(input, context, {}, &numRead);
  if (!attr || !mlir::isa<mlir::StringAttr>(attr)) {
    LOG(ERROR) << "Input is not parsable as a MLIR StringAttr.";
    return nullptr;
  }
  auto str_attr = mlir::cast<mlir::StringAttr>(attr);

  mlir::DialectRegistry registry;
  RegisterMlirInputDialects(registry);
  context->appendDialectRegistry(registry);

           << "requires device ordinal from device " << logical_core_name
           << " to be present in 'tf.device.replicate' op";
  }
  llvm::StringRef tpu_device =
      tpu_replica.cast<ArrayAttr>()[replica_id].cast<StringAttr>().getValue();
  return tensorflow::GetDeviceOrdinalFromDeviceString(op->getLoc(), tpu_device,
                                                      &device_ordinal);
}

// Checks whether the given op contains QuantizationTrait::FullyQuantizable.
inline bool HasQuantizableTrait(Operation* op) {
  return op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
         op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue().str() ==
             QuantTraitValues[QuantizationTrait::FullyQuantizable];
}

// returns `std::nullopt`.
std::optional<llvm::StringRef> GetTpuClusterName(Operation* op) {
  if (auto tpu_status = op->getAttrOfType<StringAttr>(kTpuStatusAttr)) {
    // Borrow cluster name from TPU status (for `TPUCompilationResult` op).
    return tpu_status.getValue();
  }
  auto device_type = op->getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr);
  if (!device_type || device_type.getValue() != TF::kTpuDevice) {