}

  StringRef getDescription() const override {
    return "Merge the save function's ops to the main function. The save "
           "function will be removed after the pass.";
  }

  void runOnOperation() override;
};

// Returns true iff func_op has either no Region or the body has no Blocks.
bool IsFuncOpEmpty(func::FuncOp func_op) {
  return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}

class WhileOutlinePass : public impl::WhileOutlinePassBase<WhileOutlinePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(WhileOutlinePass)
  explicit WhileOutlinePass() {}

 private:
  void runOnOperation() override;

  // Outlines the regions of the WhileOp's cond and body and insert function
  // calls instead,
  void OutlineWhile(WhileOp while_op);

  // Get unique name by using the loc to name mapping.

struct LegalizeWhilePass
    : public impl::LegalizeWhilePassBase<LegalizeWhilePass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeWhilePass)
  void RunOnFunction(func::FuncOp func);

  void runOnOperation() override {
    for (auto op : getOperation().getOps<func::FuncOp>()) RunOnFunction(op);
  }
};

}  // namespace

// Inserts call to the given function into the 'region'.

          LegalizeStablehloCustomCallToCompositePass> {
  using LegalizeStablehloCustomCallToCompositePassBase::
      LegalizeStablehloCustomCallToCompositePassBase;

  void runOnOperation() override {
    MLIRContext *context = &getContext();

    ConversionTarget target(*context);
    target.addLegalDialect<mlir::stablehlo::StablehloDialect>();
    target.addLegalDialect<mlir::func::FuncDialect>();

class UnfoldLargeSplatConstantPass
    : public impl::UnfoldLargeSplatConstantPassBase<
          UnfoldLargeSplatConstantPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnfoldLargeSplatConstantPass)

  void runOnOperation() override {
    auto module = getOperation();

    mlir::OpBuilder op_builder(&module.getBodyRegion());
    module.walk([&](mlir::arith::ConstantOp const_op) {
      MaybeUnfoldLargeSplatConstant(&op_builder, const_op);

  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, arith::ArithDialect>();
  }

  void runOnOperation() override;
};

// Generate an einsum equation from the given DotDimensionNumber.
std::string CreateEinsumEquation(
    const xla::DotDimensionNumbers& dot_dimension_numbers, const int lhs_rank,

class PushTransposeThroughEwisePass
    : public impl::PushTransposeThroughEwisePassBase<
          PushTransposeThroughEwisePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PushTransposeThroughEwisePass)
  void runOnOperation() override;
};

// Compute the permutation that would take `arr` to the identity.
llvm::SmallVector<int32_t> InvertPermutation(llvm::SmallVector<int32_t> arr) {
  llvm::SmallVector<int32_t> inverse_arr(arr.size());

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<TFL::TensorFlowLiteDialect, TF::TensorFlowDialect>();
  }
  void runOnOperation() override;
};

LogicalResult FindImplicityProducers(
    const std::vector<uint64_t> &explicitly_consumed_ids,
    std::vector<bool> &is_consumed_id,
    const std::vector<std::vector<uint64_t>> &dependency_graph) {

    registry.insert<TF::TensorFlowDialect>();
    registry.insert<quant::QuantizationDialect>();
    registry.insert<quantfork::QuantizationForkDialect>();
  }

 private:
  void runOnOperation() override;

  Option<DebuggerType> debugger_type_{
      *this, "debugger_type",
      llvm::cl::init(DebuggerConfig::DEBUGGER_TYPE_UNSPECIFIED),
      llvm::cl::values(

  }

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

  void runOnOperation() override;

 private:
  enum TestCase {
    TEST_CASE_MIN_MAX,
    TEST_CASE_AVERAGE_MIN_MAX,
    TEST_CASE_HISTOGRAM_PERCENTILE,
    TEST_CASE_HISTOGRAM_MSE_BRUTEFORCE,