mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;

  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(mlir_module_str, &context, &mlir_module));
  TF_RETURN_IF_ERROR(SetupReturnValues(*mlir_module, metadata_proto));
  TF_RETURN_IF_ERROR(SetupArguments(*mlir_module, arg_shapes, metadata_proto));