code: def prepare_wheel(
      sources_path: pathlib.Path, *, cpu, include_gpu_plugin_extension
  ):
      """Assembles a source tree for the wheel in `sources_path`."""
      copy_runfiles = functools.partial(
          build_utils.copy_file, runfiles=r
      )
  
      verify_mac_libraries_dont_reference_chkstack()
      copy_runfiles(
          dst_dir=sources_path,
          src_files=[
              "__main__/jaxlib/tools/LICENSE.txt",
              "__main__/jaxlib/README.md",
              "__main__/jaxlib/setup.py",
          ],
      )
      write_setup_cfg(sources_path, cpu)
  
      jaxlib_dir = sources_path / "jaxlib"
      copy_runfiles(
          "__main__/jaxlib/init.py",
          dst_dir=jaxlib_dir,
          dst_filename="__init__.py",
      )
      copy_runfiles(
          dst_dir=jaxlib_dir,
          src_files=[
              f"__main__/jaxlib/cpu_feature_guard.{pyext}",
              f"__main__/jaxlib/utils.{pyext}",
              "__main__/jaxlib/lapack.py",
              "__main__/jaxlib/hlo_helpers.py",
              "__main__/jaxlib/ducc_fft.py",
              "__main__/jaxlib/gpu_prng.py",
              "__main__/jaxlib/gpu_linalg.py",
              "__main__/jaxlib/gpu_rnn.py",
              "__main__/jaxlib/gpu_triton.py",
              "__main__/jaxlib/gpu_common_utils.py",
              "__main__/jaxlib/gpu_solver.py",
              "__main__/jaxlib/gpu_sparse.py",
              "__main__/jaxlib/version.py",
              "__main__/jaxlib/xla_client.py",
              f"__main__/jaxlib/xla_extension.{pyext}",
          ],
      )
      # This file is required by PEP-561. It marks jaxlib as package containing
      # type stubs.
      with open(jaxlib_dir / "py.typed", "w"):
          pass
      patch_copy_xla_extension_stubs(jaxlib_dir)
  
      copy_runfiles(
          dst_dir=jaxlib_dir / "cpu",
          src_files=[
              f"__main__/jaxlib/cpu/_lapack.{pyext}",
              f"__main__/jaxlib/cpu/_ducc_fft.{pyext}",
          ],
      )
  
      if (
          exists(f"__main__/jaxlib/cuda/_solver.{pyext}")
          and not include_gpu_plugin_extension
      ):
          copy_runfiles(
              dst_dir=jaxlib_dir / "cuda" / "nvvm" / "libdevice",
              src_files=[
                  "local_config_cuda/cuda/cuda/nvvm/libdevice/libdevice.10.bc"
              ],
          )
          copy_runfiles(
              dst_dir=jaxlib_dir / "cuda",
              src_files=[
                  f"__main__/jaxlib/cuda/_solver.{pyext}",
                  f"__main__/jaxlib/cuda/_blas.{pyext}",
                  f"__main__/jaxlib/cuda/_linalg.{pyext}",
                  f"__main__/jaxlib/cuda/_prng.{pyext}",
                  f"__main__/jaxlib/cuda/_rnn.{pyext}",
                  f"__main__/jaxlib/cuda/_sparse.{pyext}",
                  f"__main__/jaxlib/cuda/_triton.{pyext}",
                  f"__main__/jaxlib/cuda/_versions.{pyext}",
              ],
          )
      if exists(f"__main__/jaxlib/rocm/_solver.{pyext}"):
          copy_runfiles(
              dst_dir=jaxlib_dir / "rocm",
              src_files=[
                  f"__main__/jaxlib/rocm/_solver.{pyext}",
                  f"__main__/jaxlib/rocm/_blas.{pyext}",
                  f"__main__/jaxlib/rocm/_linalg.{pyext}",
                  f"__main__/jaxlib/rocm/_prng.{pyext}",
                  f"__main__/jaxlib/rocm/_sparse.{pyext}",
                  f"__main__/jaxlib/rocm/_triton.{pyext}",
              ],
          )
  
      mosaic_python_dir = jaxlib_dir / "mosaic" / "python"
      copy_runfiles(
          dst_dir=mosaic_python_dir,
          src_files=[
              "__main__/jaxlib/mosaic/python/layout_defs.py",
              "__main__/jaxlib/mosaic/python/tpu.py",
          ],
      )
      # TODO (sharadmv,skyewm): can we avoid patching this file?
      patch_copy_mlir_import(
          "__main__/jaxlib/mosaic/python/_tpu_gen.py",
          dst_dir=mosaic_python_dir,
      )
  
      copy_runfiles(
          dst_dir=jaxlib_dir / "mlir",
          src_files=[
              "__main__/jaxlib/mlir/ir.py",
              "__main__/jaxlib/mlir/passmanager.py",
          ],
      )
      copy_runfiles(
          dst_dir=jaxlib_dir / "mlir" / "dialects",
          src_files=[
              "__main__/jaxlib/mlir/dialects/_arith_enum_gen.py",
              "__main__/jaxlib/mlir/dialects/_arith_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_builtin_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_chlo_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_func_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_math_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_memref_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_mhlo_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_ods_common.py",
              "__main__/jaxlib/mlir/dialects/_scf_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_sparse_tensor_enum_gen.py",
              "__main__/jaxlib/mlir/dialects/_sparse_tensor_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_stablehlo_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/_vector_enum_gen.py",
              "__main__/jaxlib/mlir/dialects/_vector_ops_gen.py",
              "__main__/jaxlib/mlir/dialects/arith.py",
              "__main__/jaxlib/mlir/dialects/builtin.py",
              "__main__/jaxlib/mlir/dialects/chlo.py",
              "__main__/jaxlib/mlir/dialects/func.py",
              "__main__/jaxlib/mlir/dialects/math.py",
              "__main__/jaxlib/mlir/dialects/memref.py",
              "__main__/jaxlib/mlir/dialects/mhlo.py",
              "__main__/jaxlib/mlir/dialects/scf.py",
              "__main__/jaxlib/mlir/dialects/sparse_tensor.py",
              "__main__/jaxlib/mlir/dialects/stablehlo.py",
              "__main__/jaxlib/mlir/dialects/vector.py",
          ],
      )
      copy_runfiles(
          dst_dir=jaxlib_dir / "mlir" / "extras",
          src_files=[
              "__main__/jaxlib/mlir/extras/meta.py",
          ],
      )
  
      if build_utils.is_windows():
          capi_so = (
              "__main__/jaxlib/mlir/_mlir_libs/jaxlib_mlir_capi.dll"
          )
      else:
          so_ext = "dylib" if _is_mac() else "so"
          capi_so = f"__main__/jaxlib/mlir/_mlir_libs/libjaxlib_mlir_capi.{so_ext}"
  
      mlir_libs_dir = jaxlib_dir / "mlir" / "_mlir_libs"
      copy_runfiles(
          dst_dir=mlir_libs_dir,
          src_files=[
              capi_so,
              "__main__/jaxlib/mlir/_mlir_libs/__init__.py",
              f"__main__/jaxlib/mlir/_mlir_libs/_mlir.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/_chlo.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/_mlirHlo.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsSparseTensor.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/_mlirSparseTensorPasses.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/_tpu_ext.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/_stablehlo.{pyext}",
              f"__main__/jaxlib/mlir/_mlir_libs/register_jax_dialects.{pyext}",
          ],
      )
  
      triton_dir = jaxlib_dir / "triton"
      copy_runfiles(
          dst_dir=triton_dir,
          src_files=[
              "__main__/jaxlib/triton/__init__.py",
              "__main__/jaxlib/triton/compat.py",
              "__main__/jaxlib/triton/dialect.py",
              f"__main__/jaxlib/triton/_triton_ext.{pyext}",
              "__main__/jaxlib/triton/_triton_ext.pyi",
          ],
      )
      patch_copy_mlir_import(
          "__main__/jaxlib/triton/_triton_enum_gen.py",
          dst_dir=triton_dir,
      )
      patch_copy_mlir_import(
          "__main__/jaxlib/triton/_triton_ops_gen.py",
          dst_dir=triton_dir,
      )
  
file_path: jax/jaxlib/tools/build_wheel.py
imports: import argparse
  import functools
  import os
  import pathlib
  import platform
  import re
  import subprocess
  import tempfile
  from bazel_tools.tools.python.runfiles import runfiles
  from jax.tools import build_utils
label_desc: Use image node_id 7 for calling a global function with example usage: prepare_wheel(sources_path) without return types