: ctx_(stablehlo::CreateMlirContextForQuantization()),
        builder_(ctx_.get()) {
    ctx_->loadDialect<
        arith::ArithDialect, mlir::stablehlo::StablehloDialect,
        func::FuncDialect, TF::TensorFlowDialect, TFL::TensorFlowLiteDialect,
        tf_saved_model::TensorFlowSavedModelDialect,
        tf_executor::TensorFlowExecutorDialect, quant::QuantizationDialect,
        quantfork::QuantizationForkDialect>();
  }

  // Register StableHLO Quantizer pass pipelines.
  mlir::quant::stablehlo::RegisterPassPipelines();

  mlir::DialectRegistry registry;
  registry.insert<mlir::scf::SCFDialect, mlir::TF::TensorFlowDialect,
                  mlir::tf_saved_model::TensorFlowSavedModelDialect,
                  mlir::func::FuncDialect, mlir::shape::ShapeDialect,
                  mlir::arith::ArithDialect, mlir::tf_type::TFTypeDialect,

    // This is a brief description of the pass.
    return "Convert tf.CustomAggregator op to quant.Stats";
  }

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

  void runOnOperation() override;
};

class ConvertCustomAggregationOpToQuantStats

namespace mlir {
namespace quant {
namespace {

using ::testing::Eq;

class GetSizeInBytesTest : public ::testing::Test {
 protected:
  GetSizeInBytesTest() : ctx_() { ctx_.loadDialect<TF::TensorFlowDialect>(); }

  MLIRContext ctx_;
};

TF::ConstOp ParseConstOp(const absl::string_view const_op_str, Block& block,
                         MLIRContext& ctx) {
  const LogicalResult parse_result =

  LogicalResult status = EvaluateOperation(inst, inputs, constants);
  results.assign(constants.begin(), constants.end());
  return status;
}

static bool init_hooks = ([] () {
  TensorFlowDialect::RegisterConstantFoldHook(ConstantFoldFallbackHook);
}(), true);

}  // namespace TF

using ::testing::IsNull;
using ::testing::NotNull;
using ::testing::SizeIs;

class ConvertAssetArgsTest : public ::testing::Test {
 protected:
  ConvertAssetArgsTest() {
    ctx_.loadDialect<func::FuncDialect, TF::TensorFlowDialect,
                     tf_saved_model::TensorFlowSavedModelDialect>();
  }

  // Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
  // module op is valid.

    "mlir::TFL::TFLDialect",
  ];
}

def LegalizeHloToTfPass : Pass<"tf-legalize-hlo", "ModuleOp"> {
  let summary = "Legalize from MHLO to the TF dialect";
  let dependentDialects = ["TF::TensorFlowDialect"];
  let constructor = "mlir::odml::CreateLegalizeHloToTfPass()";
}

def LegalizeHloToTfLitePass : Pass<"tfl-legalize-hlo", "mlir::ModuleOp"> {
  let summary = "Legalize from MHLO to the TFLite dialect";

               << "': " << error.message() << "\n";
    return 1;
  }

  // Load the MLIR module.
  mlir::DialectRegistry registry;
  registry.insert<mlir::TF::TensorFlowDialect, mlir::TFL::TensorFlowLiteDialect,
                  mlir::arith::ArithDialect, mlir::func::FuncDialect>();
  mlir::MLIRContext context(registry);

  llvm::SourceMgr source_mgr;

  void SetUp() override {
    context_ = std::make_unique<mlir::MLIRContext>();
    context_->loadDialect<mlir::arith::ArithDialect, mlir::func::FuncDialect,
                          mlir::TF::TensorFlowDialect, TensorFlowLiteDialect>();
    builder_ = std::make_unique<mlir::Builder>(context_.get());

    fused_max_unpooling_func_ =
        createMaxUnpoolingFunc(builder_.get(), {2, 4, 4, 2}, {2, 2, 2, 2});

void populateCanonicalizationPatterns(func::FuncOp func,
                                      RewritePatternSet &patterns) {
  MLIRContext *context = func.getContext();
  mlir::Dialect *tf = context->getLoadedDialect<mlir::TF::TensorFlowDialect>();
  // Load all official canonicalization patterns. Here we skip the
  // canonicalization of the ops in the tf dialect, because they couldn't
  // propagate the attributes correctly. These optimization will be played by