jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def bytenet_internal(inputs, targets, hparams):
  """ByteNet, main step used for training."""
  with tf.variable_scope("bytenet"):
    # Flatten inputs and extend length by 50%.
    inputs = tf.expand_dims(common_layers.flatten4d3d(inputs), axis=2)
    extend_length = tf.cast(0.5 * tf.to_float(tf.shape(inputs)[1]), tf.int32)
    inputs_shape = inputs.shape.as_list()
    inputs = tf.pad(inputs, [[0, 0], [0, extend_length], [0, 0], [0, 0]])
    inputs_shape[1] = None
    inputs.set_shape(inputs_shape)  # Don't lose the other shapes when padding.
    # Pad inputs and targets to be the same length, divisible by 50.
    inputs, targets = common_layers.pad_to_same_length(
        inputs, targets, final_length_divisible_by=50)
    final_encoder = residual_dilated_conv(inputs, hparams.num_block_repeat,
                                          "SAME", "encoder", hparams)

    shifted_targets = common_layers.shift_right(targets)
    kernel = (hparams.kernel_height, hparams.kernel_width)
    decoder_start = common_layers.conv_block(
        tf.concat([final_encoder, shifted_targets], axis=3),
        hparams.hidden_size, [((1, 1), kernel)],
        padding="LEFT")

    return residual_dilated_conv(decoder_start, hparams.num_block_repeat,
                                 "LEFT", "decoder", hparams)
//  [python]
def normalize_image(image, original_minval, original_maxval, target_minval,
                    target_maxval):
  """Normalizes pixel values in the image.

  Moves the pixel values from the current [original_minval, original_maxval]
  range to a the [target_minval, target_maxval] range.

  Args:
    image: rank 3 float32 tensor containing 1
           image -> [height, width, channels].
    original_minval: current image minimum value.
    original_maxval: current image maximum value.
    target_minval: target image minimum value.
    target_maxval: target image maximum value.

  Returns:
    image: image which is the same shape as input image.
  """
  with tf.name_scope('NormalizeImage', values=[image]):
    original_minval = float(original_minval)
    original_maxval = float(original_maxval)
    target_minval = float(target_minval)
    target_maxval = float(target_maxval)
    image = tf.to_float(image)
    image = tf.subtract(image, original_minval)
    image = tf.multiply(image, (target_maxval - target_minval) /
                        (original_maxval - original_minval))
    image = tf.add(image, target_minval)
    return image
//  [python]
def normalize_image(image, original_minval, original_maxval, target_minval,
                    target_maxval):
  """Normalizes pixel values in the image.

  Moves the pixel values from the current [original_minval, original_maxval]
  range to a the [target_minval, target_maxval] range.

  Args:
    image: rank 3 float32 tensor containing 1
           image -> [height, width, channels].
    original_minval: current image minimum value.
    original_maxval: current image maximum value.
    target_minval: target image minimum value.
    target_maxval: target image maximum value.

  Returns:
    image: image which is the same shape as input image.
  """
  with tf.name_scope('NormalizeImage', values=[image]):
    original_minval = float(original_minval)
    original_maxval = float(original_maxval)
    target_minval = float(target_minval)
    target_maxval = float(target_maxval)
    image = tf.cast(image, tf.float32)
    image = tf.subtract(image, original_minval)
    image = tf.multiply(image, (target_maxval - target_minval) /
                        (original_maxval - original_minval))
    image = tf.add(image, target_minval)
    return image
//  [python]
def image_to_float(image):
  """Used in Faster R-CNN. Casts image pixel values to float.

  Args:
    image: input image which might be in tf.uint8 or sth else format

  Returns:
    image: image in tf.float32 format.
  """
  with tf.name_scope('ImageToFloat', values=[image]):
    image = tf.to_float(image)
    return image
//  [python]
def image_to_float(image):
  """Used in Faster R-CNN. Casts image pixel values to float.

  Args:
    image: input image which might be in tf.uint8 or sth else format

  Returns:
    image: image in tf.float32 format.
  """
  with tf.name_scope('ImageToFloat', values=[image]):
    image = tf.cast(image, tf.float32)
    return image
//  [python]
def heatmap_from_uvd_op(uvd_pts, cfg, gaussian_filter):
    '''we firstly construct a sparse tensor from the coordinate
    val: the value at the center of corresponding point
    '''
    with tf.name_scope('preprocess'):
        uvd_pts = tf.reshape(uvd_pts, (-1,3))
        num_pt = uvd_pts.shape[0]
        num_pt_op = tf.to_int64(num_pt)

        nn = tf.range(num_pt, dtype=tf.int64)
        nn = tf.reshape(nn, (-1,1))

        xx = uvd_pts[:,0]
        xx = tf.clip_by_value(xx, 0, cfg.w-1)
        xx = tf.to_int64(xx)
        xx = tf.reshape(xx, (-1,1))

        yy = uvd_pts[:,1]
        yy = tf.clip_by_value(yy, 0, cfg.h-1)
        yy = tf.to_int64(yy)
        yy = tf.reshape(yy, (-1,1))
        indices = tf.concat([nn,yy,xx], axis=1)

        val = 1.0
        raw_hm = tf.sparse_to_dense(sparse_indices=indices,
                                    output_shape=[num_pt_op,cfg.h,cfg.w],
                                    sparse_values=val)
        raw_hm = tf.expand_dims(raw_hm, axis=[-1])
        raw_hm = tf.cast(raw_hm, tf.float32)
        
        hm = tf.nn.conv2d(raw_hm, gaussian_filter, strides=[1,1,1,1], 
                         padding='SAME', data_format='NHWC')
        hm = tf.nn.conv2d(hm, gaussian_filter, strides=[1,1,1,1], 
                         padding='SAME', data_format='NHWC')
        hm = tf.divide(hm, tf.reduce_max(hm))
        
        # shuffle dimensions of hm
        hm_list = tf.unstack(hm, axis=0)
        hm = tf.concat(hm_list, axis=2)
        return hm
