#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "tensorflow/core/platform/status.h"

namespace tensorflow {

// Prints a MLIR module `module_op` and returns it as a string.
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,

#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/platform/errors.h"

namespace tensorflow {

std::string SerializeMlirModule(mlir::ModuleOp module_op) {
  std::string serialized_mlir_module;
  llvm::raw_string_ostream os(serialized_mlir_module);
  mlir::OpPrintingFlags print_flags;
  print_flags.enableDebugInfo();

  if (!mlir::SetTPUInfeedLayout(mlir_module))
    return errors::Internal("Failed to set layouts attribute");

  TF_ASSIGN_OR_RETURN(
      auto compiled_mlir,
      CompileSerializedMlirToXlaHlo(
          SerializeMlirModule(mlir_module.get()), arg_shapes, device_type,
          use_tuple_args, true, shape_determination_fns, compilation_result,
          custom_legalization_passes, metadata.module_name(),
          lower_to_xla_hlo));

  return module_ref;
}

static mlir::LogicalResult MlirModuleToSerializedMlirStringAttrTranslate(
    mlir::ModuleOp module_op, llvm::raw_ostream& output) {
  output << "\"";
  std::string serialized_module = SerializeMlirModule(module_op);
  llvm::printEscapedString(serialized_module, output);
  output << "\"";
  return mlir::success();
}

static mlir::LogicalResult MlirTfToHloTextTranslateFunction(

      clone.setPublic();
    } else {
      clone.setPrivate();
    }
    symbol_table.insert(clone);
  }

  *serialized_func_module =
      tensorflow::SerializeMlirModule(module_for_func.get());
  return success();
}

// Create a `tf._TPUCompileMlir` that contains a MLIR module that is
// functionally equivalent to the function referenced by cluster_func.

      use_tuple_args, enable_op_fallback, use_return_tuple,
      shape_determination_fns, custom_legalization_passes, module_name,
      lower_to_xla_hlo));

  auto mlir_compilation = SerializeMlirModule(module_op);

  // Only attempt to fill in the compilation result's IO info if lowering.
  if (lower_to_xla_hlo) {
    TF_RETURN_IF_ERROR(PopulateCollectiveInfo(module_op, compilation_result));

      ModuleOp::create(mlir::UnknownLoc::get(func.getContext()));
  SymbolTable symbol_table(module_for_func.get());

  symbol_table.insert(func);
  *serialized_func_module =
      tensorflow::SerializeMlirModule(module_for_func.get());
}

// Returns whether `op` or ops nested in `op` are outside compiled.
bool HasOutsideCompilationNested(Operation* op) {
  return op
      ->walk([&](Operation* walked_op) {