jax.experimental.compilation_cache 模块

JAX 磁盘编译缓存。

API

jax.experimental.compilation_cache.compilation_cache.set_cache_dir(path)

设置持久化编译缓存目录。

调用此函数后，jit 编译的函数将被保存在 path 中，因此如果进程被重启或再次运行，则无需重新编译。这也会告诉 Jax 在编译之前从何处查找已编译的函数。

返回类型:

无

jax.experimental.compilation_cache.compilation_cache.reset_cache()

恢复到原始的、未初始化的状态。

返回类型:

无