//  [python]
def heatmap_from_uvd_op(uvd_pts, cfg, gaussian_filter):
    '''we firstly construct a sparse tensor from the coordinate
    val: the value at the center of corresponding point
    '''
    with tf.name_scope('preprocess'):
        uvd_pts = tf.reshape(uvd_pts, (-1,3))
        num_pt = uvd_pts.shape[0]
        num_pt_op = tf.cast(num_pt, tf.int64)

        nn = tf.range(num_pt, dtype=tf.int64)
        nn = tf.reshape(nn, (-1,1))

        xx = uvd_pts[:,0]
        xx = tf.clip_by_value(xx, 0, cfg.w-1)
        xx = tf.cast(xx, tf.int64)
        xx = tf.reshape(xx, (-1,1))

        yy = uvd_pts[:,1]
        yy = tf.clip_by_value(yy, 0, cfg.h-1)
        yy = tf.cast(yy, tf.int64)
        yy = tf.reshape(yy, (-1,1))
        indices = tf.concat([nn,yy,xx], axis=1)

        val = 1.0
        raw_hm = tf.sparse_to_dense(sparse_indices=indices,
                                    output_shape=[num_pt_op,cfg.h,cfg.w],
                                    sparse_values=val)
        raw_hm = tf.expand_dims(raw_hm, axis=[-1])
        raw_hm = tf.cast(raw_hm, tf.float32)
        
        hm = tf.nn.conv2d(raw_hm, gaussian_filter, strides=[1,1,1,1], 
                         padding='SAME', data_format='NHWC')
        hm = tf.nn.conv2d(hm, gaussian_filter, strides=[1,1,1,1], 
                         padding='SAME', data_format='NHWC')
        hm = tf.divide(hm, tf.reduce_max(hm))
        
        # shuffle dimensions of hm
        hm_list = tf.unstack(hm, axis=0)
        hm = tf.concat(hm_list, axis=2)
        return hm
//  [python]
def mentornet_nn(input_features,
                 label_embedding_size=2,
                 epoch_embedding_size=5,
                 num_label_embedding=2,
                 num_fc_nodes=20):
  """The neural network form of the MentorNet.

  An implementation of the mentornet. The details are in:
  Jiang, Lu, et al. "MentorNet: Learning Data-Driven Curriculum for Very Deep
  Neural Networks on Corrupted Labels." ICML. 2018.

  Args:
    input_features: a [batch_size, 4] tensor. Each dimension corresponds to
      0: loss, 1: loss difference to the moving average, 2: label and 3: epoch,
      where epoch is an integer between 0 and 99 (the first and the last epoch).
    label_embedding_size: the embedding size for the label feature.
    epoch_embedding_size: the embedding size for the epoch feature.
    num_label_embedding: the number of different labels.
    num_fc_nodes: number of hidden nodes in the fc layer.
  Returns:
    v: [batch_size, 1] weight vector.
  """
  batch_size = int(input_features.get_shape()[0])

  losses = tf.reshape(input_features[:, 0], [-1, 1])
  loss_diffs = tf.reshape(input_features[:, 1], [-1, 1])
  labels = tf.to_int32(tf.reshape(input_features[:, 2], [-1, 1]))
  epochs = tf.to_int32(tf.reshape(input_features[:, 3], [-1, 1]))
  epochs = tf.minimum(epochs, tf.ones([batch_size, 1], dtype=tf.int32) * 99)

  if len(losses.get_shape()) <= 1:
    num_steps = 1
  else:
    num_steps = int(losses.get_shape()[1])

  with tf.variable_scope('mentornet', reuse=tf.AUTO_REUSE):
    label_embedding = tf.get_variable(
        'label_embedding', [num_label_embedding, label_embedding_size])
    epoch_embedding = tf.get_variable(
        'epoch_embedding', [100, epoch_embedding_size])

    lstm_inputs = tf.stack([losses, loss_diffs], axis=1)
    lstm_inputs = tf.squeeze(lstm_inputs)
    lstm_inputs = [lstm_inputs]

    forward_cell = tf.contrib.rnn.BasicLSTMCell(1, forget_bias=0.0)
    backward_cell = tf.contrib.rnn.BasicLSTMCell(1, forget_bias=0.0)

    _, out_state_fw, out_state_bw = tf.contrib.rnn.static_bidirectional_rnn(
        forward_cell,
        backward_cell,
        inputs=lstm_inputs,
        dtype=tf.float32,
        sequence_length=np.ones(batch_size) * num_steps)

    label_inputs = tf.squeeze(tf.nn.embedding_lookup(label_embedding, labels))
    epoch_inputs = tf.squeeze(tf.nn.embedding_lookup(epoch_embedding, epochs))

    h = tf.concat([out_state_fw[0], out_state_bw[0]], 1)
    feat = tf.concat([label_inputs, epoch_inputs, h], 1)
    feat_dim = int(feat.get_shape()[1])

    fc_1 = tf.add(
        tf.matmul(
            feat,
            tf.get_variable(
                'Variable',
                initializer=tf.random_normal([feat_dim, num_fc_nodes]))),
        tf.get_variable(
            'Variable_1', initializer=tf.random_normal([num_fc_nodes])))
    fc_1 = tf.nn.tanh(fc_1)
    # Output layer with linear activation
    out_layer = tf.matmul(
        fc_1,
        tf.get_variable(
            'Variable_2', initializer=tf.random_normal([num_fc_nodes, 1])) +
        tf.get_variable('Variable_3', initializer=tf.random_normal([1])))
    return out_layer
