Clear caches on jax exit.

PiperOrigin-RevId: 681529928
openxla · Oct 2, 2024 · 6b2237c · 6b2237c
1 parent c88b612
commit 6b2237c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/xla/python/pjit.cc b/xla/python/pjit.cc
@@ -135,7 +135,10 @@ class PjitFunctionCache {
 
   int Size() const { return lru_list_.Size(); }
   int Capacity() const { return lru_list_.Capacity(); }
-  void Clear() { lru_list_.Clear(); }
+  void Clear() {
+    lru_list_.Clear();
+    functions_.clear();
+  }
 
  private:
   struct Key {
@@ -347,6 +350,7 @@ class PjitFunctionStore {
     for (auto* function : compiled_functions_) {
       function->ClearCache();
     }
+    compiled_functions_.clear();
   }
 
  private: