#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/OwningOpRef.h"  // from @llvm-project

namespace mlir {

// Set layouts attribute of tf.InfeedDequeueTuple ops.
bool SetTPUInfeedLayout(OwningOpRef<ModuleOp>& mlir_module);

// Try to determine the right TPU infeed layout.
FailureOr<Attribute> GetTPUInfeedLayout(ArrayRef<Type> types,
                                        OpBuilder& rewriter);

  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(computation.mlir_module, &context, &mlir_module));
  if (!mlir::SetTPUInfeedLayout(mlir_module))
    return errors::Internal("Failed to set layouts attribute");

  TF_ASSIGN_OR_RETURN(
      auto compiled_mlir,
      CompileSerializedMlirToXlaHlo(

  } else {
    // types.size() == 1 and types[0] == TokenType
    // For this case, we return an empty array attribute.
    return rewriter.getArrayAttr({});
  }
}

bool SetTPUInfeedLayout(mlir::OwningOpRef<mlir::ModuleOp> &mlir_module) {
  auto res = mlir_module->walk([&](mlir::TF::InfeedDequeueTupleOp op) {
    mlir::OpBuilder builder(op.getContext());
    std::vector<mlir::Type> result_types;

  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
  TF_RETURN_IF_ERROR(DeserializeMlirModule(mlir_computation.mlir_module,
                                           &context, &mlir_module));
  if (!mlir::SetTPUInfeedLayout(mlir_module))
    return errors::Internal("Failed to set layouts attribute");

  if (VLOG_IS_ON(2)) {
    tensorflow::DumpMlirOpToFile("legalize_with_old_bridge", mlir_module.get());
  }