From f5b9c10197a4985d26e7ecc87aeb7b0f3528cc0e Mon Sep 17 00:00:00 2001
From: Surbhi Jain <sjsurbhi@google.com>
Date: Fri, 7 Aug 2026 20:35:58 +0000
Subject: [PATCH] Add JAX compilation cache clear fixture to prevent XLA memory
 accumulation

---
 .github/scripts/analyze_code_changes.sh        |  8 ++++++++
 tests/conftest.py                              | 18 ++++++++++++++++++
 tests/integration/deepseek_scan_engram_test.py |  3 ---
 tests/integration/estimator_test.py            |  2 --
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/.github/scripts/analyze_code_changes.sh b/.github/scripts/analyze_code_changes.sh
index 509eb339aa..bc6b2445b7 100755
--- a/.github/scripts/analyze_code_changes.sh
+++ b/.github/scripts/analyze_code_changes.sh
@@ -25,6 +25,7 @@ set -e
 # - Pathways only: Run Pathways tests only
 # - Post-training only: Run post-training tests and notebooks only
 # - General inference only: Run pretrain TPU/CPU suites only
+# - Tests only: Run test suites only (no notebooks)
 # - Default fallback: Run all suites (fail-open for core/shared code)
 
 # Helper to output a key-value flag to GITHUB_OUTPUT (if set) and stdout
@@ -137,6 +138,13 @@ if matches_only_domain 'src/maxtext/inference/|tests/inference/'; then
   exit 0
 fi
 
+# Tests and test-tooling only changes (skips notebooks as tutorials are unaffected)
+if matches_only_domain '(^tests/|^\.github/scripts/|^pytest\.ini$|^\.coveragerc$)'; then
+  echo "Only test files and test configurations changed, enabling test suites (skipping notebooks)."
+  enable_flags run_tests run_pretrain_tests run_posttrain_tests run_pathways_tests run_gpu_tests
+  exit 0
+fi
+
 # Default fallback: run all domain suites
 echo "General source changes detected, enabling all domain suites."
 set_all_flags "true"
diff --git a/tests/conftest.py b/tests/conftest.py
index 16c9cb2908..dfa8c70898 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -20,6 +20,7 @@
 are not marked.
 """
 
+import gc
 import pytest
 import sys
 import warnings
@@ -222,6 +223,23 @@ def pytest_configure(config):
     config.addinivalue_line("markers", m)
 
 
+@pytest.fixture(autouse=True, scope="module")
+def clear_jax_compilation_cache():
+  """Clears JAX compilation caches and forces garbage collection after each test module.
+
+  Prevents XLA memory accumulation and compilation cache buildup across
+  long test runs while preserving intra-module compilation cache reuse.
+  """
+  yield
+  try:
+    has_accelerator = any(d.platform in ("tpu", "gpu", "proxy") for d in jax.devices())
+    if has_accelerator:
+      jax.clear_caches()
+      gc.collect()
+  except Exception:  # pylint: disable=broad-exception-caught
+    pass
+
+
 @pytest.fixture(autouse=True)
 def handle_skip_on_tpu7x(request):
   """Dynamically skip tests marked with skip_on_tpu7x if running on TPU7x."""
diff --git a/tests/integration/deepseek_scan_engram_test.py b/tests/integration/deepseek_scan_engram_test.py
index 6891fd83e2..a9d0b12e92 100644
--- a/tests/integration/deepseek_scan_engram_test.py
+++ b/tests/integration/deepseek_scan_engram_test.py
@@ -14,7 +14,6 @@
 
 """Unit tests for DeepSeek Engram across scanned decoder layers."""
 
-import gc
 import os
 import unittest
 from unittest.mock import patch
@@ -164,8 +163,6 @@ def batch_decode(self, token_ids, *args, **kwargs):
     del variables
     del params
     del decoder
-    jax.clear_caches()
-    gc.collect()
 
   @pytest.mark.tpu_only
   @patch("transformers.AutoTokenizer.from_pretrained")
diff --git a/tests/integration/estimator_test.py b/tests/integration/estimator_test.py
index 53c2d3f8b6..c4998f45cf 100644
--- a/tests/integration/estimator_test.py
+++ b/tests/integration/estimator_test.py
@@ -96,7 +96,6 @@ def test_is_oom_returns_bool(self):
     ]
     policy = RematPolicy(tensor_names=tensor_names, initial_level=Action.REMAT)
 
-    jax.clear_caches()
     result = is_oom(base_argv, policy, pdb=2.0)
 
     self.assertIsInstance(result, bool)
@@ -120,7 +119,6 @@ def test_search_policy_only_small_model(self):
         "out_proj",
     ]
 
-    jax.clear_caches()
     result = search_policy_only(tensor_names, base_argv, pdb=2.0)
 
     # Should return a RematPolicy