def rewire_summary_calls(use_tpu):
  """Rewires histogram() to nop and scalar() to tpu_summary_scalar.

  Summaries are not yet supported on TPUs (b/64980426). So we provide this
  context manager that can swap out tf.summary.histogram
  to be no-op and tf.summary.scalar to be tpu_summary_scalar. When the context
  exits, the rewiring is reset.
  See tpu_summary_scalar() for option to use.

  Args:
    use_tpu: Whether the model is being executed on TPUs.

  Yields:
    None.
  """
  # Store the original functions so we can put them back when the context ends.
  if use_tpu:
    original_tf_summary_scalar = tf.summary.scalar
    original_tf_summary_histogram = tf.summary.histogram
    original_summary_scalar = summary.scalar
    original_summary_histogram = summary.histogram

    tf.summary.scalar = summary.scalar = _tpu_summary_no_op
    tf.summary.histogram = summary.histogram = _tpu_summary_no_op
    try:
      yield
    finally:
      tf.summary.scalar = original_tf_summary_scalar
      tf.summary.histogram = original_tf_summary_histogram
      summary.scalar = original_summary_scalar
      summary.histogram = original_summary_histogram
  else:
    yield
Rewires histogram() to nop and scalar() to tpu_summary_scalar.

  Summaries are not yet supported on TPUs (b/64980426). So we provide this
  context manager that can swap out tf.summary.histogram
  to be no-op and tf.summary.scalar to be tpu_summary_scalar. When the context
  exits, the rewiring is reset.
  See tpu_summary_scalar() for option to use.

  Args:
    use_tpu: Whether the model is being executed on TPUs.

  Yields:
    None.
  
language/xsp/model/tpu_utils.py