// ops.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateTPURewritePass(
    llvm::StringRef module_name = llvm::StringRef());

// Creates a pass that adds ops which perform formatting on variables at
// run-time according to compilation result.
std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUVariableRuntimeReformattingPass();

#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/status.h"

namespace tensorflow {
namespace tf2xla {
namespace v1 {
namespace {

using mlir::ModuleOp;

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

static constexpr char kCompilationStreamz[] =

      return %0 : tensor<1x3xf32>
    }
  }
)mlir";

TEST_F(LiftAsFunctionCallTest, LiftedFunctionSucceeds) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleLifted);
  ASSERT_TRUE(module_op);

  auto composite_dot_general_fn =
      module_op->lookupSymbol<func::FuncOp>("composite_dot_general_fn_1");

// 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++) {

std::string SerializeMlirModule(mlir::ModuleOp module_op);

// Parses a MLIR module from `mlir_module_string` into `mlir_module` with
// context `mlir_context`.
Status DeserializeMlirModule(llvm::StringRef serialized_mlir_module,
                             mlir::MLIRContext* mlir_context,
                             mlir::OwningOpRef<mlir::ModuleOp>* mlir_module);

}  // namespace tensorflow

#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,

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