search

conda-forge / jax-feedstock

A conda-smithy repository for jax.
BSD 3-Clause "New" or "Revised" License
5 stars 15 forks source link

import jax;key = jax.random.PRNGKey(0) fails with Illegal instruction on QEMU-emulated linux-aarch64 #144

Open traversaro opened 1 month ago

traversaro commented 1 month ago

Solution to issue cannot be found in the documentation.

Issue

The following instructions on linux-64 fail with illegal instruction:

docker run --rm -it --platform linux/arm64 condaforge/miniforge3
conda install jax
python -c "import jax;key = jax.random.PRNGKey(0)"
Illegal instruction

I was not able to test on a physical linux-aarch64 system, so I am not sure if that happens also there. With the pip wheel, this does not happens.

Installed packages

(base) root@3d003c853c00:/# conda list
# packages in environment at /opt/conda:
#
# Name                    Version                   Build  Channel
_openmp_mutex             4.5                       2_gnu    conda-forge
archspec                  0.2.3              pyhd8ed1ab_0    conda-forge
boltons                   24.0.0             pyhd8ed1ab_0    conda-forge
brotli-python             1.1.0           py310hbb3657e_1    conda-forge
bzip2                     1.0.8                h31becfc_5    conda-forge
c-ares                    1.28.1               h31becfc_0    conda-forge
ca-certificates           2024.2.2             hcefe29a_0    conda-forge
certifi                   2024.2.2           pyhd8ed1ab_0    conda-forge
cffi                      1.16.0          py310hce94938_0    conda-forge
charset-normalizer        3.3.2              pyhd8ed1ab_0    conda-forge
colorama                  0.4.6              pyhd8ed1ab_0    conda-forge
conda                     24.5.0          py310h4c7bcd0_0    conda-forge
conda-libmamba-solver     24.1.0             pyhd8ed1ab_0    conda-forge
conda-package-handling    2.2.0              pyh38be061_0    conda-forge
conda-package-streaming   0.9.0              pyhd8ed1ab_0    conda-forge
distro                    1.9.0              pyhd8ed1ab_0    conda-forge
fmt                       10.2.1               h2a328a1_0    conda-forge
frozendict                2.4.4           py310hb52b2da_0    conda-forge
icu                       73.2                 h787c7f5_0    conda-forge
idna                      3.7                pyhd8ed1ab_0    conda-forge
importlib-metadata        7.1.0              pyha770c72_0    conda-forge
importlib_metadata        7.1.0                hd8ed1ab_0    conda-forge
jax                       0.4.27             pyhd8ed1ab_0    conda-forge
jaxlib                    0.4.23          cpu_py310h6928413_2    conda-forge
jsonpatch                 1.33               pyhd8ed1ab_0    conda-forge
jsonpointer               2.4             py310h4c7bcd0_3    conda-forge
keyutils                  1.6.1                h4e544f5_0    conda-forge
krb5                      1.21.2               hc419048_0    conda-forge
ld_impl_linux-aarch64     2.40                 h2d8c526_0    conda-forge
libabseil                 20240116.2      cxx17_h2f0025b_0    conda-forge
libarchive                3.7.2                hd2f85e0_1    conda-forge
libblas                   3.9.0           22_linuxaarch64_openblas    conda-forge
libcblas                  3.9.0           22_linuxaarch64_openblas    conda-forge
libcurl                   8.7.1                h4e8248e_0    conda-forge
libedit                   3.1.20191231         he28a2e2_2    conda-forge
libev                     4.33                 h31becfc_2    conda-forge
libffi                    3.4.2                h3557bc0_5    conda-forge
libgcc-ng                 13.2.0               hf8544c7_5    conda-forge
libgfortran-ng            13.2.0               he9431aa_7    conda-forge
libgfortran5              13.2.0               h87d9d71_7    conda-forge
libgomp                   13.2.0               hf8544c7_5    conda-forge
libgrpc                   1.62.2               h98a9317_0    conda-forge
libiconv                  1.17                 h31becfc_2    conda-forge
liblapack                 3.9.0           22_linuxaarch64_openblas    conda-forge
libmamba                  1.5.8                hea3be6c_0    conda-forge
libmambapy                1.5.8           py310h5938bc3_0    conda-forge
libnghttp2                1.58.0               hb0e430d_1    conda-forge
libnsl                    2.0.1                h31becfc_0    conda-forge
libopenblas               0.3.27          pthreads_h5a5ec62_0    conda-forge
libprotobuf               4.25.3               h648ac29_0    conda-forge
libre2-11                 2023.09.01           h9d008c2_2    conda-forge
libsolv                   0.7.28               hd84c7bf_2    conda-forge
libsqlite                 3.45.2               h194ca79_0    conda-forge
libssh2                   1.11.0               h492db2e_0    conda-forge
libstdcxx-ng              13.2.0               h9a76618_5    conda-forge
libuuid                   2.38.1               hb4cce97_0    conda-forge
libxcrypt                 4.4.36               h31becfc_1    conda-forge
libxml2                   2.12.6               h3091e33_1    conda-forge
libzlib                   1.2.13               h31becfc_5    conda-forge
lz4-c                     1.9.4                hd600fc2_0    conda-forge
lzo                       2.10              h516909a_1000    conda-forge
mamba                     1.5.8           py310hcbdc16a_0    conda-forge
menuinst                  2.0.2           py310h4c7bcd0_0    conda-forge
ml_dtypes                 0.4.0           py310h34310a0_1    conda-forge
ncurses                   6.4.20240210         h0425590_0    conda-forge
numpy                     1.26.4          py310hcbab775_0    conda-forge
openssl                   3.3.0                h68df207_3    conda-forge
opt-einsum                3.3.0                hd8ed1ab_2    conda-forge
opt_einsum                3.3.0              pyhc1e730c_2    conda-forge
packaging                 24.0               pyhd8ed1ab_0    conda-forge
pip                       24.0               pyhd8ed1ab_0    conda-forge
platformdirs              4.2.0              pyhd8ed1ab_0    conda-forge
pluggy                    1.4.0              pyhd8ed1ab_0    conda-forge
pybind11-abi              4                    hd8ed1ab_3    conda-forge
pycosat                   0.6.6           py310hb299538_0    conda-forge
pycparser                 2.22               pyhd8ed1ab_0    conda-forge
pysocks                   1.7.1              pyha2e5f31_6    conda-forge
python                    3.10.14         hbbe8eec_0_cpython    conda-forge
python_abi                3.10                    4_cp310    conda-forge
re2                       2023.09.01           h9caee61_2    conda-forge
readline                  8.2                  h8fc344f_1    conda-forge
reproc                    14.2.4.post0         h31becfc_1    conda-forge
reproc-cpp                14.2.4.post0         h2f0025b_1    conda-forge
requests                  2.31.0             pyhd8ed1ab_0    conda-forge
ruamel.yaml               0.18.6          py310hb299538_0    conda-forge
ruamel.yaml.clib          0.2.8           py310hb299538_0    conda-forge
scipy                     1.13.1          py310h70fbbe5_0    conda-forge
setuptools                69.5.1             pyhd8ed1ab_0    conda-forge
tk                        8.6.13               h194ca79_0    conda-forge
tqdm                      4.66.2             pyhd8ed1ab_0    conda-forge
truststore                0.8.0              pyhd8ed1ab_0    conda-forge
tzdata                    2024a                h0c530f3_0    conda-forge
urllib3                   2.2.1              pyhd8ed1ab_0    conda-forge
wheel                     0.43.0             pyhd8ed1ab_1    conda-forge
xz                        5.2.6                h9cdd2b7_0    conda-forge
yaml-cpp                  0.8.0                h2f0025b_0    conda-forge
zipp                      3.17.0             pyhd8ed1ab_0    conda-forge
zstandard                 0.22.0          py310h468e293_0    conda-forge
zstd                      1.5.5                h4c53e97_0    conda-forge

