How to use JAX with xprof

[Findus23]

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

[Findus23]

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

[biao-geng-me]

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]

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]

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.

[wchen99998]

Same issues when I was using python 3.13

[biao-geng-me]

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

[biao-geng-me]

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.

[Matt-Hurd]

In the future, feel free to open up issues on the XProf GitHub repo (https://github.com/openxla/xprof) where we're much more likely to see and respond in a more timely manner.

For the issues in this thread:
ValueError: Protocol not known: c should be fixed by openxla/xprof@95c5984.

Unable to load profiler plugin. Import error: cannot import name '_pywrap_profiler_plugin' from 'xprof.convert' is caused by two issues. One of them is resolved in the above commit. The second issue comes from a mistake I made that led to our build system not creating builds for Python 3.13. This will be resolved once openxla/xprof#1682 is merged. It'll be merged by the End of Week, just needing a quick review on it. I'll be able to generate Python 3.13 versions for the latest release as well.

For now, the preferred resolution is to use the latest release (or nightly), and Python 3.9-3.12.