code: def main():
      cwd = os.getcwd()
      parser = argparse.ArgumentParser(
          description="Builds jaxlib from source.", epilog=EPILOG
      )
      add_boolean_argument(
          parser,
          "verbose",
          default=False,
          help_str="Should we produce verbose debugging output?",
      )
      parser.add_argument(
          "--bazel_path",
          help="Path to the Bazel binary to use. The default is to find bazel via "
          "the PATH; if none is found, downloads a fresh copy of bazel from "
          "GitHub.",
      )
      parser.add_argument(
          "--python_bin_path",
          help="Path to Python binary to use. The default is the Python "
          "interpreter used to run the build script.",
      )
      parser.add_argument(
          "--target_cpu_features",
          choices=["release", "native", "default"],
          default="release",
          help="What CPU features should we target? 'release' enables CPU "
          "features that should be enabled for a release build, which on "
          "x86-64 architectures enables AVX. 'native' enables "
          "-march=native, which generates code targeted to use all "
          "features of the current machine. 'default' means don't opt-in "
          "to any architectural features and use whatever the C compiler "
          "generates by default.",
      )
      add_boolean_argument(
          parser,
          "enable_mkl_dnn",
          default=True,
          help_str="Should we build with MKL-DNN enabled?",
      )
      add_boolean_argument(
          parser,
          "enable_cuda",
          help_str="Should we build with CUDA enabled? Requires CUDA and CuDNN.",
      )
      add_boolean_argument(
          parser,
          "build_gpu_plugin",
          default=False,
          help_str=(
              "Are we building the gpu plugin in addition to jaxlib? The GPU "
              "plugin is still experimental and is not ready for use yet."
          ),
      )
      add_boolean_argument(
          parser,
          "build_cuda_kernel_plugin",
          default=False,
          help_str=(
              "Are we building the cuda kernel plugin? jaxlib will not be built "
              "when this flag is True."
          ),
      )
      add_boolean_argument(
          parser,
          "build_cuda_pjrt_plugin",
          default=False,
          help_str=(
              "Are we building the cuda pjrt plugin? jaxlib will not be built "
              "when this flag is True."
          ),
      )
      parser.add_argument(
          "--gpu_plugin_cuda_version",
          choices=["11", "12"],
          default="12",
          help="Which CUDA major version the gpu plugin is for.",
      )
      add_boolean_argument(
          parser,
          "enable_rocm",
          help_str="Should we build with ROCm enabled?",
      )
      add_boolean_argument(
          parser,
          "enable_nccl",
          default=True,
          help_str="Should we build with NCCL enabled? Has no effect for non-CUDA "
          "builds.",
      )
      add_boolean_argument(
          parser,
          "remote_build",
          default=False,
          help_str="Should we build with RBE (Remote Build Environment)?",
      )
      parser.add_argument(
          "--cuda_path", default=None, help="Path to the CUDA toolkit."
      )
      parser.add_argument(
          "--cudnn_path", default=None, help="Path to CUDNN libraries."
      )
      parser.add_argument(
          "--cuda_version",
          default=None,
          help="CUDA toolkit version, e.g., 11.1",
      )
      parser.add_argument(
          "--cudnn_version", default=None, help="CUDNN version, e.g., 8"
      )
      # Caution: if changing the default list of CUDA capabilities, you should also
      # update the list in .bazelrc, which is used for wheel builds.
      parser.add_argument(
          "--cuda_compute_capabilities",
          default=None,
          help="A comma-separated list of CUDA compute capabilities to support.",
      )
      parser.add_argument(
          "--rocm_amdgpu_targets",
          default="gfx900,gfx906,gfx908,gfx90a,gfx1030",
          help="A comma-separated list of ROCm amdgpu targets to support.",
      )
      parser.add_argument(
          "--rocm_path", default=None, help="Path to the ROCm toolkit."
      )
      parser.add_argument(
          "--bazel_startup_options",
          action="append",
          default=[],
          help="Additional startup options to pass to bazel.",
      )
      parser.add_argument(
          "--bazel_options",
          action="append",
          default=[],
          help="Additional options to pass to bazel.",
      )
      parser.add_argument(
          "--output_path",
          default=os.path.join(cwd, "dist"),
          help="Directory to which the jaxlib wheel should be written",
      )
      parser.add_argument(
          "--target_cpu",
          default=None,
          help="CPU platform to target. Default is the same as the host machine. "
          "Currently supported values are 'darwin_arm64' and 'darwin_x86_64'.",
      )
      parser.add_argument(
          "--editable",
          action="store_true",
          help="Create an 'editable' jaxlib build instead of a wheel.",
      )
      add_boolean_argument(
          parser,
          "configure_only",
          default=False,
          help_str="If true, writes a .bazelrc file but does not build jaxlib.",
      )
      args = parser.parse_args()
  
      logging.basicConfig()
      if args.verbose:
          logger.setLevel(logging.DEBUG)
  
      if is_windows() and args.enable_cuda:
          if args.cuda_version is None:
              parser.error(
                  "--cuda_version is needed for Windows CUDA build."
              )
          if args.cudnn_version is None:
              parser.error(
                  "--cudnn_version is needed for Windows CUDA build."
              )
  
      if args.enable_cuda and args.enable_rocm:
          parser.error(
              "--enable_cuda and --enable_rocm cannot be enabled at the same time."
          )
  
      print(BANNER)
  
      output_path = os.path.abspath(args.output_path)
      os.chdir(os.path.dirname(__file__ or args.prog) or ".")
  
      host_cpu = platform.machine()
      wheel_cpus = {
          "darwin_arm64": "arm64",
          "darwin_x86_64": "x86_64",
          "ppc": "ppc64le",
          "aarch64": "aarch64",
      }
      # TODO(phawkins): support other bazel cpu overrides.
      wheel_cpu = (
          wheel_cpus[args.target_cpu]
          if args.target_cpu is not None
          else host_cpu
      )
  
      # Find a working Bazel.
      bazel_path, bazel_version = get_bazel_path(args.bazel_path)
      print(f"Bazel binary path: {bazel_path}")
      print(f"Bazel version: {bazel_version}")
  
      python_bin_path = get_python_bin_path(args.python_bin_path)
      print(f"Python binary path: {python_bin_path}")
      python_version = get_python_version(python_bin_path)
      print(
          "Python version: {}".format(
              ".".join(map(str, python_version))
          )
      )
      check_python_version(python_version)
  
      numpy_version = check_numpy_version(python_bin_path)
      print(f"NumPy version: {numpy_version}")
      check_package_is_installed(
          python_bin_path, python_version, "wheel"
      )
      check_package_is_installed(
          python_bin_path, python_version, "build"
      )
      check_package_is_installed(
          python_bin_path, python_version, "setuptools"
      )
  
      print(
          "MKL-DNN enabled: {}".format(
              "yes" if args.enable_mkl_dnn else "no"
          )
      )
      print(f"Target CPU: {wheel_cpu}")
      print(f"Target CPU features: {args.target_cpu_features}")
  
      cuda_toolkit_path = args.cuda_path
      cudnn_install_path = args.cudnn_path
      rocm_toolkit_path = args.rocm_path
      print(
          "CUDA enabled: {}".format("yes" if args.enable_cuda else "no")
      )
      if args.enable_cuda:
          if cuda_toolkit_path:
              print(f"CUDA toolkit path: {cuda_toolkit_path}")
          if cudnn_install_path:
              print(f"CUDNN library path: {cudnn_install_path}")
          if args.cuda_compute_capabilities is not None:
              print(
                  f"CUDA compute capabilities: {args.cuda_compute_capabilities}"
              )
          if args.cuda_version:
              print(f"CUDA version: {args.cuda_version}")
          if args.cudnn_version:
              print(f"CUDNN version: {args.cudnn_version}")
          print(
              "NCCL enabled: {}".format(
                  "yes" if args.enable_nccl else "no"
              )
          )
  
      print(
          "ROCm enabled: {}".format("yes" if args.enable_rocm else "no")
      )
      if args.enable_rocm:
          if rocm_toolkit_path:
              print(f"ROCm toolkit path: {rocm_toolkit_path}")
          print(f"ROCm amdgpu targets: {args.rocm_amdgpu_targets}")
  
      write_bazelrc(
          python_bin_path=python_bin_path,
          remote_build=args.remote_build,
          cuda_toolkit_path=cuda_toolkit_path,
          cudnn_install_path=cudnn_install_path,
          cuda_version=args.cuda_version,
          cudnn_version=args.cudnn_version,
          rocm_toolkit_path=rocm_toolkit_path,
          cpu=args.target_cpu,
          cuda_compute_capabilities=args.cuda_compute_capabilities,
          rocm_amdgpu_targets=args.rocm_amdgpu_targets,
          bazel_options=args.bazel_options,
          target_cpu_features=args.target_cpu_features,
          wheel_cpu=wheel_cpu,
          enable_mkl_dnn=args.enable_mkl_dnn,
          enable_cuda=args.enable_cuda,
          enable_nccl=args.enable_nccl,
          enable_rocm=args.enable_rocm,
          build_gpu_plugin=args.build_gpu_plugin,
      )
  
      if args.configure_only:
          return
  
      print(
          "\nBuilding XLA and installing it in the jaxlib source tree..."
      )
  
      if (
          not args.build_cuda_kernel_plugin
          and not args.build_cuda_pjrt_plugin
      ):
          command = (
              [bazel_path]
              + args.bazel_startup_options
              + ["run", "--verbose_failures=true"]
              + [
                  "//jaxlib/tools:build_wheel",
                  "--",
                  f"--output_path={output_path}",
                  f"--jaxlib_git_hash={get_githash()}",
                  f"--cpu={wheel_cpu}",
              ]
          )
          if args.build_gpu_plugin:
              command.append("--include_gpu_plugin_extension")
          if args.editable:
              command += ["--editable"]
          print(" ".join(command))
          shell(command)
  
      if args.build_gpu_plugin or args.build_cuda_kernel_plugin:
          build_cuda_kernels_command = (
              [bazel_path]
              + args.bazel_startup_options
              + ["run", "--verbose_failures=true"]
              + [
                  "//jaxlib/tools:build_cuda_kernels_wheel",
                  "--",
                  f"--output_path={output_path}",
                  f"--jaxlib_git_hash={get_githash()}",
                  f"--cpu={wheel_cpu}",
                  f"--cuda_version={args.gpu_plugin_cuda_version}",
              ]
          )
          if args.editable:
              build_cuda_kernels_command.append("--editable")
          print(" ".join(build_cuda_kernels_command))
          shell(build_cuda_kernels_command)
  
      if args.build_gpu_plugin or args.build_cuda_pjrt_plugin:
          build_pjrt_plugin_command = (
              [bazel_path]
              + args.bazel_startup_options
              + ["run", "--verbose_failures=true"]
              + [
                  "//jaxlib/tools:build_gpu_plugin_wheel",
                  "--",
                  f"--output_path={output_path}",
                  f"--jaxlib_git_hash={get_githash()}",
                  f"--cpu={wheel_cpu}",
                  f"--cuda_version={args.gpu_plugin_cuda_version}",
              ]
          )
          if args.editable:
              build_pjrt_plugin_command.append("--editable")
          print(" ".join(build_pjrt_plugin_command))
          shell(build_pjrt_plugin_command)
  
      shell([bazel_path] + args.bazel_startup_options + ["shutdown"])
  
imports: import argparse
  import collections
  import hashlib
  import os
  import platform
  import re
  import shutil
  import stat
  import subprocess
  import sys
  import textwrap
  import urllib.request
  import logging
label_desc: Use image node_id 16 for calling a global function with example usage: main() without return types