How to use JAX with xprof #29751

Findus23 · 2025-06-25T17:06:33Z

Findus23
Jun 25, 2025

As this is possibly not a bug in jax, I will create this here instead as feedback to #29129 and #29509.

No matter which combinations of packages I try, I can't get xprof to work with jax.

Following https://github.com/jax-ml/jax/blob/test_772225721/docs/profiling.md and using a fresh virtual env with python 3.13.3 on Debian Testing:

(tempenv-6850222121ccb) ➜  ~/tmp pip install jax
[...]
(tempenv-6850222121ccb) ➜  ~/tmp pip install xprof-nightly
[...]
(tempenv-6850222121ccb) ➜  ~/tmp python
Python 3.13.3 (main, Apr 10 2025, 21:38:51) [GCC 14.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
... 
... jax.profiler.start_trace("/tmp/profile-data")
... 
... # Run the operations to be profiled
... key = jax.random.key(0)
... x = jax.random.normal(key, (5000, 5000))
... y = x @ x
... y.block_until_ready()
... 
... jax.profiler.stop_trace()
... 
2025-06-25 19:03:43.533532: E external/xla/xla/python/profiler/internal/python_hooks.cc:412] Can't import tensorflow.python.profiler.trace
2025-06-25 19:03:44.122895: E external/xla/xla/python/profiler/internal/python_hooks.cc:412] Can't import tensorflow.python.profiler.trace
>>> 
(tempenv-6850222121ccb) ➜  ~/tmp xprof --port 8791 /tmp/profile-data

Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
Unable to load profiler plugin. Import error: cannot import name '_pywrap_profiler_plugin' from 'xprof.convert' (unknown location)
Traceback (most recent call last):
  File "/home/lukas/.virtualenvs/tempenv-6850222121ccb/bin/xprof", line 8, in <module>
    sys.exit(main())
             ~~~~^^
  File "/home/lukas/.virtualenvs/tempenv-6850222121ccb/lib/python3.13/site-packages/xprof/server.py", line 171, in main
    launch_server(logdir, port)
    ~~~~~~~~~~~~~^^^^^^^^^^^^^^
  File "/home/lukas/.virtualenvs/tempenv-6850222121ccb/lib/python3.13/site-packages/xprof/server.py", line 108, in launch_server
    run_server(plugin, _get_wildcard_address(port), port)
    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lukas/.virtualenvs/tempenv-6850222121ccb/lib/python3.13/site-packages/xprof/server.py", line 51, in run_server
    app = make_wsgi_app(plugin)
  File "/home/lukas/.virtualenvs/tempenv-6850222121ccb/lib/python3.13/site-packages/xprof/server.py", line 32, in make_wsgi_app
    apps = plugin.get_plugin_apps()
           ^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'get_plugin_apps'

No matter which xprof I use, I always get an error: cannot import name '_pywrap_profiler_plugin' from 'xprof.convert' error on startup.

>>> from xprof.convert import _pywrap_profiler_plugin
Traceback (most recent call last):
  File "<python-input-0>", line 1, in <module>
    from xprof.convert import _pywrap_profiler_plugin
ImportError: cannot import name '_pywrap_profiler_plugin' from 'xprof.convert' (unknown location)
>>> import xprof.convert
>>> help(xprof.convert)
Help on package xprof.convert in xprof:

NAME
    xprof.convert

PACKAGE CONTENTS
    csv_writer
    dcn_collective_stats_proto_to_gviz
    diagnostics
    inference_stats_proto_to_gviz
    input_pipeline_proto_to_gviz
    overview_page_proto_to_gviz
    raw_to_tool_data
    trace_events_json

FILE
    (built-in)

