code: def _as_default(
      summary_writer: tf.summary.SummaryWriter, auto_flush: bool
  ):
      """No-flush variation of summary_writer.as_default()."""
      context_manager = summary_writer.as_default()
      try:
          context_manager.__enter__()
          yield summary_writer
      finally:
          old_flush = summary_writer.flush
          new_flush = old_flush if auto_flush else lambda: None
          summary_writer.flush = new_flush
          context_manager.__exit__()
          summary_writer.flush = old_flush
  
file_path: flax/flax/metrics/tensorboard.py
imports: import contextlib
  import functools
  import os
  import numpy
  import tensorflow
  from tensorboard.plugins.hparams import api
  from flax import io
label_desc: Use image node_id 2 for calling a global function with example usage: _as_default(summary_writer, auto_flush) without return types