op_result.replaceAllUsesExcept(cast, cast.getDefiningOp());
      }
    }
    return WalkResult::advance();
  });

  if (result.wasInterrupted()) return signalPassFailure();
}

}  // namespace

  auto func = getOperation();
  auto* ctx = func.getContext();
  patterns.add<PushDownDequantize>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    signalPassFailure();
  }
}
}  // namespace

// Creates an instance of the TensorFlow Lite optimize op order pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeOpOrderPass() {

  RewritePatternSet patterns(ctx);
  patterns.add<MergeFusionWithUniformDequantizePattern>(ctx);
  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    signalPassFailure();
  }
}

}  // namespace

  for (auto &entry_func : entry_funcs) {
    if (EncapsulateFirstXlaCompilablePartitionedCalls(
            entry_func, symbol_table_collection, symtab)
            .failed()) {
      return signalPassFailure();
    }
  }
}

std::unique_ptr<OperationPass<ModuleOp>> CreateXlaClusterFormationPass() {
  return std::make_unique<XlaClusterFormationPass>();
}

}  // namespace internal

    target.addLegalOp<TF::StatefulPartitionedCallOp>();
  }

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
  if (failed(applyPartialConversion(func, target, frozen_patterns))) {
    return signalPassFailure();
  }
}

struct TFToStablehloOptions : public PassPipelineOptions<TFToStablehloOptions> {
  Option<bool> skip_quantization_ops{*this, "skip-quantization-ops",

  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError()
        << "quant-convert-tf-custom-aggregator-op-to-quant-stats failed.";
    signalPassFailure();
  }
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertCustomAggregationOpToQuantStatsPass() {
  return std::make_unique<ConvertCustomAggregationOpToQuantStatsPass>();

  MLIRContext* ctx = module_op.getContext();

  func::FuncOp main_func_op = FindMainFuncOp(module_op);
  if (!main_func_op) {
    module_op.emitError("Main function op not found.");
    return signalPassFailure();
  }

  GraphOp main_graph_op = GetGraphOpFromFuncOp(main_func_op);
  if (!main_graph_op) return;

  SessionInitializerOp session_init_op = GetSessionInitializerOp(module_op);

  if (failed(applyPartialConversion(getOperation(), target,
                                    std::move(patterns)))) {
    getOperation().emitError("mhlo to TFLite legalization failed.");
    signalPassFailure();
  }
}
}  // namespace

void PopulateLegalizeHloToTFLitePatterns(RewritePatternSet* patterns,
                                         MLIRContext* context) {

        assert(replicated_input->use_empty());
        replicated_input->erase();
        return WalkResult::advance();
      });

  if (result.wasInterrupted()) {
    signalPassFailure();
    return;
  }

  getOperation()->walk([](TF::TPUPartitionedInputV2Op partitioned_input) {
    if (partitioned_input->use_empty()) partitioned_input->erase();
  });
}

}  // namespace

      has_invalid_ops = true;
      IncrementCounterFor(mlir_failed_legalization_op_count, op);
    }
    return WalkResult::advance();
  });

  if (has_invalid_ops) signalPassFailure();
}

}  // namespace

std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateVerifyTFXLALegalizationPass(bool legalize_chlo) {
  return std::make_unique<VerifyTFXLALegalization>(legalize_chlo);