We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Error caused by shutil.rmtree 
```
Traceback (most recent call last):
  File "\\?\C:\Users\cloud\AppData\Local\Temp\Bazel.runfiles_vfpgffuf\runfiles\__main__\build\install_xla_in_source_tree.py", line 83, in <module>
    shutil.rmtree(jaxlib_dir)
  File "C:\Users\cloud\miniconda3\lib\shutil.py", line 516, in rmtree
    return _rmtree_unsafe(path, onerror)
  File "C:\Users\cloud\miniconda3\lib\shutil.py", line 400, in _rmtree_unsafe
    onerror(os.unlink, fullname, sys.exc_info())
  File "C:\Users\cloud\miniconda3\lib\shutil.py", line 398, in _rmtree_unsafe
    os.unlink(fullname)
WindowsError: [Error 5] Access is denied.: 'D:\\jax\\build\\jaxlib\\cublas_kernels.pyd'
```

This only happens on rebuild.

The reason is `shutil.rmtree` will not delete readonly file on Windows.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `build/build_wheel.py`
Content:
```
1 # Copyright 2020 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 # Script that builds a jaxlib wheel, intended to be run via bazel run as part
16 # of the jaxlib build process.
17 
18 # Most users should not run this script directly; use build.py instead.
19 
20 import argparse
21 import functools
22 import glob
23 import os
24 import platform
25 import shutil
26 import subprocess
27 import sys
28 import tempfile
29 
30 from bazel_tools.tools.python.runfiles import runfiles
31 
32 parser = argparse.ArgumentParser()
33 parser.add_argument(
34   "--sources_path",
35   default=None,
36   help="Path in which the wheel's sources should be prepared. Optional. If "
37        "omitted, a temporary directory will be used.")
38 parser.add_argument(
39   "--output_path",
40   default=None,
41   required=True,
42   help="Path to which the output wheel should be written. Required.")
43 args = parser.parse_args()
44 
45 r = runfiles.Create()
46 
47 
48 def _is_windows():
49   return sys.platform.startswith("win32")
50 
51 
52 def _copy_so(src_file, dst_dir, dst_filename=None):
53   src_filename = os.path.basename(src_file)
54   if not dst_filename:
55     if _is_windows() and src_filename.endswith(".so"):
56       dst_filename = src_filename[:-3] + ".pyd"
57     else:
58       dst_filename = src_filename
59   dst_file = os.path.join(dst_dir, dst_filename)
60   shutil.copy(src_file, dst_file)
61 
62 
63 def _copy_normal(src_file, dst_dir, dst_filename=None):
64   src_filename = os.path.basename(src_file)
65   dst_file = os.path.join(dst_dir, dst_filename or src_filename)
66   shutil.copy(src_file, dst_file)
67 
68 
69 def copy_file(src_file, dst_dir, dst_filename=None):
70   if src_file.endswith(".so"):
71     _copy_so(src_file, dst_dir, dst_filename=dst_filename)
72   else:
73     _copy_normal(src_file, dst_dir, dst_filename=dst_filename)
74 
75 def patch_copy_xla_client_py(dst_dir):
76   with open(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/xla_client.py")) as f:
77     src = f.read()
78     src = src.replace("from tensorflow.compiler.xla.python import xla_extension as _xla",
79                       "from . import xla_extension as _xla")
80     with open(os.path.join(dst_dir, "xla_client.py"), "w") as f:
81       f.write(src)
82 
83 
84 def patch_copy_tpu_client_py(dst_dir):
85   with open(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client.py")) as f:
86     src = f.read()
87     src = src.replace("from tensorflow.compiler.xla.python import xla_extension as _xla",
88                       "from . import xla_extension as _xla")
89     src = src.replace("from tensorflow.compiler.xla.python import xla_client",
90                       "from . import xla_client")
91     src = src.replace(
92         "from tensorflow.compiler.xla.python.tpu_driver.client import tpu_client_extension as _tpu_client",
93         "from . import tpu_client_extension as _tpu_client")
94     with open(os.path.join(dst_dir, "tpu_client.py"), "w") as f:
95       f.write(src)
96 
97 
98 def prepare_wheel(sources_path):
99   """Assembles a source tree for the wheel in `sources_path`."""
100   jaxlib_dir = os.path.join(sources_path, "jaxlib")
101   os.makedirs(jaxlib_dir)
102   copy_to_jaxlib = functools.partial(copy_file, dst_dir=jaxlib_dir)
103 
104   copy_file(r.Rlocation("__main__/jaxlib/setup.py"), dst_dir=sources_path)
105   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/init.py"), dst_filename="__init__.py")
106   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/lapack.so"))
107   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/_pocketfft.so"))
108   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/pocketfft_flatbuffers_py_generated.py"))
109   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/pocketfft.py"))
110   if r.Rlocation("__main__/jaxlib/cusolver_kernels.so") is not None:
111     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cusolver_kernels.so"))
112     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cublas_kernels.so"))
113     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cusolver_kernels.so"))
114     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cuda_prng_kernels.so"))
115   if r.Rlocation("__main__/jaxlib/cusolver_kernels.pyd") is not None:
116     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cusolver_kernels.pyd"))
117     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cublas_kernels.pyd"))
118     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cusolver_kernels.pyd"))
119     copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cuda_prng_kernels.pyd"))
120   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/version.py"))
121   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cusolver.py"))
122   copy_to_jaxlib(r.Rlocation("__main__/jaxlib/cuda_prng.py"))
123 
124   if _is_windows():
125     copy_to_jaxlib(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/xla_extension.pyd"))
126   else:
127     copy_to_jaxlib(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/xla_extension.so"))
128   patch_copy_xla_client_py(jaxlib_dir)
129 
130   if not _is_windows():
131     copy_to_jaxlib(r.Rlocation("org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client_extension.so"))
132     patch_copy_tpu_client_py(jaxlib_dir)
133 
134 
135 def build_wheel(sources_path, output_path):
136   """Builds a wheel in `output_path` using the source tree in `sources_path`."""
137   platform_name = {
138     "Linux": "manylinux2010",
139     "Darwin": "macosx_10_9",
140     "Windows": "win",
141   }[platform.system()]
142   cpu_name = "amd64" if platform.system() == "Windows" else "x86_64"
143   python_tag_arg = (f"--python-tag=cp{sys.version_info.major}"
144                     f"{sys.version_info.minor}")
145   platform_tag_arg = f"--plat-name={platform_name}_{cpu_name}"
146   cwd = os.getcwd()
147   os.chdir(sources_path)
148   subprocess.run([sys.executable, "setup.py", "bdist_wheel",
149                  python_tag_arg, platform_tag_arg])
150   os.chdir(cwd)
151   for wheel in glob.glob(os.path.join(sources_path, "dist", "*.whl")):
152     output_file = os.path.join(output_path, os.path.basename(wheel))
153     sys.stderr.write(f"Output wheel: {output_file}\n\n")
154     sys.stderr.write(f"To install the newly-built jaxlib wheel, run:\n")
155     sys.stderr.write(f"  pip install {output_file}\n\n")
156     shutil.copy(wheel, output_path)
157 
158 
159 tmpdir = None
160 sources_path = args.sources_path
161 if sources_path is None:
162   tmpdir = tempfile.TemporaryDirectory(prefix="jaxlib")
163   sources_path = tmpdir.name
164 
165 try:
166   os.makedirs(args.output_path, exist_ok=True)
167   prepare_wheel(sources_path)
168   build_wheel(sources_path, args.output_path)
169 finally:
170   if tmpdir:
171     tmpdir.cleanup()
172 
173
```
--- 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/build/build_wheel.py b/build/build_wheel.py
--- a/build/build_wheel.py
+++ b/build/build_wheel.py
@@ -57,13 +57,19 @@
     else:
       dst_filename = src_filename
   dst_file = os.path.join(dst_dir, dst_filename)