//  [python]
def mentornet_nn(input_features,
                 label_embedding_size=2,
                 epoch_embedding_size=5,
                 num_label_embedding=2,
                 num_fc_nodes=20):
  """The neural network form of the MentorNet.

  An implementation of the mentornet. The details are in:
  Jiang, Lu, et al. "MentorNet: Learning Data-Driven Curriculum for Very Deep
  Neural Networks on Corrupted Labels." ICML. 2018.

  Args:
    input_features: a [batch_size, 4] tensor. Each dimension corresponds to
      0: loss, 1: loss difference to the moving average, 2: label and 3: epoch,
      where epoch is an integer between 0 and 99 (the first and the last epoch).
    label_embedding_size: the embedding size for the label feature.
    epoch_embedding_size: the embedding size for the epoch feature.
    num_label_embedding: the number of different labels.
    num_fc_nodes: number of hidden nodes in the fc layer.
  Returns:
    v: [batch_size, 1] weight vector.
  """
  batch_size = int(input_features.get_shape()[0])

  losses = tf.reshape(input_features[:, 0], [-1, 1])
  loss_diffs = tf.reshape(input_features[:, 1], [-1, 1])
  labels = tf.cast(tf.reshape(input_features[:, 2], [-1, 1]), tf.int32)
  epochs = tf.cast(tf.reshape(input_features[:, 3], [-1, 1]), tf.int32)
  epochs = tf.minimum(epochs, tf.ones([batch_size, 1], dtype=tf.int32) * 99)

  if len(losses.get_shape()) <= 1:
    num_steps = 1
  else:
    num_steps = int(losses.get_shape()[1])

  with tf.variable_scope('mentornet', reuse=tf.AUTO_REUSE):
    label_embedding = tf.get_variable(
        'label_embedding', [num_label_embedding, label_embedding_size])
    epoch_embedding = tf.get_variable(
        'epoch_embedding', [100, epoch_embedding_size])

    lstm_inputs = tf.stack([losses, loss_diffs], axis=1)
    lstm_inputs = tf.squeeze(lstm_inputs)
    lstm_inputs = [lstm_inputs]

    forward_cell = tf.contrib.rnn.BasicLSTMCell(1, forget_bias=0.0)
    backward_cell = tf.contrib.rnn.BasicLSTMCell(1, forget_bias=0.0)

    _, out_state_fw, out_state_bw = tf.contrib.rnn.static_bidirectional_rnn(
        forward_cell,
        backward_cell,
        inputs=lstm_inputs,
        dtype=tf.float32,
        sequence_length=np.ones(batch_size) * num_steps)

    label_inputs = tf.squeeze(tf.nn.embedding_lookup(label_embedding, labels))
    epoch_inputs = tf.squeeze(tf.nn.embedding_lookup(epoch_embedding, epochs))

    h = tf.concat([out_state_fw[0], out_state_bw[0]], 1)
    feat = tf.concat([label_inputs, epoch_inputs, h], 1)
    feat_dim = int(feat.get_shape()[1])

    fc_1 = tf.add(
        tf.matmul(
            feat,
            tf.get_variable(
                'Variable',
                initializer=tf.random_normal([feat_dim, num_fc_nodes]))),
        tf.get_variable(
            'Variable_1', initializer=tf.random_normal([num_fc_nodes])))
    fc_1 = tf.nn.tanh(fc_1)
    # Output layer with linear activation
    out_layer = tf.matmul(
        fc_1,
        tf.get_variable(
            'Variable_2', initializer=tf.random_normal([num_fc_nodes, 1])) +
        tf.get_variable('Variable_3', initializer=tf.random_normal([1])))
    return out_layer
//  [python]
def random_image_scale(image,
                       masks=None,
                       min_scale_ratio=0.5,
                       max_scale_ratio=2.0,
                       seed=None,
                       preprocess_vars_cache=None):
  """Scales the image size.

  Args:
    image: rank 3 float32 tensor contains 1 image -> [height, width, channels].
    masks: (optional) rank 3 float32 tensor containing masks with
      size [height, width, num_masks]. The value is set to None if there are no
      masks.
    min_scale_ratio: minimum scaling ratio.
    max_scale_ratio: maximum scaling ratio.
    seed: random seed.
    preprocess_vars_cache: PreprocessorCache object that records previously
                           performed augmentations. Updated in-place. If this
                           function is called multiple times with the same
                           non-null cache, it will perform deterministically.

  Returns:
    image: image which is the same rank as input image.
    masks: If masks is not none, resized masks which are the same rank as input
      masks will be returned.
  """
  with tf.name_scope('RandomImageScale', values=[image]):
    result = []
    image_shape = tf.shape(image)
    image_height = image_shape[0]
    image_width = image_shape[1]
    generator_func = functools.partial(
        tf.random_uniform, [],
        minval=min_scale_ratio, maxval=max_scale_ratio,
        dtype=tf.float32, seed=seed)
    size_coef = _get_or_create_preprocess_rand_vars(
        generator_func, preprocessor_cache.PreprocessorCache.IMAGE_SCALE,
        preprocess_vars_cache)

    image_newysize = tf.to_int32(
        tf.multiply(tf.to_float(image_height), size_coef))
    image_newxsize = tf.to_int32(
        tf.multiply(tf.to_float(image_width), size_coef))
    image = tf.image.resize_images(
        image, [image_newysize, image_newxsize], align_corners=True)
    result.append(image)
    if masks is not None:
      masks = tf.image.resize_images(
          masks, [image_newysize, image_newxsize],
          method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
          align_corners=True)
      result.append(masks)
    return tuple(result)
//  [python]
def random_image_scale(image,
                       masks=None,
                       min_scale_ratio=0.5,
                       max_scale_ratio=2.0,
                       seed=None,
                       preprocess_vars_cache=None):
  """Scales the image size.

  Args:
    image: rank 3 float32 tensor contains 1 image -> [height, width, channels].
    masks: (optional) rank 3 float32 tensor containing masks with
      size [height, width, num_masks]. The value is set to None if there are no
      masks.
    min_scale_ratio: minimum scaling ratio.
    max_scale_ratio: maximum scaling ratio.
    seed: random seed.
    preprocess_vars_cache: PreprocessorCache object that records previously
                           performed augmentations. Updated in-place. If this
                           function is called multiple times with the same
                           non-null cache, it will perform deterministically.

  Returns:
    image: image which is the same rank as input image.
    masks: If masks is not none, resized masks which are the same rank as input
      masks will be returned.
  """
  with tf.name_scope('RandomImageScale', values=[image]):
    result = []
    image_shape = tf.shape(image)
    image_height = image_shape[0]
    image_width = image_shape[1]
    generator_func = functools.partial(
        tf.random_uniform, [],
        minval=min_scale_ratio, maxval=max_scale_ratio,
        dtype=tf.float32, seed=seed)
    size_coef = _get_or_create_preprocess_rand_vars(
        generator_func, preprocessor_cache.PreprocessorCache.IMAGE_SCALE,
        preprocess_vars_cache)

    image_newysize = tf.cast(
        tf.multiply(tf.to_float(image_height), size_coef), tf.int32)
    image_newxsize = tf.cast(
        tf.multiply(tf.to_float(image_width), size_coef), tf.int32)
    image = tf.image.resize_images(
        image, [image_newysize, image_newxsize], align_corners=True)
    result.append(image)
    if masks is not None:
      masks = tf.image.resize_images(
          masks, [image_newysize, image_newxsize],
          method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
          align_corners=True)
      result.append(masks)
    return tuple(result)