Environment info

(base) root@3d003c853c00:/# conda list
# packages in environment at /opt/conda:
#
# Name                    Version                   Build  Channel
_openmp_mutex             4.5                       2_gnu    conda-forge
archspec                  0.2.3              pyhd8ed1ab_0    conda-forge
boltons                   24.0.0             pyhd8ed1ab_0    conda-forge
brotli-python             1.1.0           py310hbb3657e_1    conda-forge
bzip2                     1.0.8                h31becfc_5    conda-forge
c-ares                    1.28.1               h31becfc_0    conda-forge
ca-certificates           2024.2.2             hcefe29a_0    conda-forge
certifi                   2024.2.2           pyhd8ed1ab_0    conda-forge
cffi                      1.16.0          py310hce94938_0    conda-forge
charset-normalizer        3.3.2              pyhd8ed1ab_0    conda-forge
colorama                  0.4.6              pyhd8ed1ab_0    conda-forge
conda                     24.5.0          py310h4c7bcd0_0    conda-forge
conda-libmamba-solver     24.1.0             pyhd8ed1ab_0    conda-forge
conda-package-handling    2.2.0              pyh38be061_0    conda-forge
conda-package-streaming   0.9.0              pyhd8ed1ab_0    conda-forge
distro                    1.9.0              pyhd8ed1ab_0    conda-forge
fmt                       10.2.1               h2a328a1_0    conda-forge
frozendict                2.4.4           py310hb52b2da_0    conda-forge
icu                       73.2                 h787c7f5_0    conda-forge
idna                      3.7                pyhd8ed1ab_0    conda-forge
importlib-metadata        7.1.0              pyha770c72_0    conda-forge
importlib_metadata        7.1.0                hd8ed1ab_0    conda-forge
jax                       0.4.27             pyhd8ed1ab_0    conda-forge
jaxlib                    0.4.23          cpu_py310h6928413_2    conda-forge
jsonpatch                 1.33               pyhd8ed1ab_0    conda-forge
jsonpointer               2.4             py310h4c7bcd0_3    conda-forge
keyutils                  1.6.1                h4e544f5_0    conda-forge
krb5                      1.21.2               hc419048_0    conda-forge
ld_impl_linux-aarch64     2.40                 h2d8c526_0    conda-forge
libabseil                 20240116.2      cxx17_h2f0025b_0    conda-forge
libarchive                3.7.2                hd2f85e0_1    conda-forge
libblas                   3.9.0           22_linuxaarch64_openblas    conda-forge
libcblas                  3.9.0           22_linuxaarch64_openblas    conda-forge
libcurl                   8.7.1                h4e8248e_0    conda-forge
libedit                   3.1.20191231         he28a2e2_2    conda-forge
libev                     4.33                 h31becfc_2    conda-forge
libffi                    3.4.2                h3557bc0_5    conda-forge
libgcc-ng                 13.2.0               hf8544c7_5    conda-forge
libgfortran-ng            13.2.0               he9431aa_7    conda-forge
libgfortran5              13.2.0               h87d9d71_7    conda-forge
libgomp                   13.2.0               hf8544c7_5    conda-forge
libgrpc                   1.62.2               h98a9317_0    conda-forge
libiconv                  1.17                 h31becfc_2    conda-forge
liblapack                 3.9.0           22_linuxaarch64_openblas    conda-forge
libmamba                  1.5.8                hea3be6c_0    conda-forge
libmambapy                1.5.8           py310h5938bc3_0    conda-forge
libnghttp2                1.58.0               hb0e430d_1    conda-forge
libnsl                    2.0.1                h31becfc_0    conda-forge
libopenblas               0.3.27          pthreads_h5a5ec62_0    conda-forge
libprotobuf               4.25.3               h648ac29_0    conda-forge
libre2-11                 2023.09.01           h9d008c2_2    conda-forge
libsolv                   0.7.28               hd84c7bf_2    conda-forge
libsqlite                 3.45.2               h194ca79_0    conda-forge
libssh2                   1.11.0               h492db2e_0    conda-forge
libstdcxx-ng              13.2.0               h9a76618_5    conda-forge
libuuid                   2.38.1               hb4cce97_0    conda-forge
libxcrypt                 4.4.36               h31becfc_1    conda-forge
libxml2                   2.12.6               h3091e33_1    conda-forge
libzlib                   1.2.13               h31becfc_5    conda-forge
lz4-c                     1.9.4                hd600fc2_0    conda-forge
lzo                       2.10              h516909a_1000    conda-forge
mamba                     1.5.8           py310hcbdc16a_0    conda-forge
menuinst                  2.0.2           py310h4c7bcd0_0    conda-forge
ml_dtypes                 0.4.0           py310h34310a0_1    conda-forge
ncurses                   6.4.20240210         h0425590_0    conda-forge
numpy                     1.26.4          py310hcbab775_0    conda-forge
openssl                   3.3.0                h68df207_3    conda-forge
opt-einsum                3.3.0                hd8ed1ab_2    conda-forge
opt_einsum                3.3.0              pyhc1e730c_2    conda-forge
packaging                 24.0               pyhd8ed1ab_0    conda-forge
pip                       24.0               pyhd8ed1ab_0    conda-forge
platformdirs              4.2.0              pyhd8ed1ab_0    conda-forge
pluggy                    1.4.0              pyhd8ed1ab_0    conda-forge
pybind11-abi              4                    hd8ed1ab_3    conda-forge
pycosat                   0.6.6           py310hb299538_0    conda-forge
pycparser                 2.22               pyhd8ed1ab_0    conda-forge
pysocks                   1.7.1              pyha2e5f31_6    conda-forge
python                    3.10.14         hbbe8eec_0_cpython    conda-forge
python_abi                3.10                    4_cp310    conda-forge
re2                       2023.09.01           h9caee61_2    conda-forge
readline                  8.2                  h8fc344f_1    conda-forge
reproc                    14.2.4.post0         h31becfc_1    conda-forge
reproc-cpp                14.2.4.post0         h2f0025b_1    conda-forge
requests                  2.31.0             pyhd8ed1ab_0    conda-forge
ruamel.yaml               0.18.6          py310hb299538_0    conda-forge
ruamel.yaml.clib          0.2.8           py310hb299538_0    conda-forge
scipy                     1.13.1          py310h70fbbe5_0    conda-forge
setuptools                69.5.1             pyhd8ed1ab_0    conda-forge
tk                        8.6.13               h194ca79_0    conda-forge
tqdm                      4.66.2             pyhd8ed1ab_0    conda-forge
truststore                0.8.0              pyhd8ed1ab_0    conda-forge
tzdata                    2024a                h0c530f3_0    conda-forge
urllib3                   2.2.1              pyhd8ed1ab_0    conda-forge
wheel                     0.43.0             pyhd8ed1ab_1    conda-forge
xz                        5.2.6                h9cdd2b7_0    conda-forge
yaml-cpp                  0.8.0                h2f0025b_0    conda-forge
zipp                      3.17.0             pyhd8ed1ab_0    conda-forge
zstandard                 0.22.0          py310h468e293_0    conda-forge
zstd                      1.5.5                h4c53e97_0    conda-forge

