code: def prepare_cuda_plugin_wheel(
      sources_path: pathlib.Path, *, cpu, cuda_version
  ):
      """Assembles a source tree for the wheel in `sources_path`."""
      copy_runfiles = functools.partial(
          build_utils.copy_file, runfiles=r
      )
  
      plugin_dir = (
          sources_path / "jax_plugins" / f"xla_cuda{cuda_version}"
      )
      copy_runfiles(
          dst_dir=sources_path,
          src_files=[
              "__main__/jax_plugins/cuda/pyproject.toml",
              "__main__/jax_plugins/cuda/setup.py",
          ],
      )
      build_utils.update_setup_with_cuda_version(
          sources_path, cuda_version
      )
      write_setup_cfg(sources_path, cpu)
      copy_runfiles(
          dst_dir=plugin_dir,
          src_files=[
              "__main__/jax_plugins/cuda/__init__.py",
              "__main__/jaxlib/version.py",
          ],
      )
      copy_runfiles(
          "xla/xla/pjrt/c/pjrt_c_api_gpu_plugin.so",
          dst_dir=plugin_dir,
          dst_filename="xla_cuda_plugin.so",
      )
  
file_path: jax/jaxlib/tools/build_gpu_plugin_wheel.py
imports: import argparse
  import functools
  import os
  import pathlib
  import tempfile
  from bazel_tools.tools.python.runfiles import runfiles
  from jax.tools import build_utils
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_cuda_plugin_wheel(sources_path) without return types