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