code: def capture_jax_logs():
      log_output = io.StringIO()
      handler = logging.StreamHandler(log_output)
      logger = logging.getLogger("jax")
  
      logger.addHandler(handler)
      try:
          yield log_output
      finally:
          logger.removeHandler(handler)
  
file_path: jax/tests/logging_test.py
imports: import contextlib
  import io
  import logging
  import platform
  import subprocess
  import sys
  import textwrap
  import unittest
  import jax
  from jax import config
  import jax._src.test_util
  from jax._src import xla_bridge
  from absl.testing import absltest
label_desc: Use image node_id 1 for calling a global function with example usage: capture_jax_logs() without return types