zoomy_jax.symsolver.compile_cache

The persistent compilation cache for programs with host callbacks.

jax (0.9) writes no persistent-cache entry for an executable with Python host callbacks (jax.debug.callback, io_callback, pure_callback: ordered or not, every one is an xla_ffi_python_*_callback custom call whose callable is bound to the executable at load), so the march of the symbolic solvers, which carries the CFL warning, the linear-solve accept gate, check.march and io.write, compiled anew in every process (ANALYSIS 65). The executable itself does not depend on the callables: the custom call names its callback by an index into the list handed to the compiler, and jaxlib’s deserialize_executable takes that list. So inside callback_cache() an executable with callbacks is cached as jax caches any other one (the same key, compilation_cache.get_cache_key, in jax_compilation_cache_dir), and a hit is loaded with the callbacks of the present lowering, which binds them in the same order.

Active only with a persistent cache directory configured (jax_compilation_cache_dir, JAX_COMPILATION_CACHE_DIR) and within callback_cache() (the solvers’ march); opt out with ZOOMY_JAX_CALLBACK_CACHE=0. Any failure (another jax version, an unreadable entry) falls back to jax’s own compile.

References

• Bradbury et al. 2018, jax-ml/jax

Module Contents

Functions

callback_cache	Compiles inside this context also cache executables with host callbacks.
install	Wrap jax’s compiler.compile_or_get_cached once (no-op on an unknown signature).

Data

logger
stats

API

zoomy_jax.symsolver.compile_cache.logger

‘getLogger(…)’

zoomy_jax.symsolver.compile_cache.stats

‘dict(…)’

zoomy_jax.symsolver.compile_cache.callback_cache()

Compiles inside this context also cache executables with host callbacks.

zoomy_jax.symsolver.compile_cache.install()

Wrap jax’s compiler.compile_or_get_cached once (no-op on an unknown signature).