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

using tensorflow::tracing::TracingOperation;
using tensorflow::tracing::TracingTensorHandle;

namespace {

void RegisterDialects(mlir::MLIRContext& ctx) {
  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  ctx.appendDialectRegistry(registry);
  ctx.loadAllAvailableDialects();
}

Status ConvertDataTypeToTensor(tensorflow::DataType dtype, Builder builder,
                               Type* type) {

};

class PrepareCompositeFunctionsPass
    : public impl::PrepareCompositeFunctionsPassBase<
          PrepareCompositeFunctionsPass> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TFL::TensorFlowLiteDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareCompositeFunctionsPass)

  explicit PrepareCompositeFunctionsPass() {}

// training quantization simpler.
class PrepareQuantizePass
    : public PassWrapper<PrepareQuantizePass, OperationPass<func::FuncOp>> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, ::mlir::quant::QuantizationDialect,
                    ::mlir::quantfork::QuantizationForkDialect>();
  }

 public:

namespace {

class BreakUpIslands : public TF::PerFunctionAggregateAnalysisConsumerPass<
                           BreakUpIslands, TF::SideEffectAnalysis> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<tf_executor::TensorFlowExecutorDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BreakUpIslands)

  //
  //  Arg  --> B  --> C  --> D --> Retval
  class Subgraph {
   public:
    // Creates a graph to build the subgraph in, if it doesn't already exist,
    // using the same op registry and versions as graph_in.
    Node* MakeNodeImage(const Graph* graph_in, Node* node);

    // Returns the graph the subgraph is being built in.
    Graph* GetGraph() const;