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()

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

返回类型:

无