Hi I’m trying to use Jax alongside pennylane but when I try to import pennylane I got this message:
AttributeError: module 'jax._src.pjit' has no attribute 'pjit_p'
Here is my pip freeze:
absl-py==2.3.1
accelerate==1.9.0
aiofiles==24.1.0
anyio==4.9.0
appdirs==1.4.4
argon2-cffi==25.1.0
argon2-cffi-bindings==21.2.0
arrow==1.3.0
asttokens==3.0.0
astunparse==1.6.3
async-lru==2.0.5
attrs==25.3.0
autograd==1.8.0
autoray==0.7.2
babel==2.17.0
beautifulsoup4==4.13.4
bleach==6.2.0
cachetools==6.1.0
certifi==2025.7.14
cffi==1.17.1
charset-normalizer==3.4.2
chex==0.1.91
comm==0.2.2
custatevec-cu12==1.9.0
cutensor-cu12==2.2.0
cutensornet-cu12==2.8.0
debugpy==1.8.15
decorator==5.2.1
defusedxml==0.7.1
diastatic-malt==2.15.2
etils==1.13.0
executing==2.2.0
fastjsonschema==2.21.1
filelock==3.18.0
flax==0.12.0
fqdn==1.5.1
fsspec==2025.7.0
gast==0.6.0
h11==0.16.0
hf-xet==1.1.5
httpcore==1.0.9
httpx==0.28.1
huggingface-hub==0.33.4
humanize==4.13.0
idna==3.10
importlib_resources==6.5.2
ipykernel==6.29.5
ipython==9.4.0
ipython_pygments_lexers==1.1.1
isoduration==20.11.0
jax==0.7.2
jaxlib==0.7.2
jedi==0.19.2
Jinja2==3.1.6
json5==0.12.0
jsonpointer==3.0.0
jsonschema==4.24.0
jsonschema-specifications==2025.4.1
jupyter-events==0.12.0
jupyter-lsp==2.2.5
jupyter_client==8.6.3
jupyter_core==5.8.1
jupyter_server==2.16.0
jupyter_server_terminals==0.5.3
jupyterlab==4.4.4
jupyterlab_pygments==0.3.0
jupyterlab_server==2.27.3
loguru==0.7.3
markdown-it-py==4.0.0
MarkupSafe==3.0.2
matplotlib-inline==0.1.7
mdurl==0.1.2
mistune==3.1.3
ml_dtypes==0.5.3
mpmath==1.3.0
msgpack==1.1.1
nbclient==0.10.2
nbconvert==7.16.6
nbformat==5.10.4
nest-asyncio==1.6.0
networkx==3.5
notebook_shim==0.2.4
numpy==2.3.3
nvidia-cublas-cu12==12.6.4.1
nvidia-cuda-cupti-cu12==12.6.80
nvidia-cuda-nvrtc-cu12==12.6.77
nvidia-cuda-runtime-cu12==12.6.77
nvidia-cudnn-cu12==9.5.1.17
nvidia-cufft-cu12==11.3.0.4
nvidia-cufile-cu12==1.11.1.6
nvidia-curand-cu12==10.3.7.77
nvidia-cusolver-cu12==11.7.1.2
nvidia-cusparse-cu12==12.5.4.2
nvidia-cusparselt-cu12==0.6.3
nvidia-nccl-cu12==2.26.2
nvidia-nvjitlink-cu12==12.6.85
nvidia-nvtx-cu12==12.6.77
opt_einsum==3.4.0
optax==0.2.6
orbax-checkpoint==0.11.25
overrides==7.7.0
packaging==25.0
pandocfilters==1.5.1
parso==0.8.4
pennylane==0.42.0
pennylane_lightning==0.42.0
pennylane_lightning_gpu==0.42.0
pennylane_lightning_tensor==0.42.0
pexpect==4.9.0
pillow==11.3.0
platformdirs==4.3.8
prometheus_client==0.22.1
prompt_toolkit==3.0.51
protobuf==6.32.1
psutil==7.0.0
ptyprocess==0.7.0
pure_eval==0.2.3
pycparser==2.22
Pygments==2.19.2
python-dateutil==2.9.0.post0
python-json-logger==3.3.0
PyYAML==6.0.2
pyzmq==27.0.0
referencing==0.36.2
regex==2024.11.6
requests==2.32.4
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rich==14.1.0
rpds-py==0.26.0
rustworkx==0.16.0
safetensors==0.5.3
scipy==1.16.0
scipy-openblas32==0.3.30.0.1
Send2Trash==1.8.3
setuptools==80.9.0
simplejson==3.20.2
six==1.17.0
sniffio==1.3.1
soupsieve==2.7
stack-data==0.6.3
sympy==1.14.0
tensorstore==0.1.77
termcolor==3.1.0
terminado==0.18.1
tinycss2==1.4.0
tokenizers==0.21.2
tomlkit==0.13.3
toolz==1.0.0
torch==2.7.1
torchaudio==2.7.1
torchvision==0.22.1
tornado==6.5.1
tqdm==4.67.1
traitlets==5.14.3
transformers==4.53.2
treescope==0.1.10
triton==3.3.1
types-python-dateutil==2.9.0.20250708
typing_extensions==4.15.0
uri-template==1.3.0
urllib3==2.5.0
wcwidth==0.2.13
webcolors==24.11.1
webencodings==0.5.1
websocket-client==1.8.0
wheel==0.45.1
zipp==3.23.0
I’ve tried to install different versions without success. The machine that I’m using has H100 gpu so with this driver get from nvidia-smi : Driver Version: 560.35.05 CUDA Version: 12.6