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;

setattr(_current_module, "initializers", _initializers)


# Do an eager load for Keras' code so that any function/method that needs to
# happen at load time will trigger, eg registration of optimizers in the
# SavedModel registry.
# See b/196254385 for more details.
try:
  if _tf_uses_legacy_keras:
    importlib.import_module("tf_keras.src.optimizers")
  else:
    importlib.import_module("keras.src.optimizers")

    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),

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Insert quantized functions into the module";
  }

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

 private:
  void runOnOperation() override;

  // Returns the function library for the given quantization method and opset
  // pair.

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

  // Explicitly disable dumping Op details on failures.
  module.getContext()->printOpOnDiagnostic(false);

  mlir::DialectRegistry registry;
  mlir::func::registerAllExtensions(registry);
  module.getContext()->appendDialectRegistry(registry);

  mlir::StatusScopedDiagnosticHandler status_handler(module.getContext(),
                                                     /*propagate=*/true);

          PyExc_ValueError,
          "Failed to parse opdefs at index %d of custom op defs argument: %s",
          i, tf_opdefs);
      return nullptr;
    }

    // Register extra opdefs to TensorFlow global op registry.
    tensorflow::OpRegistry::Global()->Register(
        [opdef](
            tensorflow::OpRegistrationData* op_reg_data) -> tensorflow::Status {
          *op_reg_data = tensorflow::OpRegistrationData(opdef);

 public:
  explicit TPUClusterFormationPass(bool strict_clusters)
      : strict_clusters_(strict_clusters) {}

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<mlir::tf_device::TensorFlowDeviceDialect>();
  }

  void runOnOperation() override;

 private:
  bool strict_clusters_;
};

    // This is a brief description of the pass.
    return "Replace quantization candidates with composite functions into the "
           "module";
  }

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

  void runOnOperation() override;

 private:
  QuantizationOptions quant_options_;
  bool test_mode_;
  Option<OpSet> op_set_{

    // This is a brief description of the pass.
    return "Quantize composite functions with QDQ input/outputs.";
  }

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

 private:
  void runOnOperation() override;