We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `extensions/cpp_extension.py`
Content:
```
1 import importlib
2 import os
3 import time
4 from abc import abstractmethod
5 from pathlib import Path
6 from typing import List
7 
8 from .base_extension import _Extension
9 
10 __all__ = ["_CppExtension"]
11 
12 
13 class _CppExtension(_Extension):
14     def __init__(self, name: str, priority: int = 1):
15         super().__init__(name, support_aot=True, support_jit=True, priority=priority)
16 
17         # we store the op as an attribute to avoid repeated building and loading
18         self.cached_op = None
19 
20         # build-related variables
21         self.prebuilt_module_path = "colossalai._C"
22         self.prebuilt_import_path = f"{self.prebuilt_module_path}.{self.name}"
23         self.version_dependent_macros = ["-DVERSION_GE_1_1", "-DVERSION_GE_1_3", "-DVERSION_GE_1_5"]
24 
25     def csrc_abs_path(self, path):
26         return os.path.join(self.relative_to_abs_path("csrc"), path)
27 
28     def relative_to_abs_path(self, code_path: str) -> str:
29         """
30         This function takes in a path relative to the colossalai root directory and return the absolute path.
31         """
32 
33         # get the current file path
34         # iteratively check the parent directory
35         # if the parent directory is "extensions", then the current file path is the root directory
36         # otherwise, the current file path is inside the root directory
37         current_file_path = Path(__file__)
38         while True:
39             if current_file_path.name == "extensions":
40                 break
41             else:
42                 current_file_path = current_file_path.parent
43         extension_module_path = current_file_path
44         code_abs_path = extension_module_path.joinpath(code_path)
45         return str(code_abs_path)
46 
47     # functions must be overrided over
48     def strip_empty_entries(self, args):
49         """
50         Drop any empty strings from the list of compile and link flags
51         """
52         return [x for x in args if len(x) > 0]
53 
54     def import_op(self):
55         """
56         This function will import the op module by its string name.
57         """
58         return importlib.import_module(self.prebuilt_import_path)
59 
60     def build_aot(self) -> "CppExtension":
61         from torch.utils.cpp_extension import CppExtension
62 
63         return CppExtension(
64             name=self.prebuilt_import_path,
65             sources=self.strip_empty_entries(self.sources_files()),
66             include_dirs=self.strip_empty_entries(self.include_dirs()),
67             extra_compile_args=self.strip_empty_entries(self.cxx_flags()),
68         )
69 
70     def build_jit(self) -> None:
71         from torch.utils.cpp_extension import load
72 
73         build_directory = _Extension.get_jit_extension_folder_path()
74         build_directory = Path(build_directory)
75         build_directory.mkdir(parents=True, exist_ok=True)
76 
77         # check if the kernel has been built
78         compiled_before = False
79         kernel_file_path = build_directory.joinpath(f"{self.name}.o")
80         if kernel_file_path.exists():
81             compiled_before = True
82 
83         # load the kernel
84         if compiled_before:
85             print(f"[extension] Loading the JIT-built {self.name} kernel during runtime now")
86         else:
87             print(f"[extension] Compiling the JIT {self.name} kernel during runtime now")
88 
89         build_start = time.time()
90         op_kernel = load(
91             name=self.name,
92             sources=self.strip_empty_entries(self.sources_files()),
93             extra_include_paths=self.strip_empty_entries(self.include_dirs()),
94             extra_cflags=self.cxx_flags(),
95             extra_ldflags=[],
96             build_directory=str(build_directory),
97         )
98         build_duration = time.time() - build_start
99 
100         if compiled_before:
101             print(f"[extension] Time taken to load {self.name} op: {build_duration} seconds")
102         else:
103             print(f"[extension] Time taken to compile {self.name} op: {build_duration} seconds")
104 
105         return op_kernel
106 
107     # functions must be overrided begin
108     @abstractmethod
109     def sources_files(self) -> List[str]:
110         """
111         This function should return a list of source files for extensions.
112         """
113 
114     @abstractmethod
115     def include_dirs(self) -> List[str]:
116         """
117         This function should return a list of include files for extensions.
118         """
119 
120     @abstractmethod
121     def cxx_flags(self) -> List[str]:
122         """
123         This function should return a list of cxx compilation flags for extensions.
124         """
125 
126     def load(self):
127         try:
128             op_kernel = self.import_op()
129         except ImportError:
130             # if import error occurs, it means that the kernel is not pre-built
131             # so we build it jit
132             op_kernel = self.build_jit()
133 
134         return op_kernel
135
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/extensions/cpp_extension.py b/extensions/cpp_extension.py
--- a/extensions/cpp_extension.py
+++ b/extensions/cpp_extension.py
@@ -126,7 +126,7 @@
     def load(self):
         try:
             op_kernel = self.import_op()
-        except ImportError:
+        except (ImportError, ModuleNotFoundError):
             # if import error occurs, it means that the kernel is not pre-built
             # so we build it jit
             op_kernel = self.build_jit()

{"golden_diff": "diff --git a/extensions/cpp_extension.py b/extensions/cpp_extension.py\n--- a/extensions/cpp_extension.py\n+++ b/extensions/cpp_extension.py\n@@ -126,7 +126,7 @@\n     def load(self):\n         try:\n             op_kernel = self.import_op()\n-        except ImportError:\n+        except (ImportError, ModuleNotFoundError):\n             # if import error occurs, it means that the kernel is not pre-built\n             # so we build it jit\n             op_kernel = self.build_jit()\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import importlib\nimport os\nimport time\nfrom abc import abstractmethod\nfrom pathlib import Path\nfrom typing import List\n\nfrom .base_extension import _Extension\n\n__all__ = [\"_CppExtension\"]\n\n\nclass _CppExtension(_Extension):\n    def __init__(self, name: str, priority: int = 1):\n        super().__init__(name, support_aot=True, support_jit=True, priority=priority)\n\n        # we store the op as an attribute to avoid repeated building and loading\n        self.cached_op = None\n\n        # build-related variables\n        self.prebuilt_module_path = \"colossalai._C\"\n        self.prebuilt_import_path = f\"{self.prebuilt_module_path}.{self.name}\"\n        self.version_dependent_macros = [\"-DVERSION_GE_1_1\", \"-DVERSION_GE_1_3\", \"-DVERSION_GE_1_5\"]\n\n    def csrc_abs_path(self, path):\n        return os.path.join(self.relative_to_abs_path(\"csrc\"), path)\n\n    def relative_to_abs_path(self, code_path: str) -> str:\n        \"\"\"\n        This function takes in a path relative to the colossalai root directory and return the absolute path.\n        \"\"\"\n\n        # get the current file path\n        # iteratively check the parent directory\n        # if the parent directory is \"extensions\", then the current file path is the root directory\n        # otherwise, the current file path is inside the root directory\n        current_file_path = Path(__file__)\n        while True:\n            if current_file_path.name == \"extensions\":\n                break\n            else:\n                current_file_path = current_file_path.parent\n        extension_module_path = current_file_path\n        code_abs_path = extension_module_path.joinpath(code_path)\n        return str(code_abs_path)\n\n    # functions must be overrided over\n    def strip_empty_entries(self, args):\n        \"\"\"\n        Drop any empty strings from the list of compile and link flags\n        \"\"\"\n        return [x for x in args if len(x) > 0]\n\n    def import_op(self):\n        \"\"\"\n        This function will import the op module by its string name.\n        \"\"\"\n        return importlib.import_module(self.prebuilt_import_path)\n\n    def build_aot(self) -> \"CppExtension\":\n        from torch.utils.cpp_extension import CppExtension\n\n        return CppExtension(\n            name=self.prebuilt_import_path,\n            sources=self.strip_empty_entries(self.sources_files()),\n            include_dirs=self.strip_empty_entries(self.include_dirs()),\n            extra_compile_args=self.strip_empty_entries(self.cxx_flags()),\n        )\n\n    def build_jit(self) -> None:\n        from torch.utils.cpp_extension import load\n\n        build_directory = _Extension.get_jit_extension_folder_path()\n        build_directory = Path(build_directory)\n        build_directory.mkdir(parents=True, exist_ok=True)\n\n        # check if the kernel has been built\n        compiled_before = False\n        kernel_file_path = build_directory.joinpath(f\"{self.name}.o\")\n        if kernel_file_path.exists():\n            compiled_before = True\n\n        # load the kernel\n        if compiled_before:\n            print(f\"[extension] Loading the JIT-built {self.name} kernel during runtime now\")\n        else:\n            print(f\"[extension] Compiling the JIT {self.name} kernel during runtime now\")\n\n        build_start = time.time()\n        op_kernel = load(\n            name=self.name,\n            sources=self.strip_empty_entries(self.sources_files()),\n            extra_include_paths=self.strip_empty_entries(self.include_dirs()),\n            extra_cflags=self.cxx_flags(),\n            extra_ldflags=[],\n            build_directory=str(build_directory),\n        )\n        build_duration = time.time() - build_start\n\n        if compiled_before:\n            print(f\"[extension] Time taken to load {self.name} op: {build_duration} seconds\")\n        else:\n            print(f\"[extension] Time taken to compile {self.name} op: {build_duration} seconds\")\n\n        return op_kernel\n\n    # functions must be overrided begin\n    @abstractmethod\n    def sources_files(self) -> List[str]:\n        \"\"\"\n        This function should return a list of source files for extensions.\n        \"\"\"\n\n    @abstractmethod\n    def include_dirs(self) -> List[str]:\n        \"\"\"\n        This function should return a list of include files for extensions.\n        \"\"\"\n\n    @abstractmethod\n    def cxx_flags(self) -> List[str]:\n        \"\"\"\n        This function should return a list of cxx compilation flags for extensions.\n        \"\"\"\n\n    def load(self):\n        try:\n            op_kernel = self.import_op()\n        except ImportError:\n            # if import error occurs, it means that the kernel is not pre-built\n            # so we build it jit\n            op_kernel = self.build_jit()\n\n        return op_kernel\n", "path": "extensions/cpp_extension.py"}], "after_files": [{"content": "import importlib\nimport os\nimport time\nfrom abc import abstractmethod\nfrom pathlib import Path\nfrom typing import List\n\nfrom .base_extension import _Extension\n\n__all__ = [\"_CppExtension\"]\n\n\nclass _CppExtension(_Extension):\n    def __init__(self, name: str, priority: int = 1):\n        super().__init__(name, support_aot=True, support_jit=True, priority=priority)\n\n        # we store the op as an attribute to avoid repeated building and loading\n        self.cached_op = None\n\n        # build-related variables\n        self.prebuilt_module_path = \"colossalai._C\"\n        self.prebuilt_import_path = f\"{self.prebuilt_module_path}.{self.name}\"\n        self.version_dependent_macros = [\"-DVERSION_GE_1_1\", \"-DVERSION_GE_1_3\", \"-DVERSION_GE_1_5\"]\n\n    def csrc_abs_path(self, path):\n        return os.path.join(self.relative_to_abs_path(\"csrc\"), path)\n\n    def relative_to_abs_path(self, code_path: str) -> str:\n        \"\"\"\n        This function takes in a path relative to the colossalai root directory and return the absolute path.\n        \"\"\"\n\n        # get the current file path\n        # iteratively check the parent directory\n        # if the parent directory is \"extensions\", then the current file path is the root directory\n        # otherwise, the current file path is inside the root directory\n        current_file_path = Path(__file__)\n        while True:\n            if current_file_path.name == \"extensions\":\n                break\n            else:\n                current_file_path = current_file_path.parent\n        extension_module_path = current_file_path\n        code_abs_path = extension_module_path.joinpath(code_path)\n        return str(code_abs_path)\n\n    # functions must be overrided over\n    def strip_empty_entries(self, args):\n        \"\"\"\n        Drop any empty strings from the list of compile and link flags\n        \"\"\"\n        return [x for x in args if len(x) > 0]\n\n    def import_op(self):\n        \"\"\"\n        This function will import the op module by its string name.\n        \"\"\"\n        return importlib.import_module(self.prebuilt_import_path)\n\n    def build_aot(self) -> \"CppExtension\":\n        from torch.utils.cpp_extension import CppExtension\n\n        return CppExtension(\n            name=self.prebuilt_import_path,\n            sources=self.strip_empty_entries(self.sources_files()),\n            include_dirs=self.strip_empty_entries(self.include_dirs()),\n            extra_compile_args=self.strip_empty_entries(self.cxx_flags()),\n        )\n\n    def build_jit(self) -> None:\n        from torch.utils.cpp_extension import load\n\n        build_directory = _Extension.get_jit_extension_folder_path()\n        build_directory = Path(build_directory)\n        build_directory.mkdir(parents=True, exist_ok=True)\n\n        # check if the kernel has been built\n        compiled_before = False\n        kernel_file_path = build_directory.joinpath(f\"{self.name}.o\")\n        if kernel_file_path.exists():\n            compiled_before = True\n\n        # load the kernel\n        if compiled_before:\n            print(f\"[extension] Loading the JIT-built {self.name} kernel during runtime now\")\n        else:\n            print(f\"[extension] Compiling the JIT {self.name} kernel during runtime now\")\n\n        build_start = time.time()\n        op_kernel = load(\n            name=self.name,\n            sources=self.strip_empty_entries(self.sources_files()),\n            extra_include_paths=self.strip_empty_entries(self.include_dirs()),\n            extra_cflags=self.cxx_flags(),\n            extra_ldflags=[],\n            build_directory=str(build_directory),\n        )\n        build_duration = time.time() - build_start\n\n        if compiled_before:\n            print(f\"[extension] Time taken to load {self.name} op: {build_duration} seconds\")\n        else:\n            print(f\"[extension] Time taken to compile {self.name} op: {build_duration} seconds\")\n\n        return op_kernel\n\n    # functions must be overrided begin\n    @abstractmethod\n    def sources_files(self) -> List[str]:\n        \"\"\"\n        This function should return a list of source files for extensions.\n        \"\"\"\n\n    @abstractmethod\n    def include_dirs(self) -> List[str]:\n        \"\"\"\n        This function should return a list of include files for extensions.\n        \"\"\"\n\n    @abstractmethod\n    def cxx_flags(self) -> List[str]:\n        \"\"\"\n        This function should return a list of cxx compilation flags for extensions.\n        \"\"\"\n\n    def load(self):\n        try:\n            op_kernel = self.import_op()\n        except (ImportError, ModuleNotFoundError):\n            # if import error occurs, it means that the kernel is not pre-built\n            # so we build it jit\n            op_kernel = self.build_jit()\n\n        return op_kernel\n", "path": "extensions/cpp_extension.py"}]}