//  [python]
def _compute_clip_window(self, image_shapes):
    """Computes clip window for non max suppression based on image shapes.

    This function assumes that the clip window's left top corner is at (0, 0).

    Args:
      image_shapes: A 2-D int32 tensor of shape [batch_size, 3] containing
      shapes of images in the batch. Each row represents [height, width,
      channels] of an image.

    Returns:
      A 2-D float32 tensor of shape [batch_size, 4] containing the clip window
      for each image in the form [ymin, xmin, ymax, xmax].
    """
    clip_heights = image_shapes[:, 0]
    clip_widths = image_shapes[:, 1]
    clip_window = tf.to_float(tf.stack([tf.zeros_like(clip_heights),
                                        tf.zeros_like(clip_heights),
                                        clip_heights, clip_widths], axis=1))
    return clip_window
//  [python]
def _compute_clip_window(self, image_shapes):
    """Computes clip window for non max suppression based on image shapes.

    This function assumes that the clip window's left top corner is at (0, 0).

    Args:
      image_shapes: A 2-D int32 tensor of shape [batch_size, 3] containing
      shapes of images in the batch. Each row represents [height, width,
      channels] of an image.

    Returns:
      A 2-D float32 tensor of shape [batch_size, 4] containing the clip window
      for each image in the form [ymin, xmin, ymax, xmax].
    """
    clip_heights = image_shapes[:, 0]
    clip_widths = image_shapes[:, 1]
    clip_window = tf.cast(tf.stack([tf.zeros_like(clip_heights),
                                        tf.zeros_like(clip_heights),
                                        clip_heights, clip_widths], axis=1), tf.float32)
    return clip_window
//  [python]
def _gather_instance_masks(self, instance_masks, classes):
    """Gathers the masks that correspond to classes.

    Args:
      instance_masks: A 4-D float32 tensor with shape
        [K, num_classes, mask_height, mask_width].
      classes: A 2-D int32 tensor with shape [batch_size, max_detection].

    Returns:
      masks: a 3-D float32 tensor with shape [K, mask_height, mask_width].
    """
    _, num_classes, height, width = instance_masks.get_shape().as_list()
    k = tf.shape(instance_masks)[0]
    instance_masks = tf.reshape(instance_masks, [-1, height, width])
    classes = tf.to_int32(tf.reshape(classes, [-1]))
    gather_idx = tf.range(k) * num_classes + classes
    return tf.gather(instance_masks, gather_idx)
//  [python]
def _gather_instance_masks(self, instance_masks, classes):
    """Gathers the masks that correspond to classes.

    Args:
      instance_masks: A 4-D float32 tensor with shape
        [K, num_classes, mask_height, mask_width].
      classes: A 2-D int32 tensor with shape [batch_size, max_detection].

    Returns:
      masks: a 3-D float32 tensor with shape [K, mask_height, mask_width].
    """
    _, num_classes, height, width = instance_masks.get_shape().as_list()
    k = tf.shape(instance_masks)[0]
    instance_masks = tf.reshape(instance_masks, [-1, height, width])
    classes = tf.cast(tf.reshape(classes, [-1]), tf.int32)
    gather_idx = tf.range(k) * num_classes + classes
    return tf.gather(instance_masks, gather_idx)
//  [python]
def train(model_config,
          input_embeddings,
          source_len,
          output_vocab_size,
          output_vocab_embeddings_table,
          target_decode_steps,
          mode,
          input_copy_mask=None):
  """Constructs encoder and decoder transformation for training and eval.

  In the shapes described below, B is batch size, L is sequence length,
  D is the dimensionality of the model embeddings, and T is the output vocab
  size.

  Args:
    model_config: ModelConfig proto.
    input_embeddings: Tensor of shape (B, L, D) representing inputs.
    source_len: Tensor of shape (B) containing length of each input sequence.
    output_vocab_size: Size of output vocabulary.
    output_vocab_embeddings_table: Tensor of shape (T, D) representing table of
      embeddings for output symbols.
    target_decode_steps: DecodeSteps representing target outputs. Each tensor
      has shape (B, L).
    mode: Enum indicating model mode, TRAIN or EVAL.
    input_copy_mask: Mask for copying actions.

  Returns:
    Tuple of (logits, predicted_ids), where logits is a tensor of shape
    (B, L, T) representing model output logits, and predicted_ids is
    a tensor of shape (B, L) containing the derived integer IDs of the
    one-best output symbol.
  """
  logits = _transformer_body(
      input_embeddings,
      source_len,
      target_decode_steps,
      mode,
      model_config,
      output_vocab_size,
      output_vocab_embeddings_table,
      input_copy_mask=input_copy_mask)

  predicted_ids = tf.to_int32(tf.argmax(logits, axis=-1))
  output_decode_steps = decode_utils.get_decode_steps(predicted_ids,
                                                      output_vocab_size,
                                                      model_config)
  predictions = decode_utils.get_predictions(output_decode_steps)
  return logits, predictions
