module_op, /*initializer_type=*/kTfSavedModelInitializerInitType);

  EXPECT_THAT(init_func_op, NotNull());
  EXPECT_THAT(init_func_op.getSymName(), "init_func");
  EXPECT_THAT(
      init_func_op->getAttrOfType<StringAttr>(kTfSavedModelInitializerTypeAttr),
      kTfSavedModelInitializerInitType);
}

TEST_F(TfSavedModelTest, GetInitializerFunctionNoMatchingInitializerType) {
  constexpr StringRef kModuleOpStr = R"mlir(

                      ArrayRef<const Device *> results) override {
    if (llvm::isa<TF::TPUExecuteOp>(op) ||
        llvm::isa<TF::TPUExecuteAndUpdateVariablesOp>(op)) {
      auto device = op->getAttrOfType<StringAttr>(kDevice);
      for (auto *operand : operands)
        propagateIfChanged(operand, operand->SetDevice(device));
    } else {
      // Propagate device through other ops. These ops might have their

// "arg0:0,arg1:0"}`, then this function returns `{"arg0:0", "arg1:0"}`.
SmallVector<StringRef> GetEntryFunctionInputs(func::FuncOp func_op) {
  auto entry_function_attr =
      func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");

  SmallVector<StringRef> inputs;
  mlir::dyn_cast_or_null<StringAttr>(entry_function_attr.get("inputs"))
      .strref()
      .split(inputs, /*Separator=*/",");

LogicalResult AssignDevicesInRegion(const Dialect* tf_dialect,
                                    tf_device::LaunchOp launch,
                                    Region& region) {
  auto parallel_group_attr =
      launch->getAttrOfType<StringAttr>(TF::kParallelExecAnnotation);
  auto result = region.walk([&](Operation* op) -> WalkResult {
    if (op->getDialect() != tf_dialect) return WalkResult::advance();

    if (parallel_group_attr) {

  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());
    // Skip non-tpu device cluster_func.
    auto cluster_id =
        cluster_func_op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
    if (!cluster_id) return WalkResult::advance();

    llvm::SmallVector<int, 4> dynamic_shape_arg_index;

    // Traverse the operands of the cluster func op and find which operand

      DeviceNameUtils::AddressSpace(parsed_name));
  return result.empty() ? kLocalhost : result;
}

std::string GetHost(Operation *op) {
  std::string device = "";
  if (StringAttr attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
    device = attr.getValue().str();
  }
  return GetHost(device);
}

// The device is considered to be on the localhost iff one of the following is
// true:

  auto clean_status = MakeCleanup([status] { TF_DeleteStatus(status); });

  // Builds TF operation and sets all the attributes.
  std::string node_name = "unnamed";
  if (auto attr = inst->getAttrOfType<mlir::StringAttr>("name")) {
    node_name = std::string(attr.getValue());
  }
  auto node_def_or = ConvertTFDialectOpToNodeDef(
      inst, node_name.c_str(), /*ignore_unregistered_attrs=*/true);

    existing_func_names.insert(func_name);
    // We may retrieve the original function's name from the attribute.
    // Functions without this attribute are ignored.
    auto original_func_name =
        func_op->getAttrOfType<StringAttr>("tf._original_func_name");
    if (original_func_name) {
      if (auto alias_itr = function_aliases.find(original_func_name.str());
          alias_itr != function_aliases.end()) {

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

    return "Inserts the main function to the module.";
  }

  void runOnOperation() override;
};

// Checks if a FuncOp is exported.
bool IsExported(func::FuncOp op) {
  auto exported_names =
      op->getAttrOfType<ArrayAttr>(kTfSavedModelExportedNamesAttr);
  return exported_names && !exported_names.empty();
}

// Check if a function is an entry function.
bool IsEntryFunction(func::FuncOp op) {