code: def pytest_collection() -> None:
      if not os.environ.get("JAX_ENABLE_TPU_XDIST", None):
          return
      # When running as an xdist worker, will be something like "gw0"
      xdist_worker_name = os.environ.get("PYTEST_XDIST_WORKER", "")
      if not xdist_worker_name.startswith("gw"):
          return
      xdist_worker_number = int(xdist_worker_name[len("gw") :])
      os.environ.setdefault(
          "TPU_VISIBLE_CHIPS", str(xdist_worker_number)
      )
      os.environ.setdefault("ALLOW_MULTIPLE_LIBTPU_LOAD", "true")
  
label_desc: Use image node_id 2 for calling a global function with example usage: pytest_collection() without return types