jax.extend.backend.register_backend_cache#

jax.extend.backend.register_backend_cache(cache, for_what)[source]#

Registers a cache with JAX’s cache management.

Parameters:

  • cache (Any) – an object supporting cache_clear(), cache_info(), and cache_keys(), like the result of functools.lru_cache().

  • for_what (str) – a string to identify what this cache is used for. This is used for debugging.