xla_extension.so: cannot enable executable stack as shared object requires: Invalid argument

[DSilva27]

Description

I got this bug today after creating a fresh environment. I tried creating the env using uv and venv, and it happened in both cases. The bug happened in both python 3.11.11 and 3.13.1 (did not check more). Here is the full error message

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/path/to/my/env/myenv_name/lib/python3.11/site-packages/jax/__init__.py", line 25, in <module>
    from jax._src.cloud_tpu_init import cloud_tpu_init as _cloud_tpu_init
  File "/path/to/my/env/myenv_name/lib/python3.11/site-packages/jax/_src/cloud_tpu_init.py", line 19, in <module>
    from jax._src import config
  File "/path/to/my/env/myenv_name/lib/python3.11/site-packages/jax/_src/config.py", line 27, in <module>
    from jax._src.lib import guard_lib
  File "/path/to/my/env/myenv_name/lib/python3.11/site-packages/jax/_src/lib/__init__.py", line 87, in <module>
    import jaxlib.xla_client as xla_client
  File "/path/to/my/env/myenv_name/lib/python3.11/site-packages/jaxlib/xla_client.py", line 33, in <module>
    from . import xla_extension as _xla
ImportError: /path/to/my/env/myenv_name/lib/python3.11/site-packages/jaxlib/xla_extension.so: cannot enable executable stack as shared object requires: Invalid argument

I managed to fix this by running execstack -c /path/to/my/env/myenv_name/lib/python3.11/site-packages/jaxlib/xla_extension.so.

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.5.1
jaxlib: 0.5.1
numpy:  2.2.3
python: 3.11.11 (main, Feb 12 2025, 14:51:05) [Clang 19.1.6 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='asusvivobook', release='6.6.75-2-MANJARO', version='#1 SMP PREEMPT_DYNAMIC Mon Feb  3 17:53:46 UTC 2025', machine='x86_64')

[nshepperd]

This seems to be a regression in jaxlib 0.5.1. The xla_extension.so indeed has the stack marked as executable:

$ objdump -x .venv/lib/python3.12/site-packages/jaxlib/xla_extension.so | grep -A1 'STACK off'
   STACK off    0x0000000000000000 vaddr 0x0000000000000000 paddr 0x0000000000000000 align 2**4
         filesz 0x0000000000000000 memsz 0x0000000000000000 flags rwx

While jaxlib 0.5.0 does not:

$ objdump -x .venv/lib/python3.12/site-packages/jaxlib/xla_extension.so | grep -A1 'STACK off'
   STACK off    0x0000000000000000 vaddr 0x0000000000000000 paddr 0x0000000000000000 align 2**4
         filesz 0x0000000000000000 memsz 0x0000000000000000 flags rw-

An unhappy coincidence given that dlopen'ing shared libraries with executable stack was officially removed from glibc less than a month ago, in version 2.41!

[nshepperd]

Also, this might actually be an XLA issue, since xla_extension.so comes from XLA. But I can't be sure, because I couldn't reproduce the issue with xla_extension.so compiling XLA by itself under some configurations, and couldn't get the compilation to work at all under other configurations. So this may also be an issue with the compilation environment of jaxlib's wheels specifically.

[sandycat27]

I have the same issue. Reverting to jax version 0.5.0 cleared the issue.