//  [python]
def train(model_config,
          input_embeddings,
          source_len,
          output_vocab_size,
          output_vocab_embeddings_table,
          target_decode_steps,
          mode,
          input_copy_mask=None):
  """Constructs encoder and decoder transformation for training and eval.

  In the shapes described below, B is batch size, L is sequence length,
  D is the dimensionality of the model embeddings, and T is the output vocab
  size.

  Args:
    model_config: ModelConfig proto.
    input_embeddings: Tensor of shape (B, L, D) representing inputs.
    source_len: Tensor of shape (B) containing length of each input sequence.
    output_vocab_size: Size of output vocabulary.
    output_vocab_embeddings_table: Tensor of shape (T, D) representing table of
      embeddings for output symbols.
    target_decode_steps: DecodeSteps representing target outputs. Each tensor
      has shape (B, L).
    mode: Enum indicating model mode, TRAIN or EVAL.
    input_copy_mask: Mask for copying actions.

  Returns:
    Tuple of (logits, predicted_ids), where logits is a tensor of shape
    (B, L, T) representing model output logits, and predicted_ids is
    a tensor of shape (B, L) containing the derived integer IDs of the
    one-best output symbol.
  """
  logits = _transformer_body(
      input_embeddings,
      source_len,
      target_decode_steps,
      mode,
      model_config,
      output_vocab_size,
      output_vocab_embeddings_table,
      input_copy_mask=input_copy_mask)

  predicted_ids = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
  output_decode_steps = decode_utils.get_decode_steps(predicted_ids,
                                                      output_vocab_size,
                                                      model_config)
  predictions = decode_utils.get_predictions(output_decode_steps)
  return logits, predictions
//  [python]
def resize_to_min_dimension(image, masks=None, min_dimension=600):
  """Resizes image and masks given the min size maintaining the aspect ratio.

  If one of the image dimensions is smaller that min_dimension, it will scale
  the image such that its smallest dimension is equal to min_dimension.
  Otherwise, will keep the image size as is.

  Args:
    image: a tensor of size [height, width, channels].
    masks: (optional) a tensors of size [num_instances, height, width].
    min_dimension: minimum image dimension.

  Returns:
    Note that the position of the resized_image_shape changes based on whether
    masks are present.
    resized_image: A tensor of size [new_height, new_width, channels].
    resized_masks: If masks is not None, also outputs masks. A 3D tensor of
      shape [num_instances, new_height, new_width]
    resized_image_shape: A 1D tensor of shape [3] containing the shape of the
      resized image.

  Raises:
    ValueError: if the image is not a 3D tensor.
  """
  if len(image.get_shape()) != 3:
    raise ValueError('Image should be 3D tensor')

  with tf.name_scope('ResizeGivenMinDimension', values=[image, min_dimension]):
    image_height = tf.shape(image)[0]
    image_width = tf.shape(image)[1]
    num_channels = tf.shape(image)[2]
    min_image_dimension = tf.minimum(image_height, image_width)
    min_target_dimension = tf.maximum(min_image_dimension, min_dimension)
    target_ratio = tf.to_float(min_target_dimension) / tf.to_float(
        min_image_dimension)
    target_height = tf.to_int32(tf.to_float(image_height) * target_ratio)
    target_width = tf.to_int32(tf.to_float(image_width) * target_ratio)
    image = tf.image.resize_bilinear(
        tf.expand_dims(image, axis=0),
        size=[target_height, target_width],
        align_corners=True)
    result = [tf.squeeze(image, axis=0)]

    if masks is not None:
      masks = tf.image.resize_nearest_neighbor(
          tf.expand_dims(masks, axis=3),
          size=[target_height, target_width],
          align_corners=True)
      result.append(tf.squeeze(masks, axis=3))

    result.append(tf.stack([target_height, target_width, num_channels]))
    return result
//  [python]
def resize_to_min_dimension(image, masks=None, min_dimension=600):
  """Resizes image and masks given the min size maintaining the aspect ratio.

  If one of the image dimensions is smaller that min_dimension, it will scale
  the image such that its smallest dimension is equal to min_dimension.
  Otherwise, will keep the image size as is.

  Args:
    image: a tensor of size [height, width, channels].
    masks: (optional) a tensors of size [num_instances, height, width].
    min_dimension: minimum image dimension.

  Returns:
    Note that the position of the resized_image_shape changes based on whether
    masks are present.
    resized_image: A tensor of size [new_height, new_width, channels].
    resized_masks: If masks is not None, also outputs masks. A 3D tensor of
      shape [num_instances, new_height, new_width]
    resized_image_shape: A 1D tensor of shape [3] containing the shape of the
      resized image.

  Raises:
    ValueError: if the image is not a 3D tensor.
  """
  if len(image.get_shape()) != 3:
    raise ValueError('Image should be 3D tensor')

  with tf.name_scope('ResizeGivenMinDimension', values=[image, min_dimension]):
    image_height = tf.shape(image)[0]
    image_width = tf.shape(image)[1]
    num_channels = tf.shape(image)[2]
    min_image_dimension = tf.minimum(image_height, image_width)
    min_target_dimension = tf.maximum(min_image_dimension, min_dimension)
    target_ratio = tf.to_float(min_target_dimension) / tf.to_float(
        min_image_dimension)
    target_height = tf.cast(tf.to_float(image_height) * target_ratio, tf.int32)
    target_width = tf.cast(tf.to_float(image_width) * target_ratio, tf.int32)
    image = tf.image.resize_bilinear(
        tf.expand_dims(image, axis=0),
        size=[target_height, target_width],
        align_corners=True)
    result = [tf.squeeze(image, axis=0)]

    if masks is not None:
      masks = tf.image.resize_nearest_neighbor(
          tf.expand_dims(masks, axis=3),
          size=[target_height, target_width],
          align_corners=True)
      result.append(tf.squeeze(masks, axis=3))

    result.append(tf.stack([target_height, target_width, num_channels]))
    return result
//  [python]
def alltoall(self, x, mesh_axis, split_axis, concat_axis):
    """Grouped alltoall (like MPI alltoall with splitting and concatenation).

    Args:
      x: a LaidOutTensor
      mesh_axis: an integer the mesh axis along which to group
      split_axis: an integer (the Tensor axis along which to split)
      concat_axis: an integer (the Tensor axis along which to concatenate)
    Returns:
      a LaidOutTensor
    """
    x = x.to_laid_out_tensor()
    t = x.one_slice
    group_assignment = self._create_group_assignment([mesh_axis])
    dtype = t.dtype
    if dtype == tf.float32:
      # There seems to be a bug with float32 alltoall.
      # Do it in bfloat16 until the bug is fixed.
      # TODO(noam): file a bug
      t = tf.to_bfloat16(t)
    t = tpu_ops.all_to_all(
        t,
        concat_dimension=concat_axis,
        split_dimension=split_axis,
        split_count=len(group_assignment[0]),
        group_assignment=group_assignment)
    t = tf.cast(t, dtype)
    x = self.LaidOutTensor([t])
    return x
