code: def prepare_wheel(sources_path: pathlib.Path, *, cpu, cuda_version):
      """Assembles a source tree for the cuda kernel wheel in `sources_path`."""
      copy_runfiles = functools.partial(
          build_utils.copy_file, runfiles=r
      )
  
      copy_runfiles(
          "__main__/jax_plugins/cuda/plugin_pyproject.toml",
          dst_dir=sources_path,
          dst_filename="pyproject.toml",
      )
      copy_runfiles(
          "__main__/jax_plugins/cuda/plugin_setup.py",
          dst_dir=sources_path,
          dst_filename="setup.py",
      )
      build_utils.update_setup_with_cuda_version(
          sources_path, cuda_version
      )
      write_setup_cfg(sources_path, cpu)
  
      plugin_dir = sources_path / f"jax_cuda{cuda_version}_plugin"
      copy_runfiles(
          dst_dir=plugin_dir / "nvvm" / "libdevice",
          src_files=[
              "local_config_cuda/cuda/cuda/nvvm/libdevice/libdevice.10.bc"
          ],
      )
      copy_runfiles(
          dst_dir=plugin_dir,
          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}",
              f"__main__/jaxlib/cuda_plugin_extension.{pyext}",
              "__main__/jaxlib/version.py",
          ],
      )
  
file_path: jax/jaxlib/tools/build_cuda_kernels_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_wheel(sources_path) without return types