jax.experimental.compilation_cache 模块#

JAX 磁盘编译缓存。

API#

jax.experimental.compilation_cache.compilation_cache.is_initialized()[源代码]#

已弃用。

返回是否启用缓存。初始化可以延迟,因此不检查初始化状态。保留名称是为了向后兼容。

返回类型:

bool

jax.experimental.compilation_cache.compilation_cache.initialize_cache(path)[源代码]#

此 API 已弃用;请改用 set_cache_dir。

设置路径。要生效,应在任何调用 get_executable_and_time() 和 put_executable_and_time() 之前调用。

返回类型:

jax.experimental.compilation_cache.compilation_cache.set_cache_dir(path)[源代码]#

设置持久编译缓存目录。

调用此方法后,jit 编译的函数将保存到 path,因此如果重新启动进程或以其他方式再次运行,则无需重新编译。 这也告诉 Jax 在编译之前在哪里查找已编译的函数。

返回类型:

jax.experimental.compilation_cache.compilation_cache.reset_cache()[源代码]#

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

返回类型: