code: def test_debug_logging(self):
      # Warmup so we don't get "No GPU/TPU" warning later.
      jax.jit(lambda x: x + 1)(1)
  
      # Nothing logged by default (except warning messages, which we don't expect
      # here).
      with capture_jax_logs() as log_output:
          jax.jit(lambda x: x + 1)(1)
      self.assertEmpty(log_output.getvalue())
  
      # Turn on all debug logging.
      config.update("jax_debug_log_modules", "jax")
      with capture_jax_logs() as log_output:
          jax.jit(lambda x: x + 1)(1)
      self.assertIn(
          "Finished tracing + transforming", log_output.getvalue()
      )
      self.assertIn("Compiling <lambda>", log_output.getvalue())
  
      # Turn off all debug logging.
      config.update("jax_debug_log_modules", None)
      with capture_jax_logs() as log_output:
          jax.jit(lambda x: x + 1)(1)
      self.assertEmpty(log_output.getvalue())
  
      # Turn on one module.
      config.update("jax_debug_log_modules", "jax._src.dispatch")
      with capture_jax_logs() as log_output:
          jax.jit(lambda x: x + 1)(1)
      self.assertIn(
          "Finished tracing + transforming", log_output.getvalue()
      )
      self.assertNotIn("Compiling <lambda>", log_output.getvalue())
  
      # Turn everything off again.
      config.update("jax_debug_log_modules", None)
      with capture_jax_logs() as log_output:
          jax.jit(lambda x: x + 1)(1)
      self.assertEmpty(log_output.getvalue())
  
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 2 for calling the LoggingTest obj's underlying member method code with example usage: obj.test_debug_logging() without return types