//  [python]
def alltoall(self, x, mesh_axis, split_axis, concat_axis):
    """Grouped alltoall (like MPI alltoall with splitting and concatenation).

    Args:
      x: a LaidOutTensor
      mesh_axis: an integer the mesh axis along which to group
      split_axis: an integer (the Tensor axis along which to split)
      concat_axis: an integer (the Tensor axis along which to concatenate)
    Returns:
      a LaidOutTensor
    """
    x = x.to_laid_out_tensor()
    t = x.one_slice
    group_assignment = self._create_group_assignment([mesh_axis])
    dtype = t.dtype
    if dtype == tf.float32:
      # There seems to be a bug with float32 alltoall.
      # Do it in bfloat16 until the bug is fixed.
      # TODO(noam): file a bug
      t = tf.cast(t, tf.bfloat16)
    t = tpu_ops.all_to_all(
        t,
        concat_dimension=concat_axis,
        split_dimension=split_axis,
        split_count=len(group_assignment[0]),
        group_assignment=group_assignment)
    t = tf.cast(t, dtype)
    x = self.LaidOutTensor([t])
    return x
//  [python]
def main(unused_argv):
  tf.logging.set_verbosity(tf.logging.INFO)
  style_files = _parse_style_files(os.path.expanduser(FLAGS.style_files))
  with tf.python_io.TFRecordWriter(
      os.path.expanduser(FLAGS.output_file)) as writer:
    for style_label, style_file in enumerate(style_files):
      tf.logging.info(
          'Processing style file %s: %s' % (style_label, style_file))
      feature = {'label': _int64_feature(style_label)}

      style_image = image_utils.load_np_image(style_file)
      buf = io.BytesIO()
      skimage.io.imsave(buf, style_image, format='JPEG')
      buf.seek(0)
      feature['image_raw'] = _bytes_feature(buf.getvalue())

      if FLAGS.compute_gram_matrices:
        with tf.Graph().as_default():
          style_end_points = learning.precompute_gram_matrices(
              tf.expand_dims(tf.to_float(style_image), 0),
              # We use 'pool5' instead of 'fc8' because a) fully-connected
              # layers are already too deep in the network to be useful for
              # style and b) they're quite expensive to store.
              final_endpoint='pool5')
          for name in style_end_points:
            feature[name] = _float_feature(
                style_end_points[name].flatten().tolist())

      example = tf.train.Example(features=tf.train.Features(feature=feature))
      writer.write(example.SerializeToString())
  tf.logging.info('Output TFRecord file is saved at %s' % os.path.expanduser(
      FLAGS.output_file))
//  [python]
def main(unused_argv):
  tf.logging.set_verbosity(tf.logging.INFO)
  style_files = _parse_style_files(os.path.expanduser(FLAGS.style_files))
  with tf.python_io.TFRecordWriter(
      os.path.expanduser(FLAGS.output_file)) as writer:
    for style_label, style_file in enumerate(style_files):
      tf.logging.info(
          'Processing style file %s: %s' % (style_label, style_file))
      feature = {'label': _int64_feature(style_label)}

      style_image = image_utils.load_np_image(style_file)
      buf = io.BytesIO()
      skimage.io.imsave(buf, style_image, format='JPEG')
      buf.seek(0)
      feature['image_raw'] = _bytes_feature(buf.getvalue())

      if FLAGS.compute_gram_matrices:
        with tf.Graph().as_default():
          style_end_points = learning.precompute_gram_matrices(
              tf.expand_dims(tf.cast(style_image, tf.float32), 0),
              # We use 'pool5' instead of 'fc8' because a) fully-connected
              # layers are already too deep in the network to be useful for
              # style and b) they're quite expensive to store.
              final_endpoint='pool5')
          for name in style_end_points:
            feature[name] = _float_feature(
                style_end_points[name].flatten().tolist())

      example = tf.train.Example(features=tf.train.Features(feature=feature))
      writer.write(example.SerializeToString())
  tf.logging.info('Output TFRecord file is saved at %s' % os.path.expanduser(
      FLAGS.output_file))
//  [python]
def attention_bias_lower_triangle(length):
  """Create a bias tensor to be added to attention logits.

  Masked elements will have an effective negative-infinity value.

  Args:
    length: Integer specifying maximum sequence length in batch.

  Returns:
    <float>[length, length]
  """
  # First, create a sequence mask, e.g.:
  # [1, 0, ..., 0]
  # ...
  # [1, 1, ..., 1]
  sequence_mask = tf.sequence_mask(tf.range(1, length + 1), length)
  # Invert to transform to attention biases.
  attention_bias = tf.to_float(tf.logical_not(sequence_mask)) * -1e9
  return attention_bias
//  [python]
def attention_bias_lower_triangle(length):
  """Create a bias tensor to be added to attention logits.

  Masked elements will have an effective negative-infinity value.

  Args:
    length: Integer specifying maximum sequence length in batch.

  Returns:
    <float>[length, length]
  """
  # First, create a sequence mask, e.g.:
  # [1, 0, ..., 0]
  # ...
  # [1, 1, ..., 1]
  sequence_mask = tf.sequence_mask(tf.range(1, length + 1), length)
  # Invert to transform to attention biases.
  attention_bias = tf.cast(tf.logical_not(sequence_mask), dtype=tf.float32) * -1e9
  return attention_bias
//  [python]
def attention_bias_ignore_padding(source_len, max_length):
  """Create a bias tensor to be added to attention logits.

  Out-of-range elements will have an effective negative-infinity value.

  Args:
    source_len: <int>[batch_size]
    max_length: Integer specifying maxmimum sequence length in batch.

  Returns:
    <float>[batch_size, 1, 1, max_length]
  """
  memory_padding = tf.to_float(
      tf.logical_not(tf.sequence_mask(source_len, maxlen=max_length)))
  ret = memory_padding * -1e9
  # Expand so tensor can be broadcast across heads and query length.
  return tf.expand_dims(tf.expand_dims(ret, axis=1), axis=1)
