jax.test_util: add capture_stdout context manager

jax/_src/test_util.py

@@ -14,12 +14,13 @@
 
 from contextlib import contextmanager
 import inspect
+import io
 import functools
 from functools import partial
 import re
 import os
 import textwrap
-from typing import Dict, List, Generator, Sequence, Tuple, Union
+from typing import Callable, Dict, List, Generator, Sequence, Tuple, Union
 import unittest
 import warnings
 import zlib
@@ -153,6 +154,14 @@ def check_eq(xs, ys, err_msg=''):
   tree_all(tree_map(assert_close, xs, ys))
 
 
+@contextmanager
+def capture_stdout() -> Generator[Callable[[], str], None, None]:
+  with unittest.mock.patch('sys.stdout', new_callable=io.StringIO) as fp:
+    def _read() -> str:
+      return fp.getvalue()
+    yield _read
+
+
 @contextmanager
 def count_device_put():
   device_put = dispatch.device_put