Don't use the trace context in the prune_closed_jaxpr_outputs cache.

This code only manipulates jaxprs, and does not trace anything. PiperOrigin-RevId: 648398046
skondla · Jul 1, 2024 · 55589fb · 55589fb
1 parent 9653f58
commit 55589fb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
@@ -1521,7 +1521,7 @@ def prune_closed_jaxpr_outputs(
 ) -> ClosedJaxpr:
   return _prune_closed_jaxpr_outputs(jaxpr, tuple(used_outputs))
 
-@weakref_lru_cache
+@partial(weakref_lru_cache, trace_context_in_key=False)
 def _prune_closed_jaxpr_outputs(
     jaxpr: ClosedJaxpr, used_outputs: tuple[bool, ...]
 ) -> ClosedJaxpr:

diff --git a/jax/_src/util.py b/jax/_src/util.py
@@ -317,7 +317,8 @@ def clear_all_caches():
 
 memoize = cache(max_size=None)
 
-def weakref_lru_cache(call: Callable, maxsize=2048):
+def weakref_lru_cache(call: Callable, maxsize=2048,
+                      trace_context_in_key: bool = True):
   """
   Least recently used cache decorator with weakref support.
 
@@ -326,7 +327,9 @@ def weakref_lru_cache(call: Callable, maxsize=2048):
   behave similar to `functools.lru_cache`.
   """
   global _weakref_lru_caches
-  cached_call = xc.weakref_lru_cache(config.trace_context, call, maxsize)
+  cached_call = xc.weakref_lru_cache(
+      config.trace_context if trace_context_in_key else _ignore,
+      call, maxsize)
   _weakref_lru_caches.add(cached_call)
   return cached_call