Memory Viewer always analyzes first HLO module

[Findus23]

Hi,

I recently stumbled across this tool trying to understand the memory allocations of a complex JAX software a bit better.

It seems like the memory_viewer output works quite well and corresponds to what I assumed from looking at the XLA dump .cpu_after_optimizations-buffer-assignment.txt files before.

But there seems to be one major bug in the tool: No matter which HLO Module is selected in the sidebar, the first HLO module is reported.

I can see in the HTTP requests and CLI output that even if I select jit_test(5) instead of jit_broadcast_in_dim(3), the latter is still used and the HTTP request is made to
http://localhost:8791/data/plugin/profile/data?run=2025_06_27_16_03_57&tag=memory_viewer&host=lukasnotebook&module_name=jit_broadcast_in_dim(3)&memory_space=0
If instead I rewrite the HTTP request in transit to instead have the right &module_name=jit_test(5) I will see the correct visualisation for this HLO module.

Using:

• Python 3.12.9 ¹
• xprof-nightly==2.21.1a20250627
• Firefox 139.0.4 and Chromium 137.0.7151.68

Footnotes

1. There seems to be another issue with python 3.13. If I use python 3.13.3 and install xprof-nightly in a fresh venv, xprof always fails to start with Import error: cannot import name '_pywrap_profiler_plugin' from 'xprof.convert' (unknown location) ↩

[Matt-Hurd]

Thanks for reporting this! I've managed to replicate the issue, and I'm investigating it.

For the Python3.13 issue, we haven't yet started supporting it fully. Our releases include a binary with python bindings that needs to be compiled for each minor Python release. We'll get Python3.13 working soon as well.