-  shutil.copy(src_file, dst_file)
+  if _is_windows():
+    shutil.copyfile(src_file, dst_file)
+  else:
+    shutil.copy(src_file, dst_file)
 
 
 def _copy_normal(src_file, dst_dir, dst_filename=None):
   src_filename = os.path.basename(src_file)
   dst_file = os.path.join(dst_dir, dst_filename or src_filename)
-  shutil.copy(src_file, dst_file)
+  if _is_windows():
+    shutil.copyfile(src_file, dst_file)
+  else:
+    shutil.copy(src_file, dst_file)
 
 
 def copy_file(src_file, dst_dir, dst_filename=None):
@@ -169,4 +175,3 @@
 finally:
   if tmpdir:
     tmpdir.cleanup()
-

{"golden_diff": "diff --git a/build/build_wheel.py b/build/build_wheel.py\n--- a/build/build_wheel.py\n+++ b/build/build_wheel.py\n@@ -57,13 +57,19 @@\n     else:\n       dst_filename = src_filename\n   dst_file = os.path.join(dst_dir, dst_filename)\n-  shutil.copy(src_file, dst_file)\n+  if _is_windows():\n+    shutil.copyfile(src_file, dst_file)\n+  else:\n+    shutil.copy(src_file, dst_file)\n \n \n def _copy_normal(src_file, dst_dir, dst_filename=None):\n   src_filename = os.path.basename(src_file)\n   dst_file = os.path.join(dst_dir, dst_filename or src_filename)\n-  shutil.copy(src_file, dst_file)\n+  if _is_windows():\n+    shutil.copyfile(src_file, dst_file)\n+  else:\n+    shutil.copy(src_file, dst_file)\n \n \n def copy_file(src_file, dst_dir, dst_filename=None):\n@@ -169,4 +175,3 @@\n finally:\n   if tmpdir:\n     tmpdir.cleanup()\n-\n", "issue": "Error caused by shutil.rmtree \n```\r\nTraceback (most recent call last):\r\n  File \"\\\\?\\C:\\Users\\cloud\\AppData\\Local\\Temp\\Bazel.runfiles_vfpgffuf\\runfiles\\__main__\\build\\install_xla_in_source_tree.py\", line 83, in <module>\r\n    shutil.rmtree(jaxlib_dir)\r\n  File \"C:\\Users\\cloud\\miniconda3\\lib\\shutil.py\", line 516, in rmtree\r\n    return _rmtree_unsafe(path, onerror)\r\n  File \"C:\\Users\\cloud\\miniconda3\\lib\\shutil.py\", line 400, in _rmtree_unsafe\r\n    onerror(os.unlink, fullname, sys.exc_info())\r\n  File \"C:\\Users\\cloud\\miniconda3\\lib\\shutil.py\", line 398, in _rmtree_unsafe\r\n    os.unlink(fullname)\r\nWindowsError: [Error 5] Access is denied.: 'D:\\\\jax\\\\build\\\\jaxlib\\\\cublas_kernels.pyd'\r\n```\r\n\r\nThis only happens on rebuild.\r\n\r\nThe reason is `shutil.rmtree` will not delete readonly file on Windows.\n", "before_files": [{"content": "# Copyright 2020 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Script that builds a jaxlib wheel, intended to be run via bazel run as part\n# of the jaxlib build process.\n\n# Most users should not run this script directly; use build.py instead.\n\nimport argparse\nimport functools\nimport glob\nimport os\nimport platform\nimport shutil\nimport subprocess\nimport sys\nimport tempfile\n\nfrom bazel_tools.tools.python.runfiles import runfiles\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\n  \"--sources_path\",\n  default=None,\n  help=\"Path in which the wheel's sources should be prepared. Optional. If \"\n       \"omitted, a temporary directory will be used.\")\nparser.add_argument(\n  \"--output_path\",\n  default=None,\n  required=True,\n  help=\"Path to which the output wheel should be written. Required.\")\nargs = parser.parse_args()\n\nr = runfiles.Create()\n\n\ndef _is_windows():\n  return sys.platform.startswith(\"win32\")\n\n\ndef _copy_so(src_file, dst_dir, dst_filename=None):\n  src_filename = os.path.basename(src_file)\n  if not dst_filename:\n    if _is_windows() and src_filename.endswith(\".so\"):\n      dst_filename = src_filename[:-3] + \".pyd\"\n    else:\n      dst_filename = src_filename\n  dst_file = os.path.join(dst_dir, dst_filename)\n  shutil.copy(src_file, dst_file)\n\n\ndef _copy_normal(src_file, dst_dir, dst_filename=None):\n  src_filename = os.path.basename(src_file)\n  dst_file = os.path.join(dst_dir, dst_filename or src_filename)\n  shutil.copy(src_file, dst_file)\n\n\ndef copy_file(src_file, dst_dir, dst_filename=None):\n  if src_file.endswith(\".so\"):\n    _copy_so(src_file, dst_dir, dst_filename=dst_filename)\n  else:\n    _copy_normal(src_file, dst_dir, dst_filename=dst_filename)\n\ndef patch_copy_xla_client_py(dst_dir):\n  with open(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/xla_client.py\")) as f:\n    src = f.read()\n    src = src.replace(\"from tensorflow.compiler.xla.python import xla_extension as _xla\",\n                      \"from . import xla_extension as _xla\")\n    with open(os.path.join(dst_dir, \"xla_client.py\"), \"w\") as f:\n      f.write(src)\n\n\ndef patch_copy_tpu_client_py(dst_dir):\n  with open(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client.py\")) as f:\n    src = f.read()\n    src = src.replace(\"from tensorflow.compiler.xla.python import xla_extension as _xla\",\n                      \"from . import xla_extension as _xla\")\n    src = src.replace(\"from tensorflow.compiler.xla.python import xla_client\",\n                      \"from . import xla_client\")\n    src = src.replace(\n        \"from tensorflow.compiler.xla.python.tpu_driver.client import tpu_client_extension as _tpu_client\",\n        \"from . import tpu_client_extension as _tpu_client\")\n    with open(os.path.join(dst_dir, \"tpu_client.py\"), \"w\") as f:\n      f.write(src)\n\n\ndef prepare_wheel(sources_path):\n  \"\"\"Assembles a source tree for the wheel in `sources_path`.\"\"\"\n  jaxlib_dir = os.path.join(sources_path, \"jaxlib\")\n  os.makedirs(jaxlib_dir)\n  copy_to_jaxlib = functools.partial(copy_file, dst_dir=jaxlib_dir)\n\n  copy_file(r.Rlocation(\"__main__/jaxlib/setup.py\"), dst_dir=sources_path)\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/init.py\"), dst_filename=\"__init__.py\")\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/lapack.so\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/_pocketfft.so\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/pocketfft_flatbuffers_py_generated.py\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/pocketfft.py\"))\n  if r.Rlocation(\"__main__/jaxlib/cusolver_kernels.so\") is not None:\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.so\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cublas_kernels.so\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.so\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cuda_prng_kernels.so\"))\n  if r.Rlocation(\"__main__/jaxlib/cusolver_kernels.pyd\") is not None:\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.pyd\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cublas_kernels.pyd\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.pyd\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cuda_prng_kernels.pyd\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/version.py\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver.py\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cuda_prng.py\"))\n\n  if _is_windows():\n    copy_to_jaxlib(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/xla_extension.pyd\"))\n  else:\n    copy_to_jaxlib(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/xla_extension.so\"))\n  patch_copy_xla_client_py(jaxlib_dir)\n\n  if not _is_windows():\n    copy_to_jaxlib(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client_extension.so\"))\n    patch_copy_tpu_client_py(jaxlib_dir)\n\n\ndef build_wheel(sources_path, output_path):\n  \"\"\"Builds a wheel in `output_path` using the source tree in `sources_path`.\"\"\"\n  platform_name = {\n    \"Linux\": \"manylinux2010\",\n    \"Darwin\": \"macosx_10_9\",\n    \"Windows\": \"win\",\n  }[platform.system()]\n  cpu_name = \"amd64\" if platform.system() == \"Windows\" else \"x86_64\"\n  python_tag_arg = (f\"--python-tag=cp{sys.version_info.major}\"\n                    f\"{sys.version_info.minor}\")\n  platform_tag_arg = f\"--plat-name={platform_name}_{cpu_name}\"\n  cwd = os.getcwd()\n  os.chdir(sources_path)\n  subprocess.run([sys.executable, \"setup.py\", \"bdist_wheel\",\n                 python_tag_arg, platform_tag_arg])\n  os.chdir(cwd)\n  for wheel in glob.glob(os.path.join(sources_path, \"dist\", \"*.whl\")):\n    output_file = os.path.join(output_path, os.path.basename(wheel))\n    sys.stderr.write(f\"Output wheel: {output_file}\\n\\n\")\n    sys.stderr.write(f\"To install the newly-built jaxlib wheel, run:\\n\")\n    sys.stderr.write(f\"  pip install {output_file}\\n\\n\")\n    shutil.copy(wheel, output_path)\n\n\ntmpdir = None\nsources_path = args.sources_path\nif sources_path is None:\n  tmpdir = tempfile.TemporaryDirectory(prefix=\"jaxlib\")\n  sources_path = tmpdir.name\n\ntry:\n  os.makedirs(args.output_path, exist_ok=True)\n  prepare_wheel(sources_path)\n  build_wheel(sources_path, args.output_path)\nfinally:\n  if tmpdir:\n    tmpdir.cleanup()\n\n", "path": "build/build_wheel.py"}], "after_files": [{"content": "# Copyright 2020 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Script that builds a jaxlib wheel, intended to be run via bazel run as part\n# of the jaxlib build process.\n\n# Most users should not run this script directly; use build.py instead.\n\nimport argparse\nimport functools\nimport glob\nimport os\nimport platform\nimport shutil\nimport subprocess\nimport sys\nimport tempfile\n\nfrom bazel_tools.tools.python.runfiles import runfiles\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\n  \"--sources_path\",\n  default=None,\n  help=\"Path in which the wheel's sources should be prepared. Optional. If \"\n       \"omitted, a temporary directory will be used.\")\nparser.add_argument(\n  \"--output_path\",\n  default=None,\n  required=True,\n  help=\"Path to which the output wheel should be written. Required.\")\nargs = parser.parse_args()\n\nr = runfiles.Create()\n\n\ndef _is_windows():\n  return sys.platform.startswith(\"win32\")\n\n\ndef _copy_so(src_file, dst_dir, dst_filename=None):\n  src_filename = os.path.basename(src_file)\n  if not dst_filename:\n    if _is_windows() and src_filename.endswith(\".so\"):\n      dst_filename = src_filename[:-3] + \".pyd\"\n    else:\n      dst_filename = src_filename\n  dst_file = os.path.join(dst_dir, dst_filename)\n  if _is_windows():\n    shutil.copyfile(src_file, dst_file)\n  else:\n    shutil.copy(src_file, dst_file)\n\n\ndef _copy_normal(src_file, dst_dir, dst_filename=None):\n  src_filename = os.path.basename(src_file)\n  dst_file = os.path.join(dst_dir, dst_filename or src_filename)\n  if _is_windows():\n    shutil.copyfile(src_file, dst_file)\n  else:\n    shutil.copy(src_file, dst_file)\n\n\ndef copy_file(src_file, dst_dir, dst_filename=None):\n  if src_file.endswith(\".so\"):\n    _copy_so(src_file, dst_dir, dst_filename=dst_filename)\n  else:\n    _copy_normal(src_file, dst_dir, dst_filename=dst_filename)\n\ndef patch_copy_xla_client_py(dst_dir):\n  with open(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/xla_client.py\")) as f:\n    src = f.read()\n    src = src.replace(\"from tensorflow.compiler.xla.python import xla_extension as _xla\",\n                      \"from . import xla_extension as _xla\")\n    with open(os.path.join(dst_dir, \"xla_client.py\"), \"w\") as f:\n      f.write(src)\n\n\ndef patch_copy_tpu_client_py(dst_dir):\n  with open(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client.py\")) as f:\n    src = f.read()\n    src = src.replace(\"from tensorflow.compiler.xla.python import xla_extension as _xla\",\n                      \"from . import xla_extension as _xla\")\n    src = src.replace(\"from tensorflow.compiler.xla.python import xla_client\",\n                      \"from . import xla_client\")\n    src = src.replace(\n        \"from tensorflow.compiler.xla.python.tpu_driver.client import tpu_client_extension as _tpu_client\",\n        \"from . import tpu_client_extension as _tpu_client\")\n    with open(os.path.join(dst_dir, \"tpu_client.py\"), \"w\") as f:\n      f.write(src)\n\n\ndef prepare_wheel(sources_path):\n  \"\"\"Assembles a source tree for the wheel in `sources_path`.\"\"\"\n  jaxlib_dir = os.path.join(sources_path, \"jaxlib\")\n  os.makedirs(jaxlib_dir)\n  copy_to_jaxlib = functools.partial(copy_file, dst_dir=jaxlib_dir)\n\n  copy_file(r.Rlocation(\"__main__/jaxlib/setup.py\"), dst_dir=sources_path)\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/init.py\"), dst_filename=\"__init__.py\")\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/lapack.so\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/_pocketfft.so\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/pocketfft_flatbuffers_py_generated.py\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/pocketfft.py\"))\n  if r.Rlocation(\"__main__/jaxlib/cusolver_kernels.so\") is not None:\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.so\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cublas_kernels.so\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.so\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cuda_prng_kernels.so\"))\n  if r.Rlocation(\"__main__/jaxlib/cusolver_kernels.pyd\") is not None:\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.pyd\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cublas_kernels.pyd\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver_kernels.pyd\"))\n    copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cuda_prng_kernels.pyd\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/version.py\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cusolver.py\"))\n  copy_to_jaxlib(r.Rlocation(\"__main__/jaxlib/cuda_prng.py\"))\n\n  if _is_windows():\n    copy_to_jaxlib(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/xla_extension.pyd\"))\n  else:\n    copy_to_jaxlib(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/xla_extension.so\"))\n  patch_copy_xla_client_py(jaxlib_dir)\n\n  if not _is_windows():\n    copy_to_jaxlib(r.Rlocation(\"org_tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/tpu_client_extension.so\"))\n    patch_copy_tpu_client_py(jaxlib_dir)\n\n\ndef build_wheel(sources_path, output_path):\n  \"\"\"Builds a wheel in `output_path` using the source tree in `sources_path`.\"\"\"\n  platform_name = {\n    \"Linux\": \"manylinux2010\",\n    \"Darwin\": \"macosx_10_9\",\n    \"Windows\": \"win\",\n  }[platform.system()]\n  cpu_name = \"amd64\" if platform.system() == \"Windows\" else \"x86_64\"\n  python_tag_arg = (f\"--python-tag=cp{sys.version_info.major}\"\n                    f\"{sys.version_info.minor}\")\n  platform_tag_arg = f\"--plat-name={platform_name}_{cpu_name}\"\n  cwd = os.getcwd()\n  os.chdir(sources_path)\n  subprocess.run([sys.executable, \"setup.py\", \"bdist_wheel\",\n                 python_tag_arg, platform_tag_arg])\n  os.chdir(cwd)\n  for wheel in glob.glob(os.path.join(sources_path, \"dist\", \"*.whl\")):\n    output_file = os.path.join(output_path, os.path.basename(wheel))\n    sys.stderr.write(f\"Output wheel: {output_file}\\n\\n\")\n    sys.stderr.write(f\"To install the newly-built jaxlib wheel, run:\\n\")\n    sys.stderr.write(f\"  pip install {output_file}\\n\\n\")\n    shutil.copy(wheel, output_path)\n\n\ntmpdir = None\nsources_path = args.sources_path\nif sources_path is None:\n  tmpdir = tempfile.TemporaryDirectory(prefix=\"jaxlib\")\n  sources_path = tmpdir.name\n\ntry:\n  os.makedirs(args.output_path, exist_ok=True)\n  prepare_wheel(sources_path)\n  build_wheel(sources_path, args.output_path)\nfinally:\n  if tmpdir:\n    tmpdir.cleanup()\n", "path": "build/build_wheel.py"}]}