class XlaCallModuleDeserializationPass
    : public impl::XlaCallModuleDeserializationPassBase<
          XlaCallModuleDeserializationPass> {
 public:
  void getDependentDialects(mlir::DialectRegistry &registry) const override {
    XlaCallModuleDeserializationPassBase::getDependentDialects(registry);
    mlir::func::registerAllExtensions(registry);
  }

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

// Applies prepare quantization on the model in TF dialect for dynamic range
// quantization case.
class PrepareQuantizeDRQPass
    : public PassWrapper<PrepareQuantizeDRQPass, OperationPass<ModuleOp>> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, ::mlir::quant::QuantizationDialect,
                    ::mlir::quantfork::QuantizationForkDialect>();
  }

 public:

  StringRef getArgument() const final { return "constant-fold-broadcast-pass"; }
  StringRef getDescription() const final {
    return "Constant folds BroadcastInDimOp before binary elementwise ops";
  }
  void getDependentDialects(::mlir::DialectRegistry &registry) const override {}

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.add<FoldBroadcastInDimBeforeMulOp>(&getContext());

    }
    return success();
  }
};

class LiftTfliteFlexOpsPass
    : public impl::LiftTfliteFlexOpsPassBase<LiftTfliteFlexOpsPass> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LiftTfliteFlexOpsPass)

 public:
  StringRef getArgument() const final { return "unfuse-mhlo-batch-norm-pass"; }
  StringRef getDescription() const final {
    return "Unfuses MHLO batch norm inference op into arithmetic ops";
  }
  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
    // TODO(b/299514833): Remove TensorFlowDialect usage.
    registry.insert<shape::ShapeDialect, mlir::TF::TensorFlowDialect>();
  }

  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Insert custom aggregation ops for the calibration procedure";
  }

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

  void runOnOperation() override;

 private:
  enum TestCase {
    TEST_CASE_MIN_MAX,

// passes. Currently, per-channel quantization only supports 1D results.
class PreprocessOpPass
    : public PassWrapper<PreprocessOpPass, OperationPass<ModuleOp>> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, QuantizationDialect,
                    quantfork::QuantizationForkDialect>();
  }

 public:

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Apply graph optimizations such as fusing and constant folding to "
           "prepare lifting.";
  }

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

  void runOnOperation() override;

 private:
  Option<OpSet> op_set_{

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct DecomposeReduceDatasetPass
    : public impl::DecomposeReduceDatasetPassBase<DecomposeReduceDatasetPass> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<tf_device::TensorFlowDeviceDialect>();
  }

  void runOnOperation() override;
};

struct ReduceWhileOperandsPass
    : public impl::ReduceWhileOperandsPassBase<ReduceWhileOperandsPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReduceWhileOperandsPass)

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

LogicalResult FindImplicityProducers(