code: def patch_copy_xla_extension_stubs(dst_dir):
      xla_extension_dir = os.path.join(dst_dir, "xla_extension")
      os.makedirs(xla_extension_dir)
      for stub_name in _XLA_EXTENSION_STUBS:
          stub_path = r.Rlocation(
              "xla/xla/python/xla_extension/" + stub_name
          )
          stub_path = str(
              stub_path
          )  # Make pytype accept os.path.exists(stub_path).
          if (
              stub_name in _OPTIONAL_XLA_EXTENSION_STUBS
              and not os.path.exists(stub_path)
          ):
              continue
          with open(stub_path) as f:
              src = f.read()
          src = src.replace(
              "from xla.python import xla_extension",
              "from .. import xla_extension",
          )
          with open(
              os.path.join(xla_extension_dir, stub_name), "w"
          ) as f:
              f.write(src)
  
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 4 for calling a global function with example usage: patch_copy_xla_extension_stubs(dst_dir) without return types