(tempenv-6850222121ccb) ➜  ~/tmp pip freeze
aiohappyeyeballs==2.6.1
aiohttp==3.12.13
aiosignal==1.3.2
attrs==25.3.0
cachetools==5.5.2
certifi==2025.6.15
charset-normalizer==3.4.2
cheroot==10.0.1
decorator==5.2.1
etils==1.12.2
frozenlist==1.7.0
fsspec==2025.5.1
gcsfs==2025.5.1
google-api-core==2.25.1
google-auth==2.40.3
google-auth-oauthlib==1.2.2
google-cloud-core==2.4.3
google-cloud-storage==3.1.1
google-crc32c==1.7.1
google-resumable-media==2.7.2
googleapis-common-protos==1.70.0
gviz-api==1.10.0
idna==3.10
importlib_resources==6.5.2
jaraco.functools==4.2.1
jax==0.6.2
jaxlib==0.6.2
MarkupSafe==3.0.2
ml_dtypes==0.5.1
more-itertools==10.7.0
multidict==6.5.1
numpy==2.3.1
oauthlib==3.3.1
opt_einsum==3.4.0
packaging==25.0
propcache==0.3.2
proto-plus==1.26.1
protobuf==6.31.1
pyasn1==0.6.1
pyasn1_modules==0.4.2
requests==2.32.4
requests-oauthlib==2.0.0
rsa==4.9.1
scipy==1.16.0
setuptools==80.9.0
six==1.17.0
typing_extensions==4.14.0
urllib3==2.5.0
Werkzeug==3.1.3
xprof-nightly==2.21.1a20250625
yarl==1.20.1
zipp==3.23.0

Answered by Findus23

Jun 25, 2025

Okay, I just found out a bit too late that if I use python 3.12 instead, things seem to be working better

View full answer

Findus23 · 2025-06-25T17:11:21Z

Findus23
Jun 25, 2025
Author

Okay, I just found out a bit too late that if I use python 3.12 instead, things seem to be working better

4 replies

biao-geng-me Jul 12, 2025

Hello. Is it working for you with Python 3.12? I was origally using Python 3.13 and had exactly the same errors as yours. I tried Python 3.12.11 after seeing this comment, but no luck for me, it is still the same error.

Findus23 Jul 14, 2025
Author

Yes, using python 3.12, everything is working for me as expected:

➜  ~ mktmpenv
Creating tempenv-62ba209210a36 via ~/.virtualenvs/virtualfish/bin/python …
(tempenv-62ba209210a36) ➜  ~ pip install xprof-nightly
[...]
(tempenv-62ba209210a36) ➜  ~ xprof ~/tmp/jax-pr/
Attempting to start XProf server:
  Log Directory: ~/tmp/jax-pr
  Port: 8791
Unable to load profiler plugin. Import error: cannot import name '_pywrap_profiler_plugin' from 'xprof.convert' (unknown location)
Traceback (most recent call last):
[...]
(tempenv-62ba209210a36) ➜  ~ [1] deactivate 
Removing temporary virtualenv tempenv-62ba209210a36
➜  ~ mktmpenv  --python=/home/lukas/.pyenv/versions/3.12.9/bin/python
Creating tempenv-2dd328365a98 via ~/.pyenv/versions/3.12.9/bin/python …
(tempenv-2dd328365a98) ➜  ~ pip install xprof-nightly
[...]
(tempenv-2dd328365a98) ➜  ~ xprof ~/tmp/jax-pr/
Attempting to start XProf server:
  Log Directory: ~/tmp/jax-pr
  Port: 8791
XProf at http://localhost:8791/ (Press CTRL+C to quit)

biao-geng-me Jul 14, 2025

Thank you very much for demonstrating it in your environment. I was able to run it inside WSL on Windows. It seems my original error was related to Windows and might have been a different error.

cjim8889 Jul 25, 2025

Same issues when I was using python 3.13

biao-geng-me · 2025-07-12T17:19:20Z

biao-geng-me
Jul 12, 2025

Just wanted to add that I came across the same errors trying to use Tensorboard/xprof on both Linux and Windows10.

1 reply

biao-geng-me Jul 14, 2025

Actually, mine seems to be a different issue. When launching xprof alone, the error is ValueError: Protocol not known: c when the code is trying to get the filesystem protocol. It's the same error using Python 3.12 or 3.13.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

How to use JAX with xprof #29751

Uh oh!

{{title}}

Uh oh!

Replies: 2 comments 5 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

How to use JAX with xprof #29751

Uh oh!

Findus23 Jun 25, 2025

Replies: 2 comments · 5 replies

Uh oh!

Findus23 Jun 25, 2025 Author

Uh oh!

biao-geng-me Jul 12, 2025

Uh oh!

Findus23 Jul 14, 2025 Author

Uh oh!

biao-geng-me Jul 14, 2025

Uh oh!

cjim8889 Jul 25, 2025

Uh oh!

biao-geng-me Jul 12, 2025

Uh oh!

biao-geng-me Jul 14, 2025

Findus23
Jun 25, 2025

Replies: 2 comments 5 replies

Findus23
Jun 25, 2025
Author

Findus23 Jul 14, 2025
Author

biao-geng-me
Jul 12, 2025