docs: def locate_cuda():
      """Locate the CUDA environment on the system
  
      If a valid cuda installation is found this returns a dict with keys 'home', 'nvcc', 'include',
      and 'lib64' and values giving the absolute path to each directory.
  
      Starts by looking for the CUDAHOME env variable. If not found, everything is based on finding
      'nvcc' in the PATH.
  
      If nvcc can't be found, this returns None
      """
      nvcc_bin = 'nvcc'
      if sys.platform.startswith("win"):
          nvcc_bin = 'nvcc.exe'
  
      # first check if the CUDAHOME env variable is in use
      if 'CUDAHOME' in os.environ:
          home = os.environ['CUDAHOME']
          nvcc = os.path.join(home, 'bin', nvcc_bin)
      elif 'CUDA_PATH' in os.environ:
          home = os.environ['CUDA_PATH']
          nvcc = os.path.join(home, 'bin', nvcc_bin)
      else:
          # otherwise, search the PATH for NVCC
          nvcc = find_in_path(nvcc_bin, os.environ['PATH'])
          if nvcc is None:
              logging.warning('The nvcc binary could not be located in your $PATH. Either add it to '
                              'your path, or set $CUDAHOME to enable CUDA extensions')
              return None
          home = os.path.dirname(os.path.dirname(nvcc))
  
      cudaconfig = {'home': home,
                    'nvcc': nvcc,
                    'include': os.path.join(home, 'include'),
                    'lib64':   os.path.join(home, 'lib64')}
      post_args = ['-gencode=arch=compute_30,code=sm_30',
                   '-gencode=arch=compute_50,code=sm_50',
                   '-gencode=arch=compute_60,code=sm_60',
                   '-gencode=arch=compute_60,code=compute_60',
                   '--ptxas-options=-v', '-O2']
  
      if sys.platform == "win32":
          cudaconfig['lib64'] = os.path.join(home, 'lib', 'x64')
          post_args += ['-Xcompiler', '/MD']
      else:
          post_args += ['-c', '--compiler-options', "'-fPIC'"]
  
      for k, v in cudaconfig.items():
          if not os.path.exists(v):
              logging.warning('The CUDA %s path could not be located in %s', k, v)
              return None
  
      cudaconfig['post_args'] = post_args
      return cudaconfig