Add JAX compilation cache clear fixture to prevent XLA memory accumulation

.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"

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."""

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")

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