`pixi` munges the jax wheel

[alanhdu]

Checks

• I have checked that this issue has not already been reported.

• I have confirmed this bug exists on the latest version of pixi, using pixi --version.

Reproducible example

Commands I ran and their output:

pixi run python -c "import jax"

pixi.toml/pyproject.toml file that reproduces my issue:

[workspace]
authors = ["Alan Du <alanhdu@gmail.com>"]
channels = ["conda-forge"]
name = "playground"
platforms = ["linux-64"]
version = "0.1.0"

[tasks]

[dependencies]
python = ">=3.14.4,<3.15"

[pypi-dependencies]
jax = ">=0.9.2, <0.10"

pixi info output:

System
------------
       Pixi version: 0.67.0
        TLS backend: rustls
           Platform: linux-64
   Virtual packages: __unix=0=0
                   : __linux=6.8.0=0
                   : __glibc=2.39=0
                   : __archspec=1=zen4
          Cache dir: /home/alan/.cache/rattler/cache
       Auth storage: /home/alan/.rattler/credentials.json
   Config locations: No config files found

Global
------------
            Bin dir: /home/alan/.pixi/bin
    Environment dir: /home/alan/.pixi/envs
       Manifest dir: /home/alan/.pixi/manifests/pixi-global.toml

Workspace
------------
               Name: playground
            Version: 0.1.0
      Manifest file: /home/alan/workspace/playground/pixi.toml
       Last updated: 09-04-2026 14:56:59

Environments
------------
        Environment: default
           Features: default
           Channels: conda-forge
   Dependency count: 1
       Dependencies: python
  PyPI Dependencies: jax
   Target platforms: linux-64
    Prefix location: /home/alan/workspace/playground/.pixi/envs/default

Other files (e.g. script files, source files, etc.):

Issue description

If you run pixi run python -c "import jax", this gives me:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
    import jax
  File "/home/alan/workspace/playground/.pixi/envs/default/lib/python3.14/site-packages/jax/__init__.py", line 25, in <module>
    from jax._src.cloud_tpu_init import cloud_tpu_init as _cloud_tpu_init
  File "/home/alan/workspace/playground/.pixi/envs/default/lib/python3.14/site-packages/jax/_src/cloud_tpu_init.py", line 20, in <module>
    from jax._src import config
  File "/home/alan/workspace/playground/.pixi/envs/default/lib/python3.14/site-packages/jax/_src/config.py", line 29, in <module>
    from jax._src import logging_config
  File "/home/alan/workspace/playground/.pixi/envs/default/lib/python3.14/site-packages/jax/_src/logging_config.py", line 17, in <module>
    from jax._src.lib import utils
  File "/home/alan/workspace/playground/.pixi/envs/default/lib/python3.14/site-packages/jax/_src/lib/__init__.py", line 89, in <module>
    import jaxlib.xla_client as xla_client  # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/alan/workspace/playground/.pixi/envs/default/lib/python3.14/site-packages/jaxlib/xla_client.py", line 28, in <module>
    from jaxlib import _jax as _xla
ImportError: libgrpc.so.52: cannot open shared object file: No such file or directory

Expected behavior

This imports without any issues at all.

[alanhdu]

FWIW, I want to use the wheel explicitly because it makes it easier to support other platforms that conda-forge doesn't support yet (Windows) and to have other environments that must use wheels (e.g. tpu) since then I only need to switch the extras parameter between the different evnrionemnts.

This does seem to be solved if I explicitly install onednn and onednn-cpu-threadpool on Linux at least -- is pixi stripping out that part of the binary wheel somehow?

[ruben-arts]

There is not really a reason for Pixi to something strange here. There are two ideas that come to mind for me:

1. Could you try to add a higher libc with pixi workspace system-requirements add glibc 2.39 then fully reinstall with rm pixi.lock && pixi clean && pixi install
2. If 1 doesn't work, try to remove the jax pypi-dependency from the manifest and use pip to install jax into the pixi environment and see if that works. If that works please report that here because that would indicate a bug on the Pixi side.