//  [python]
def attention_bias_ignore_padding(source_len, max_length):
  """Create a bias tensor to be added to attention logits.

  Out-of-range elements will have an effective negative-infinity value.

  Args:
    source_len: <int>[batch_size]
    max_length: Integer specifying maxmimum sequence length in batch.

  Returns:
    <float>[batch_size, 1, 1, max_length]
  """
  memory_padding = tf.cast(
      tf.logical_not(tf.sequence_mask(source_len, maxlen=max_length)), dtype=tf.float32)
  ret = memory_padding * -1e9
  # Expand so tensor can be broadcast across heads and query length.
  return tf.expand_dims(tf.expand_dims(ret, axis=1), axis=1)
//  [python]
def ndcg(labels, predictions,
                  metrics_collections=None,
                  updates_collections=None,
                  name=None,
                  top_k_int=1):
  # pylint: disable=unused-argument
  """
  Compute full normalized discounted cumulative gain (ndcg) based on predictions
  ndcg = dcg_k/idcg_k, k is a cut off ranking postion
  There are a few variants of ndcg
  The dcg (discounted cumulative gain) formula used in
  twml.contrib.metrics.ndcg is::

    \\sum_{i=1}^k \frac{2^{relevance\\_score} -1}{\\log_{2}(i + 1)}

  k is the length of items to be ranked in a batch/query
  Notice that whether k will be replaced with a fixed value requires discussions
  The scores in predictions are transformed to order and relevance scores to calculate ndcg
  A relevance score means how relevant a DataRecord is to a particular query

  Arguments:
    labels: the ground truth value.
    predictions: the predicted values, whose shape must match labels. Ignored for CTR computation.
    metrics_collections: optional list of collections to add this metric into.
    updates_collections: optional list of collections to add the associated update_op into.
    name: an optional variable_scope name.

  Returns:
    ndcg: A `Tensor` representing the ndcg score.
    update_op: A update operation used to accumulate data into this metric.
  """
  with tf.variable_scope(name, 'ndcg', (labels, predictions)):
    label_scores = tf.to_float(labels, name='label_to_float')
    predicted_scores = tf.to_float(predictions, name='predictions_to_float')

    if context.executing_eagerly():
      raise RuntimeError('ndcg is not supported when eager execution '
                         'is enabled.')

    total_ndcg = _metric_variable([], dtypes.float32, name='total_ndcg')
    count_query = _metric_variable([], dtypes.float32, name='query_count')

    # actual ndcg cutoff position top_k_int
    max_prediction_size = array_ops.size(predicted_scores)
    top_k_int = tf.minimum(max_prediction_size, top_k_int)
    # the ndcg score of the batch
    ndcg = math_fns.cal_ndcg(label_scores,
      predicted_scores, top_k_int=top_k_int)
    # add ndcg of the current batch to total_ndcg
    update_total_op = state_ops.assign_add(total_ndcg, ndcg)
    with ops.control_dependencies([ndcg]):
      # count_query stores the number of queries
      # count_query increases by 1 for each batch/query
      update_count_op = state_ops.assign_add(count_query, 1)

    mean_ndcg = math_fns.safe_div(total_ndcg, count_query, 'mean_ndcg')
    update_op = math_fns.safe_div(update_total_op, update_count_op, 'update_mean_ndcg_op')

    if metrics_collections:
      ops.add_to_collections(metrics_collections, mean_ndcg)

    if updates_collections:
      ops.add_to_collections(updates_collections, update_op)

    return mean_ndcg, update_op
//  [python]
def ndcg(labels, predictions,
                  metrics_collections=None,
                  updates_collections=None,
                  name=None,
                  top_k_int=1):
  # pylint: disable=unused-argument
  """
  Compute full normalized discounted cumulative gain (ndcg) based on predictions
  ndcg = dcg_k/idcg_k, k is a cut off ranking postion
  There are a few variants of ndcg
  The dcg (discounted cumulative gain) formula used in
  twml.contrib.metrics.ndcg is::

    \\sum_{i=1}^k \frac{2^{relevance\\_score} -1}{\\log_{2}(i + 1)}

  k is the length of items to be ranked in a batch/query
  Notice that whether k will be replaced with a fixed value requires discussions
  The scores in predictions are transformed to order and relevance scores to calculate ndcg
  A relevance score means how relevant a DataRecord is to a particular query

  Arguments:
    labels: the ground truth value.
    predictions: the predicted values, whose shape must match labels. Ignored for CTR computation.
    metrics_collections: optional list of collections to add this metric into.
    updates_collections: optional list of collections to add the associated update_op into.
    name: an optional variable_scope name.

  Returns:
    ndcg: A `Tensor` representing the ndcg score.
    update_op: A update operation used to accumulate data into this metric.
  """
  with tf.variable_scope(name, 'ndcg', (labels, predictions)):
    label_scores = tf.cast(labels, tf.float32, name='label_to_float')
    predicted_scores = tf.cast(predictions, tf.float32, name='predictions_to_float')

    if context.executing_eagerly():
      raise RuntimeError('ndcg is not supported when eager execution '
                         'is enabled.')

    total_ndcg = _metric_variable([], dtypes.float32, name='total_ndcg')
    count_query = _metric_variable([], dtypes.float32, name='query_count')

    # actual ndcg cutoff position top_k_int
    max_prediction_size = array_ops.size(predicted_scores)
    top_k_int = tf.minimum(max_prediction_size, top_k_int)
    # the ndcg score of the batch
    ndcg = math_fns.cal_ndcg(label_scores,
      predicted_scores, top_k_int=top_k_int)
    # add ndcg of the current batch to total_ndcg
    update_total_op = state_ops.assign_add(total_ndcg, ndcg)
    with ops.control_dependencies([ndcg]):
      # count_query stores the number of queries
      # count_query increases by 1 for each batch/query
      update_count_op = state_ops.assign_add(count_query, 1)

    mean_ndcg = math_fns.safe_div(total_ndcg, count_query, 'mean_ndcg')
    update_op = math_fns.safe_div(update_total_op, update_count_op, 'update_mean_ndcg_op')

    if metrics_collections:
      ops.add_to_collections(metrics_collections, mean_ndcg)

    if updates_collections:
      ops.add_to_collections(updates_collections, update_op)

    return mean_ndcg, update_op
