Regarding installation issue of jax

[SoumyaCYZ]

Hi,
I have created a conda environment of python3.13.3 (name jax_env) to install jax on my supercomputer. I have tried the " pip install -U "jax[cuda12]" command. However, I have encountered the following error.

ERROR: Exception:
Traceback (most recent call last):
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/base_command.py", line 160, in exc_logging_wrapper
    status = run_func(*args)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/req_command.py", line 247, in wrapper
    return func(self, options, args)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/commands/install.py", line 326, in run
    session = self.get_default_session(options)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/req_command.py", line 98, in get_default_session
    self._session = self.enter_context(self._build_session(options))
                                       ~~~~~~~~~~~~~~~~~~~^^^^^^^^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/req_command.py", line 125, in _build_session
    session = PipSession(
        cache=os.path.join(cache_dir, "http") if cache_dir else None,
    ...<3 lines>...
        ssl_context=ssl_context,
    )
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/network/session.py", line 343, in __init__
    self.headers["User-Agent"] = user_agent()
                                 ~~~~~~~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/network/session.py", line 175, in user_agent
    setuptools_dist = get_default_environment().get_distribution("setuptools")
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 188, in get_distribution
    return next(matches, None)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 185, in <genexpr>
    for distribution in self.iter_all_distributions()
                        ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/base.py", line 612, in iter_all_distributions
    for dist in self._iter_distributions():
                ~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 176, in _iter_distributions
    for dist in finder.find_eggs(location):
                ~~~~~~~~~~~~~~~~^^^^^^^^^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 144, in find_eggs
    yield from self._find_eggs_in_dir(location)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 111, in _find_eggs_in_dir
    from pip._vendor.pkg_resources import find_distributions
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_vendor/pkg_resources/__init__.py", line 2164, in <module>
    register_finder(pkgutil.ImpImporter, find_on_path)
                    ^^^^^^^^^^^^^^^^^^^
AttributeError: module 'pkgutil' has no attribute 'ImpImporter'. Did you mean: 'zipimporter'?
Traceback (most recent call last):
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/bin/pip", line 10, in <module>
    sys.exit(main())
             ~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/main.py", line 70, in main
    return command.main(cmd_args)
           ~~~~~~~~~~~~^^^^^^^^^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/base_command.py", line 101, in main
    return self._main(args)
           ~~~~~~~~~~^^^^^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/base_command.py", line 216, in _main
    self.handle_pip_version_check(options)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/req_command.py", line 179, in handle_pip_version_check
    session = self._build_session(
        options,
    ...<6 lines>...
        fallback_to_certifi=True,
    )
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/cli/req_command.py", line 125, in _build_session
    session = PipSession(
        cache=os.path.join(cache_dir, "http") if cache_dir else None,
    ...<3 lines>...
        ssl_context=ssl_context,
    )
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/network/session.py", line 343, in __init__
    self.headers["User-Agent"] = user_agent()
                                 ~~~~~~~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/network/session.py", line 175, in user_agent
    setuptools_dist = get_default_environment().get_distribution("setuptools")
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 188, in get_distribution
    return next(matches, None)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 185, in <genexpr>
    for distribution in self.iter_all_distributions()
                        ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/base.py", line 612, in iter_all_distributions
    for dist in self._iter_distributions():
                ~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 176, in _iter_distributions
    for dist in finder.find_eggs(location):
                ~~~~~~~~~~~~~~~~^^^^^^^^^^
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 144, in find_eggs
    yield from self._find_eggs_in_dir(location)
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_internal/metadata/importlib/_envs.py", line 111, in _find_eggs_in_dir
    from pip._vendor.pkg_resources import find_distributions
  File "/home/chemistry/phd/cyz218385/.conda/envs/jax_env/lib/python3.1/site-packages/pip/_vendor/pkg_resources/__init__.py", line 2164, in <module>
    register_finder(pkgutil.ImpImporter, find_on_path)
                    ^^^^^^^^^^^^^^^^^^^
AttributeError: module 'pkgutil' has no attribute 'ImpImporter'. Did you mean: 'zipimporter'?

It'll be a great help if you help to solve this issue.
Regards
Soumya

[dfm]

Hi - The issue here is coming from pip rather than JAX, so there must be something wrong with your environment. I'd recommend creating a new environment and making sure that you can pip install other packages first.