namespace tensorflow {
namespace {
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GetMlirModuleFromString(
    llvm::StringRef string, mlir::MLIRContext* context) {
  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
  auto status =

           py::return_value_policy::reference);

  py::class_<mlir::ModuleOp>(m, "ModuleOp")
      .def("create",
           [](mlir::Location loc) { return mlir::ModuleOp::create(loc); })
      .def("push_back",
           [](mlir::ModuleOp& m, mlir::func::FuncOp f) { m.push_back(f); })
      .def("dump", &mlir::ModuleOp::dump)
      .def("getAsStr", [](mlir::ModuleOp& m) {
        std::string str;
        llvm::raw_string_ostream os(str);

// Passes only used for testing purposes.

include "mlir/Pass/PassBase.td"

def TestPreCalibrationComponentPass : Pass<"stablehlo-test-pre-calibration-component", "mlir::ModuleOp"> {
  let summary = "Test-only pass to test the PreCalibrationComponent.";
  let description = [{
    Runs the pre calibration passes for post-training quantization with default
    configuration.
  }];

#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace {

using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::PassManager;
using ::tensorflow::monitoring::testing::CellReader;

absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* context) {
  mlir::DialectRegistry mlir_registry;

  // debugging purposes.
  static constexpr absl::string_view kName = "quant_ptq_pre_calibration";

  explicit PreCalibrationComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  absl::Nonnull<MLIRContext*> ctx_;
};

}  // namespace mlir::quant::stablehlo

std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeTFOpsPass(
    std::optional<ModuleOp> tfr_module = std::nullopt);

// Rewrites quantized operands and results with their storage types.
// This pass should be run at module level after decomposition, if there are
// quantized operands or results.
std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteQuantizedIOPass();

// Raise to TF ops.

  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);
  auto module =
      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
  ASSERT_TRUE(module);

  std::vector<std::string> inputs;
  std::vector<std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
      in_specs;

  // debugging purposes.
  static constexpr absl::string_view kName = "quant_ptq_post_calibration";

  explicit PostCalibrationComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

  void AddPasses(
      OpPassManager& pm,
      const ::stablehlo::quantization::QuantizationSpecs& specs,

  return absl::OkStatus();
}

// Converts arg_shapes to xla::Shape's and store into xla_input_shapes.
Status GetXlaInputShapes(
    mlir::ModuleOp module, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
    bool use_tuple_args,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    std::vector<xla::Shape>* xla_input_shapes) {

// Legalizes TF/XLA communication ops (TF dialect) to HLO dialect communication
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass();

// Legalizes TF/XLA collective ops (TF dialect) to HLO dialect collective
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass();

// Verifies that the TF/XLA ops have all been lowered to MHLO.