remote call site

jax-ml · Dec 20, 2023 · aaa2487 · aaa2487
1 parent 963339f
commit aaa2487
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
@@ -903,7 +903,8 @@ def _process_in_axis_resources(in_shardings_thunk, in_layouts_thunk, in_avals,
 def explain_cache_miss(f: Callable, unseen_f: bool, cache: dict, key: tuple,
                        result: tuple):
   if config.check_tracer_leaks.value: return
-  print(f"TRACING CACHE MISS because:")
+  callsite = source_info_util.summarize(source_info_util.current())
+  print(f"TRACING CACHE MISS at {callsite} because:")
 
   def unpack(key):
     transforms, (), _, (in_type, debug_info, _), *_ = key