}

  if (input_mlir) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::TFL::TensorFlowLiteDialect, mlir::arith::ArithDialect,
                    mlir::func::FuncDialect>();
    mlir::func::registerAllExtensions(registry);
    context->appendDialectRegistry(registry);
    source_mgr->AddNewSourceBuffer(std::move(buffer), llvm::SMLoc());
    return mlir::OwningOpRef<mlir::ModuleOp>(

// Sparsifies model to encode sparse tensors with proper format. Throws error if
// sparsification fails.
PyObject* MlirSparsifyModel(PyObject* data);

// Registers the given custom opdefs to TensorFlow global op registry.
PyObject* RegisterCustomOpdefs(PyObject* list);

// Returns the collected TFLite conversion errors.
std::vector<std::string> RetrieveCollectedErrors();

// Returns MLIR string dump of the given Flatbuffer model.

    // 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

                               std::vector<TensorShape>& arg_shapes,
                               tpu::TPUCompileMetadataProto& metadata_proto) {
  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;

  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(mlir_module_str, &context, &mlir_module));

  llvm::cl::ParseCommandLineOptions(argc, argv,
                                    "TF Saved Model to Stablehlo converter\n");

  mlir::DialectRegistry registry;
  RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);
  context.loadAllAvailableDialects();

  bool is_input_mlir_module = absl::EndsWith(input_path, ".mlir");
  std::vector<std::string> exported_model_signatures_in_vector =

  for (const auto& entry : denylisted_ops) {
    denylisted_mlir_op_names.insert(TfLiteToMlir(entry));
  }

  DialectRegistry registry;
  registry.insert<mlir::TFL::TensorFlowLiteDialect>();
  MLIRContext context(registry);
  StatusScopedDiagnosticHandler statusHandler(&context,
                                              /*propagate=*/true);

    StringPiece tfr_raw_text, mlir::MLIRContext* mlir_ctx) {
  mlir_ctx->allowUnregisteredDialects(/*allow=*/true);
  // Load dialects involved in the conversion
  mlir::DialectRegistry registry;
  // clang-format off
  registry.insert<mlir::arith::ArithDialect,
                  mlir::func::FuncDialect,
                  mlir::scf::SCFDialect,
                  mlir::shape::ShapeDialect,

  EXPECT_EQ(non_categorized_count, 428);
}

// Just a counter test to see which ops have duplicate lowerings. This isn't a
// correctness test versus a test to easily ensure the op registry is kept
// in sync.
TEST_F(LegalizationOpConfigTest, CountTypesWhichHaveBothMlirAndTf2xlaFallback) {
  int double_lowering_count = 0;

  DialectRegistry dialect_registry;

    return "Replace quantization candidates with composite functions into the "
           "module for post-training dynamic range case";
  }

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

  void runOnOperation() override;

 private:
  Option<OpSet> target_opset_{
      *this, "target-opset", llvm::cl::init(OpSet::TF),

// name should be unique to the graph it is being inserted to. If the
// `ignore_unregistered_attrs` argument is set to true, the attributes which are
// not in the op registry will be ignored. If the `ignore_unregistered_attrs`
// argument is not set to true, _output_shapes attribute is added to nodes with
// ShapedType for the leading values with ShapedType in the results of the