[jax] Add computation name to cache hit logging.

PiperOrigin-RevId: 414697336
jax-ml · Dec 7, 2021 · 56f029f · 56f029f
1 parent 803d3f2
commit 56f029f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
@@ -524,7 +524,8 @@ def compile_or_get_cached(backend, computation, compile_options):
     if cc.is_initialized() and backend.platform == 'tpu':
         cached_executable = cc.get_executable(computation, compile_options, backend)
         if cached_executable is not None:
-            logging.info('Persistent compilation cache hit')
+            logging.info('Persistent compilation cache hit for %s.',
+                         computation.name())
             return cached_executable
         else:
             compiled = backend_compile(backend, computation, compile_options)