//  [python]
def _classification_loss(self, pred_label, gt_label, num_matched_boxes):
    """Computes the classification loss.

    Computes the classification loss with hard negative mining.
    Args:
      pred_label: a flatten tensor that includes all predicted class. The shape
        is [batch_size, num_anchors, num_classes].
      gt_label: a tensor that represents the classification groundtruth targets.
        The shape is [batch_size, num_anchors, 1].
      num_matched_boxes: the number of anchors that are matched to a groundtruth
        targets. This is used as the loss normalizater.

    Returns:
      box_loss: a float32 representing total box regression loss.
    """
    cross_entropy = tf.losses.sparse_softmax_cross_entropy(
        gt_label, pred_label, reduction=tf.losses.Reduction.NONE)

    mask = tf.greater(tf.squeeze(gt_label), 0)
    float_mask = tf.cast(mask, tf.float32)

    # Hard example mining
    neg_masked_cross_entropy = cross_entropy * (1 - float_mask)
    relative_position = contrib_framework.argsort(
        contrib_framework.argsort(
            neg_masked_cross_entropy, direction='DESCENDING'))
    num_neg_boxes = tf.minimum(
        tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE,
        ssd_constants.NUM_SSD_BOXES)
    top_k_neg_mask = tf.cast(tf.less(
        relative_position,
        tf.tile(num_neg_boxes[:, tf.newaxis], (1, ssd_constants.NUM_SSD_BOXES))
    ), tf.float32)

    class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)

    return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def _classification_loss(self, pred_label, gt_label, num_matched_boxes):
    """Computes the classification loss.

    Computes the classification loss with hard negative mining.
    Args:
      pred_label: a flatten tensor that includes all predicted class. The shape
        is [batch_size, num_anchors, num_classes].
      gt_label: a tensor that represents the classification groundtruth targets.
        The shape is [batch_size, num_anchors, 1].
      num_matched_boxes: the number of anchors that are matched to a groundtruth
        targets. This is used as the loss normalizater.

    Returns:
      box_loss: a float32 representing total box regression loss.
    """
    cross_entropy = tf.losses.sparse_softmax_cross_entropy(
        gt_label, pred_label, reduction=tf.losses.Reduction.NONE)

    mask = tf.greater(tf.squeeze(gt_label), 0)
    float_mask = tf.cast(mask, tf.float32)

    # Hard example mining
    neg_masked_cross_entropy = cross_entropy * (1 - float_mask)
    relative_position = contrib_framework.argsort(
        contrib_framework.argsort(
            neg_masked_cross_entropy, direction='DESCENDING'))
    num_neg_boxes = tf.minimum(
        tf.cast(num_matched_boxes, tf.int32) * ssd_constants.NEGS_PER_POSITIVE,
        ssd_constants.NUM_SSD_BOXES)
    top_k_neg_mask = tf.cast(tf.less(
        relative_position,
        tf.tile(num_neg_boxes[:, tf.newaxis], (1, ssd_constants.NUM_SSD_BOXES))
    ), tf.float32)

    class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)

    return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def compute_loss(self, decoder_output, _features, labels):
    """Computes the loss for this model.

    Returns a tuple `(losses, loss)`, where `losses` are the per-batch
    losses and loss is a single scalar tensor to minimize.
    """
    #pylint: disable=R0201
    # Calculate loss per example-timestep of shape [B, T]
    
    losses = seq2seq_losses.cross_entropy_sequence_loss(
        logits=decoder_output.logits[:, :, :],
        targets=tf.transpose(labels["target_ids"][:, 1:], [1, 0]),
        sequence_length=labels["target_len"] - 1)

    # Calculate the average log perplexity
    loss = tf.reduce_sum(losses) / tf.to_float(
        tf.reduce_sum(labels["target_len"] - 1))

    return losses, loss
//  [python]
def compute_loss(self, decoder_output, _features, labels):
    """Computes the loss for this model.

    Returns a tuple `(losses, loss)`, where `losses` are the per-batch
    losses and loss is a single scalar tensor to minimize.
    """
    #pylint: disable=R0201
    # Calculate loss per example-timestep of shape [B, T]
    
    losses = seq2seq_losses.cross_entropy_sequence_loss(
        logits=decoder_output.logits[:, :, :],
        targets=tf.transpose(labels["target_ids"][:, 1:], [1, 0]),
        sequence_length=labels["target_len"] - 1)

    # Calculate the average log perplexity
    loss = tf.reduce_sum(losses) / tf.cast(
        tf.reduce_sum(labels["target_len"] - 1), tf.float32)

    return losses, loss
//  [python]
def _create_predictions(self, decoder_output, features, labels, losses=None):
    """Creates the dictionary of predictions that is returned by the model.
    """
    predictions = {}

    # Add features and, if available, labels to predictions
    predictions.update(_flatten_dict({"features": features}))
    if labels is not None:
      predictions.update(_flatten_dict({"labels": labels}))

    if losses is not None:
      predictions["losses"] = _transpose_batch_time(losses)

    # Decoders returns output in time-major form [T, B, ...]
    # Here we transpose everything back to batch-major for the user
    output_dict = collections.OrderedDict(
        zip(decoder_output._fields, decoder_output))
    decoder_output_flat = _flatten_dict(output_dict)
    
    decoder_output_flat = {
        k: _transpose_batch_time(v)
        for k, v in decoder_output_flat.items()
    }
    predictions.update(decoder_output_flat)

    # If we predict the ids also map them back into the vocab and process them
    if "predicted_ids" in predictions.keys():
      vocab_tables = graph_utils.get_dict_from_collection("vocab_tables")
      target_id_to_vocab = vocab_tables["target_id_to_vocab"]
      predicted_tokens = target_id_to_vocab.lookup(
          tf.to_int64(predictions["predicted_ids"]))
      # Raw predicted tokens
      predictions["predicted_tokens"] = predicted_tokens

    return predictions