TF_ASSIGN_OR_RETURN(
        module_, test::GetMlirModuleFromString(module_string, &context_));

    context_.loadAllAvailableDialects();
    return absl::OkStatus();
  }

  Status LegalizeSingleOp(Operation& op) {
    SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);

    OpBuilder op_builder(&op);
    EmptyPatternRewriter pattern_rewriter(op_builder);

    LogicalResult result =