Unusual memory usage on gfx1151

[jrhip]

Description

JAX is getting faster and more stable on gfx1151 but a persistent issue I've had is that upon initialising the rocm device, even doing:

import jax

print(jax.devices())

Uses ~50.5 GB of memory which isn't attributed to the python process. This isn't released until killing the python process.

Other than that, it seems to work correctly without error messages or incorrect calculation results.

Pytorch doesn't have this behaviour on the same rocm version.

I get the following messages on running the above which don't seem particularly related but may be useful:

W0827 22:42:57.684494   45101 registration.cpp:741] initialized tool configure for XLA-with-rocprofiler-sdk :: 0x7f0e65503af0 :: 0x55b8765937c0 :: 0x7f0e65503eb0 :: 0x7f0e655036e0
W0827 22:42:57.684519   45101 registration.cpp:791] invoking tool initialize for XLA-with-rocprofiler-sdk :: 0x7f0e65503af0 :: 0x55b8765937c0 :: 0x7f0e65503eb0 :: 0x7f0e655036e0
W0827 22:42:57.903784   45272 rocm_pcie_bandwidth.cc:75] rsmi_dev_gpu_metrics_info_get failed for 0000:c2:00.0: RSMI_STATUS_UNEXPECTED_DATA: Data read (usually from a file) or provided to function is not what was expected
W0827 22:42:57.903812   45272 rocm_executor.cc:1176] Could not determine PCIe bandwidth for device 0 via rocm_smi. Assuming PCIe Gen4 x16.
W0827 22:42:57.931862   45101 rocm_pcie_bandwidth.cc:75] rsmi_dev_gpu_metrics_info_get failed for 0000:c2:00.0: RSMI_STATUS_UNEXPECTED_DATA: Data read (usually from a file) or provided to function is not what was expected
W0827 22:42:57.931870   45101 rocm_executor.cc:1176] Could not determine PCIe bandwidth for device 0 via rocm_smi. Assuming PCIe Gen4 x16.

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

jax-rocm10-pjrt         0.11.1+rocm10.1.0a20260827
jax-rocm10-plugin       0.11.1+rocm10.1.0a20260827

jax:    0.11.1
jaxlib: 0.11.1
numpy:  2.4.4
python: 3.14.3 | packaged by Anaconda, Inc. | (main, Feb 24 2026, 22:51:43) [GCC 14.3.0]
device info: AMD Radeon 8060S Graphics-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='fedora', release='7.1.10-200.fc44.x86_64', version='#1 SMP PREEMPT_DYNAMIC Sun Aug 23 16:15:11 UTC 2026', machine='x86_64')