//  [python]
def create_host_call(model_dir):
    """Construct a host_call writing scalar summaries.

  Args:
    model_dir: String containing path to train

  Returns:
    (fn, args) Pair to be called by TPUEstimator as the host_call.
  """
    graph = tf.get_default_graph()
    summaries = graph.get_collection(tf.GraphKeys.SUMMARIES)
    gs_t = tf.reshape(tf.cast(tf.train.get_global_step(), tf.int32), [1])
    summary_kwargs = collections.OrderedDict()
    for t in summaries:
        # TODO(aidangomez): enable ImageSummary support when we have a faster method
        # see @shibow's comment in cl/202344570
        if t.op.type not in ['ScalarSummary']:
            tf.logging.warn(
                'Ignoring unsupported tf.Summary type %s' % t.op.type
            )
            continue

        name = t.op.name
        tensor = t.op.inputs[1]
        if t.op.type == 'ScalarSummary':
            assert tensor.shape.is_compatible_with([])
            if tensor.dtype == tf.int64:
                tensor = tf.cast(tensor, tf.int32)
            summary_kwargs['ScalarSummary' + name] = tf.reshape(tensor, [1])
        elif t.op.type == 'ImageSummary':
            # TODO(aidangomez): as we move to support more types, update
            # common_layers.tpu_safe_image_summary
            if tensor.dtype != tf.float32:
                tf.logging.warn(
                    'Currently T2T on TPU only supports ImageSummary of '
                    'tf.float32-type Tensors. Skipping Tensor '
                    '%s with dtype %s...' % (tensor.name, tensor.dtype)
                )
                continue
            # tensor = tf.to_float(tensor)
            summary_kwargs['ImageSummary' + name] = tensor
    # When no supported summaries are found, don't create host_call. Otherwise,
    # TPU outfeed queue would enqueue global_step while host_call doesn't dequeue
    # it, eventually causing hang.
    if not summary_kwargs:
        return None
    summary_kwargs['global_step'] = gs_t
    log_info('summary_kwargs %s' % str(summary_kwargs))

    def host_call_fn(**kwargs):
        """Training host call. Creates summaries for training metrics.

    Args:
      **kwargs: Dict of {str: Tensor} , with `Tensor` of shape `[batch]`. Must
        contain key "global_step" with value of current global_step Tensor.

    Returns:
      List of summary ops to run on the CPU host.
    """
        gs = tf.to_int64(kwargs.pop('global_step')[0])
        with contrib.summary().create_file_writer(model_dir).as_default():
            with contrib.summary().always_record_summaries():
                # We need to use tf.contrib.summary in order to feed the `step`.
                for name, value in sorted(six.iteritems(kwargs)):
                    if name.startswith('ScalarSummary'):
                        name = name[len('ScalarSummary') :]
                        contrib.summary().scalar(
                            name, tf.reduce_mean(tf.cast(value, tf.float32)), step = gs
                        )
                    elif name.startswith('ImageSummary'):
                        name = name[len('ImageSummary') :]
                        contrib.summary().image(name, value, step = gs)

                return contrib.summary().all_summary_ops()

    return (host_call_fn, summary_kwargs)