// graph VersionDef proto.
void PopulateTfVersions(mlir::ModuleOp module, const VersionDef& versions);

// Extracts TensorFlow GraphDef version information from the given module.
// Returns failure if version attribute is missing or any of the sub attributes
// are invalid.
mlir::LogicalResult ExtractTfVersions(mlir::ModuleOp module,
                                      VersionDef* versions);

    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
    if (!mlir_module_) {
      return absl::Status(
          absl::StatusCode::kNotFound,
          absl::StrCat("Could not find MLIR module at ", mlir_module_path));
    }
    return absl::OkStatus();
  }

  DialectRegistry registry_;
  MLIRContext context_;
  OwningOpRef<mlir::ModuleOp> mlir_module_;
  Env* env_;
  std::string test_dir_;

}

mlir::OwningOpRef<mlir::ModuleOp> HloToMlirHloTranslateFunction(
    mlir::StringRef input, mlir::MLIRContext* context,
    bool import_all_computations) {
  xla::HloProto hlo_proto;
  std::string content(input.data(), input.size());
  if (!LoadHloProto(content, &hlo_proto)) {
    LOG(ERROR) << "Failed to load proto";
    return nullptr;
  }

  mlir::OwningOpRef<mlir::ModuleOp> module =

    if (resource_analyzer.IsPotentiallyWritten(arg)) {
      return false;
    }
  }
  return true;
}

GlobalTensorUsesMap CreateGlobalTensorUsesMap(ModuleOp module) {
  GlobalTensorUsesMap global_tensor_uses;

  SymbolTable symbol_table(module);
  for (auto func : module.getOps<func::FuncOp>()) {
    for (size_t i = 0, e = func.getNumArguments(); i < e; i++) {

#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/util/tensor_bundle/byte_swap_tensor.h"

namespace tensorflow {

static absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GraphdefToMlirImport(
    llvm::StringRef input, const std::vector<std::string>& input_arrays,
    const std::vector<std::string>& input_dtypes,
    const std::vector<std::optional<std::vector<int>>>& input_shapes,

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

    mlir::func::registerAllExtensions(registry);
    context->appendDialectRegistry(registry);
    source_mgr->AddNewSourceBuffer(std::move(buffer), llvm::SMLoc());
    return mlir::OwningOpRef<mlir::ModuleOp>(
        mlir::parseSourceFile<mlir::ModuleOp>(*source_mgr, context));
  }

  mlir::Location loc =
      mlir::FileLineColLoc::get(context, input_filename, /*line=*/0,
                                /*column=*/0);

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 =