module.getLoc(), name,
      FunctionType::get(module.getContext(), /*inputs=*/{}, /*results=*/{}));

  // Make savedmodel verification happy.
  init->setAttr("tf_saved_model.exported_names",
                builder.getStrArrayAttr({name}));

  builder.setInsertionPointToStart(init.addEntryBlock());

  for (func::FuncOp func : tf_saved_model::GetInitializerFunctions(module)) {
    if (func.getNumArguments() != 0) {

  // Set the insertion point right before the return op.
  builder.setInsertionPoint(&region.back().back());

  StringAttr output_file_path_attr = builder.getStringAttr(output_file_path);
  ArrayAttr ids_attr = builder.getStrArrayAttr(ids);
  ArrayAttr calibration_methods_attr =
      builder.getI32ArrayAttr(calibration_methods);
  builder.create<TF::CalibrationStatisticsSaverOp>(
      region.getLoc(), statistics_outputs, output_file_path_attr, ids_attr,

                   builder.getIntegerAttr(builder.getIntegerType(64), 1));
  cluster->setAttr(kTopologyAttr, builder.getStringAttr(""));
  cluster->setAttr(kDeviceAssignmentAttr,
                   builder.getStrArrayAttr(llvm::ArrayRef<llvm::StringRef>(
                       {"bad_device_assigment"})));

  mlir::TF::RuntimeDevices runtime_devices;
  std::string host_device;
  EXPECT_TRUE(mlir::failed(

// contain the given string.
void SetExportedName(func::FuncOp f, StringRef name) {
  OpBuilder b(f);
  f->removeAttr(kTfSavedModelExportedNamesAttr);
  f->setAttr(kTfSavedModelExportedNamesAttr, b.getStrArrayAttr({name}));
}

// Convert a savedmodel public function into a private function.
// This means we need to remove any attributes that are only allowed
// on exported (public) functions.
void Unexport(func::FuncOp f) {

                        main_func.getFunctionType().getResults()));
  // Name of the main function for the eventual executable needs to be set.
  main_func->setAttr(kTfSavedModelExportedNamesAttr,
                     builder.getStrArrayAttr({func_name}));
  main_func->setAttr(
      kTfSavedModelInitializerTypeAttr,
      builder.getStringAttr(kTfSavedModelInitializerRestoreType));
  module->setAttr("tf_saved_model.semantics",

    new_arg_attrs.emplace_back(arg_attr);
  }

  const NamedAttribute index_path_attr(
      builder.getStringAttr(kTfSavedModelIndexPathAttr),
      builder.getStrArrayAttr({index_path}));

  new_arg_attrs.emplace_back(index_path_attr);
  return new_arg_attrs;
}

// Strips the "assets/" directory prefix, if `filename` begins with it. The

  // Update input/output sharding attributes on tf_device.cluster_func op.
  cluster_func->setAttr(tensorflow::kInputShardingAttr,
                        GetStrArrayAttr(builder, sharding_for_args));
  cluster_func->setAttr(tensorflow::kOutputShardingAttr,
                        GetStrArrayAttr(builder, sharding_for_rets));
  return mlir::success();
}

void TPUShardingIdentificationPass::runOnOperation() {

  auto func = builder.create<func::FuncOp>(module->getLoc(),
                                           kSessionInitFuncName, func_type);
  func->setAttr(kTfSavedModelExportedNamesAttr,
                builder.getStrArrayAttr({kSessionInitFuncName}));
  func->setAttr(kTfSavedModelInitializerTypeAttr,
                builder.getStringAttr(kTfSavedModelInitializerRestoreType));
  func.setVisibility(mlir::func::FuncOp::Visibility::Public);

              temp.push_back(conf_str);
            }
            ArrayRef<StringRef> values(temp);
            auto named_attr =
                builder->getNamedAttr(key, builder->getStrArrayAttr(values));
            attrs.push_back(named_attr);
          }

          break;
        }
        case flexbuffers::FBT_VECTOR: {
          if (std::string{key} == "dimension_numbers") {

  // newly created argument.
  main_func_op.setArgAttr(new_file_prefix_arg.getArgNumber(),
                          /*name=*/kTfSavedModelIndexPathAttr,
                          /*value=*/builder.getStrArrayAttr({kTfFilePrefix}));

  // Append the "__tf_file_prefix:0" to the "tf.entry_function" attribute's
  // item keyed by "inputs".
  AddEntryFunctionInput(Twine(kTfFilePrefix).concat(":0").str(), main_func_op);