context_->loadDialect<arith::ArithDialect, mlir::func::FuncDialect,
                          tensor::TensorDialect, mlir::TF::TensorFlowDialect,
                          TensorFlowLiteDialect>();
    builder_ = std::make_unique<mlir::Builder>(context_.get());
    fused_lstm_func_ = createLstmCompositeFunc(builder_.get(), false, false);
    fused_lstm_func_cifg_ =

          registry.insert<mlir::BuiltinDialect>();
          registry.insert<mlir::arith::ArithDialect>();
          registry.insert<mlir::func::FuncDialect>();
          registry.insert<mlir::TFL::TensorFlowLiteDialect>();
          return registry;
        }()) {
    context_.loadAllAvailableDialects();

    mlir::OpBuilder builder(&context_);
    module_ = builder.create<mlir::ModuleOp>(builder.getUnknownLoc());

 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(
    const std::vector<uint64_t> &explicitly_consumed_ids,
    std::vector<bool> &is_consumed_id,

    return "Fallback TF-Quantization-unsupported ops to TFLite Flex ops.";
  }

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

 private:
  // The mode of TF Quantization, might indicate different users/devices.
  Option<std::string> mode_{*this, "mode",

    : 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() {}

 private:

    target.addDynamicallyLegalDialect<TensorFlowLiteDialect>([](Operation* op) {
      auto tfl_op = dyn_cast_or_null<TflRuntimeVerifyOpInterface>(op);
      if (!tfl_op) return false;
      return succeeded(tfl_op.VerifyTflRuntimeConstraints(
          op, /*emit_error_on_verify_fail=*/false));
    });
  } else {
    target.addLegalDialect<TensorFlowLiteDialect>();
  }

  registry.insert<mlir::arith::ArithDialect, mlir::func::FuncDialect,
                  mlir::quant::QuantizationDialect,
                  mlir::quantfork::QuantizationForkDialect,
                  mlir::TFL::TensorFlowLiteDialect, mlir::TF::TensorFlowDialect,
                  mlir::stablehlo::StablehloDialect, mlir::vhlo::VhloDialect>();
  mlir::func::registerAllExtensions(registry);
  context->appendDialectRegistry(registry);

    }
    tf_dialect_ =
        module.getContext()->getOrLoadDialect<mlir::TF::TensorFlowDialect>();
    tfl_dialect_ = module.getContext()
                       ->getOrLoadDialect<mlir::TFL::TensorFlowLiteDialect>();
    vhlo_dialect_ =
        module.getContext()->getOrLoadDialect<mlir::vhlo::VhloDialect>();
    // Right now the TF executor dialect is still needed to build NodeDef.
    module.getContext()

        /*max_bcast_rank=*/6);
  }
  return false;
}

//===----------------------------------------------------------------------===//
// TensorFlowLiteDialect
//===----------------------------------------------------------------------===//

struct TensorFlowLiteInlinerInterface : public DialectInlinerInterface {
  using DialectInlinerInterface::DialectInlinerInterface;