mlir::DialectRegistry registry;
  registry.insert<mlir::func::FuncDialect, mlir::stablehlo::StablehloDialect,
                  mlir::TFL::TFLDialect, mlir::arith::ArithDialect,
                  mlir::TF::TensorFlowDialect, mlir::chlo::ChloDialect>();

  return failed(
      mlir::MlirOptMain(argc, argv, "ODML Converter Driver\n", registry));

  registry.insert<mlir::tensor::TensorDialect>();
  registry.insert<mlir::tosa::TosaDialect>();
  registry.insert<mlir::xla_framework::XLAFrameworkDialect,
                  mlir::TF::TensorFlowDialect, mlir::tf_type::TFTypeDialect>();
}

TEST_F(LegalizationOpConfigTest, CountLoweringsSet) {
  int mlir_lowering_count = 0;
  int tf2xla_fallback_count = 0;
  int non_categorized_count = 0;

  DialectRegistry dialect_registry;
  dialect_registry.insert<mlir::TF::TensorFlowDialect>();

  MLIRContext context(dialect_registry);
  context.loadAllAvailableDialects();

  for (auto operation : context.getRegisteredOperations()) {
    if (IsTypeLegalizedWithMlir(operation.getTypeID())) {

  MLIRContext context;
  OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
  OpBuilder builder(&module->getBodyRegion());
  context.loadDialect<TF::TensorFlowDialect>();

  Value value = CreateConstValue<int8_t>(builder, module->getLoc(),
                                         unpacked_shape, unpacked_values);

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

  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";

  if (auto cost_function =
          mlir::dyn_cast<tfrt::compiler::CostFunctionInterface>(op)) {
    cost_map_[op] = cost_function.cost();
    return;
  }

  if (!llvm::isa<mlir::TF::TensorFlowDialect>(op->getDialect())) {
    cost_map_[op] = max_arg_size_;
    return;
  }

  // Try to use its cost function if it is registered.
  const auto& registry = GetCostFunctionRegistry();