Add JAX compilation cache clear fixture to prevent XLA memory accumulation

[SurbhiJainUSC]

Description

This PR an autouse pytest fixture (clear_jax_compilation_cache) in tests/conftest.py that invokes jax.clear_caches() and gc.collect() after every test module execution to prevent XLA memory accumulation.

Key benefits:

• Prevents compilation cache buildup and XLA memory fragmentation across long-running test suites (especially in TPU and Pathways CI environments).
• Eliminates memory-related test flakiness and OOM failures.
• Cleans up redundant manual jax.clear_caches() and gc.collect() calls in tests/integration/estimator_test.py and tests/integration/deepseek_scan_engram_test.py.

Tests

CI tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a global autouse pytest fixture in tests/conftest.py to clear JAX compilation caches and force garbage collection, while removing redundant manual cache clearing and garbage collection calls from individual integration tests. The feedback points out that running this cleanup before and after every single test (including CPU-only tests) will severely degrade test suite performance due to constant recompilation. It is recommended to restrict this cleanup to run only after tests, and only when executing on TPU or GPU accelerators.

[github-actions]

🤖 Hi @SurbhiJainUSC, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request introduces an autouse pytest fixture (clear_jax_compilation_cache) to automatically clear JAX compilation caches and force garbage collection after each test on hardware accelerators (TPU, GPU, and proxy platforms). The objective is to prevent memory accumulation and fragmentation in CI/CD environments, and the overall quality and cleanliness of the implementation is excellent.

🔍 General Feedback

• Simplifies Test Codebase: Removing manual jax.clear_caches() and gc.collect() calls from individual test files (like estimator_test.py and deepseek_scan_engram_test.py) significantly reduces boilerplate and improves maintenance.
• Accelerator-Scoped Cleanups: Correctly restricts the automatic cache clearing to hardware accelerator environments, ensuring CPU-only tests or local CPU runs are unaffected.
• Compilation/Performance Overhead: Note that clearing JAX caches after every test can lead to increased compilation times if multiple test cases within the same session or module compile identical functions. However, this is a very reasonable trade-off for CI stability and memory reliability in resource-constrained TPU/Pathways environments.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!