Use capture_stderr instead of packing sys.stderr

This method is more reliable as it can capture stderr writes from the whole process
and not only those coming from Python.

I noticed this test was failing sporadically. I'm not sure if this will fix it,
but at least it might tell us why (thanks to e.g. the weakref check).

PiperOrigin-RevId: 707557909

Author: apaszke

tests/garbage_collection_guard_test.py

@@ -14,8 +14,7 @@
 """Tests for garbage allocation guard."""
 
 import gc
-import io
-from unittest import mock
+import weakref
 
 from absl.testing import absltest
 import jax
@@ -40,42 +39,36 @@ def _create_array_cycle():
   n2 = GarbageCollectionGuardTestNodeHelper(jax.jit(lambda: jnp.zeros((2, 2)))())
   n1.next = n2
   n2.next = n1
+  return weakref.ref(n1)
 
 
 class GarbageCollectionGuardTest(jtu.JaxTestCase):
 
   def test_gced_array_is_not_logged_by_default(self):
     # Create a reference cycle of two jax.Arrays.
-    _create_array_cycle()
-
-    # Use mock_stderr to be able to inspect stderr.
-    mock_stderr = io.StringIO()
-    with mock.patch("sys.stderr", mock_stderr):
-      # Trigger a garbage collection, which will garbage collect the arrays
-      # in the cycle.
+    ref = _create_array_cycle()
+    with jtu.capture_stderr() as stderr:
+      self.assertIsNotNone(ref())  # Cycle still alive.
       gc.collect()
+      self.assertIsNone(ref())  # Cycle collected.
     # Check that no error message is logged because
     # `array_garbage_collection_guard` defaults to `allow`.
     self.assertNotIn(
-        "`jax.Array` was deleted by the Python garbage collector",
-        mock_stderr.getvalue(),
+        "`jax.Array` was deleted by the Python garbage collector", stderr(),
     )
 
   def test_gced_array_is_logged(self):
-    # Use mock_stderr to be able to inspect stderr.
-    mock_stderr = io.StringIO()
-
     with config.array_garbage_collection_guard("log"):
-      # Create a reference cycle of two jax.Arrays.
-      _create_array_cycle()
-      with mock.patch("sys.stderr", mock_stderr):
+      with jtu.capture_stderr() as stderr:
+        # Create a reference cycle of two jax.Arrays.
+        ref = _create_array_cycle()
+        self.assertIsNotNone(ref())  # Cycle still alive.
         gc.collect()
-
+        self.assertIsNone(ref())  # Cycle collected.
     # Verify that an error message is logged because two jax.Arrays were garbage
     # collected.
     self.assertIn(
-        "`jax.Array` was deleted by the Python garbage collector",
-        mock_stderr.getvalue(),
+        "`jax.Array` was deleted by the Python garbage collector", stderr()
     )