You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
The persistent compilation cache simply doesn't work - it used to work well with older versions of jax but it seems some breaking changes have occurred in the past weeks.
The problem is that the compilation_cache folder is never created, and I can confirm from the lack of speedup that jax is definitely not using the persistent cache.
Reproduce:
importjaxjax.config.update("jax_compilation_cache_dir", './jax-cache')
importos@jax.jitdefsome_op(A, B):
return (A @ B) * (A+B)
A=jax.numpy.ones((128,))
B=jax.numpy.zeros((128,))
some_op(A, B)
print(os.listdir('./')) # no folder gets created
I don't think its an environment issue - I have a docker image as well that can reproduce it.
System info (python version, jaxlib version, accelerator, etc.)
I tested the provided reproducible code on my Mac-book Pro with M1 Pro chip using jax versions 0.4.26 and 0.4.27.dev20240503 and the corresponding jaxlib versions 0.4.26 and 0.4.27.dev20240503, respectively. In both the cases, a folder named 'jax-cache' was created. Please find the below screenshot for reference.
Description
The persistent compilation cache simply doesn't work - it used to work well with older versions of jax but it seems some breaking changes have occurred in the past weeks.
The problem is that the
compilation_cache
folder is never created, and I can confirm from the lack of speedup that jax is definitely not using the persistent cache.Reproduce:
I don't think its an environment issue - I have a docker image as well that can reproduce it.
System info (python version, jaxlib version, accelerator, etc.)
jax: 0.4.26
jaxlib: 0.4.25
numpy: 1.26.2
python: 3.10.12 (main, Jul 5 2023, 15:02:25) [Clang 14.0.6 ]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Darwin', node='Neels-MacBook-Air.local', release='23.2.0', version='Darwin Kernel Version 23.2.0: Wed Nov 15 21:59:33 PST 2023; root:xnu-10002.61.3~2/RELEASE_ARM64_T8112', machine='arm64')
The text was updated successfully, but these errors were encountered: