Disabling memory preallocation fails since last release

[Justin-Tan]

Hi all,

Usually I disable memory preallocation on GPU using the following snippet at the start of any script:

import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"]="false"

This works for jax <= 0.4.20, but not for the current 0.4.23 version. I first noticed this when upgrading and have done a complete reinstall of jax in a new virtualenv on my environment and the issue persists. I'm on Linux using cuda12 drivers and an Nvidia A5000 if that makes a difference.

[hawkinsp]

Can you check something for me: how did you install JAX? Share the output of pip list | grep jax.

(I can't see why it would have stopped working, but one hypothesis is you used the newer plugin installation which we documented briefly although removed from our instructions again because we found a small defect.)

[Justin-Tan]

Hi, I just installed Jax via the standard pip route, the output of pip list | grep jax in the virtualenv that is not able to disable memory preallocation is:

jax                      0.4.23
jaxlib                   0.4.23+cuda12.cudnn89

and in the virtualenv that has disabling memory preallocation working:

jax                       0.4.20
jaxlib                    0.4.20+cuda12.cudnn89

[hawkinsp]

Yes, this looks to be a bug (#19035). The fix will be in the next release.