}

  return success();
}

void LiftHashTableOpsAsArgsPass::runOnOperation() {
  auto module_op = getOperation();

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (failed(LiftHashTableOpsToArguments(module_op, func_op))) {
      signalPassFailure();
      return;
    }
  }
}

static PassRegistration<LiftHashTableOpsAsArgsPass> pass;

}  // namespace

  func::FuncOp func_op = getOperation();
  MLIRContext* context = func_op.getContext();
  RewritePatternSet patterns(context);

  BFloat16TypeConverter converter;
  patterns.add<BFloat16TypePattern, BitcastConvertOpPattern>(converter,
                                                             context);
  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,

}

bool IsMhloUniformQuantizedOp(Operation& op) {
  return llvm::isa<mhlo::UniformQuantizeOp, mhlo::UniformDequantizeOp>(op);
}

void VerifyQuantLegalization::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    // Verify all uq and qint types are lowered.
    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||
        llvm::any_of(op->getResultTypes(), IsQuantType) ||

  OpPassManager &func_pm = pm.nest<func::FuncOp>();

  // First operates on the executor dialect:
  // - remove dead islands.
  // - fuse islands as much as possible.
  // - materialize the eventual "pass-through" ops by inlining their content.
  func_pm.addPass(tf_executor::CreateTFExecutorGraphPruningPass());
  func_pm.addPass(tf_executor::CreateTFExecutorIslandCoarseningPass());

  pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsInClusterPass());
  // Encode this in its own scope so that func_pm is not mistakenly used
  // later on.
  {
    OpPassManager& func_pm = pm.nest<FuncOp>();
    func_pm.addPass(mlir::TFTPU::CreateTPUHostComputationExpansionPass());
    func_pm.addPass(mlir::TFTPU::CreateTPUUpdateEmbeddingEnqueueOpInputsPass());
  }

  // functions and avoid converting ConstOps inside initializer functions.
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    // Do not unfreeze constants under these functions.
    if (func_op.getSymName().contains("while_body")) continue;
    if (func_op.getSymName().contains("while_cond")) continue;
    absl::c_copy_if(func_op.getOps<TF::ConstOp>(),
                    std::back_inserter(target_const_ops),

// tf_device.replicate.
Operation* FindCompilePredecessor(Operation* func_op,
                                  Operation* preprocess_op) {
  llvm::DenseSet<Block*> blocks = GetAllBlocksBetween(preprocess_op, func_op);

  llvm::DenseMap<Block*, Operation*> scope;
  Operation* o = preprocess_op;
  while (o && o != func_op) {
    scope[o->getBlock()] = o;
    o = o->getParentOp();
  }

  absl::flat_hash_set<FunctionName> existing_func_names;
  module_op->walk([&](func::FuncOp func_op) {
    FunctionName func_name = func_op.getSymName().str();
    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) {

  auto func_op_quant_scale_spec = GetStableHloQuantConstraints;

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    // The function might contain more stats ops than required, and it will
    // introduce requantize if the calibration stats have conflicts. This tries
    // to remove all the redundant stats ops.
    RemoveRedundantStatsOps(func_op, func_op_quant_spec,
                            func_op_quant_scale_spec);

 private:
  void runOnOperation() override;
};

void ConvertXlaCallModuleOpToBfloat16Pass::runOnOperation() {
  Operation* func_op = getOperation();
  SymbolTableCollection symbol_table;
  OpBuilder builder(&getContext());

  auto result = func_op->walk([&](TF::XlaCallModuleOp op) {
    // Converts the serialized StableHLO module to bfloat16.
    auto result =