I did not do a detailed analysis, but the issue seems present also in older jax/jaxsim versions.

traversaro commented 1 month ago

More detailed backtrace:

(base) root@3d003c853c00:/# python -c "import faulthandler; faulthandler.enable(); import jax; key = jax.random.PRNGKey(0)"
Fatal Python error: Illegal instruction

Current thread 0x0000005501e7a420 (most recent call first):
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1186 in __call__
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/profiler.py", line 335 in wrapper
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/pjit.py", line 1488 in _pjit_call_impl_python
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/pjit.py", line 1534 in call_impl_cache_miss
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/pjit.py", line 1558 in _pjit_call_impl
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/core.py", line 879 in process_primitive
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/core.py", line 391 in bind_with_trace
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/core.py", line 2789 in bind
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/pjit.py", line 182 in _python_pjit_helper
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/pjit.py", line 305 in cache_miss
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 179 in reraise_with_filtered_traceback
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 86 in apply_primitive
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/core.py", line 879 in process_primitive
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/core.py", line 391 in bind_with_trace
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/core.py", line 387 in bind
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/lax/lax.py", line 559 in _convert_element_type
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2592 in array
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2667 in asarray
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/prng.py", line 634 in random_seed
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/random.py", line 203 in _key
  File "/opt/conda/lib/python3.10/site-packages/jax/_src/random.py", line 241 in PRNGKey
  File "<string>", line 1 in <module>

Extension modules: numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator, jaxlib.cpu_feature_guard, zstandard.backend_c (total: 15)
Illegal instruction