jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _advance(step, beam_log_probs, previous_refs,
             area_logits, areas, batch_size, beam_size, append_refs=True,
             condition=None):
  """Advance one element in the tuple for a decoding step.

  Args:
    step: the current decoding step.
    beam_log_probs: [batch_size * beam_size]
    previous_refs: [batch_size * beam_size, input_length - 1, 2]
    area_logits: [batch_size * beam_size, num_areas]
    areas: the areas.
    batch_size: the batch size.
    beam_size: the beam_size.
    append_refs: returning references or ids.
    condition: conditional probability mask in shape [batch_size * beam_size].
  Returns:
    beam_log_probs: [batch_size * beam_size]
    references in shape of [batch_size * beam_size, input_length, 2] or
        ids in shape of [batch_size * beam_size]
  """
  with tf.control_dependencies([
      tf.equal(tf.shape(beam_log_probs), (batch_size * beam_size,))]):
    num_expansions = tf.minimum(beam_size, tf.shape(area_logits)[-1])
    # [batch_size * beam_size, num_expansions]
    area_log_probs = common_layers.log_prob_from_logits(area_logits)
    if condition is not None:
      area_log_probs = area_log_probs * tf.to_float(
          tf.expand_dims(condition, 1))
    top_area_log_probs, top_area_ids = tf.nn.top_k(
        area_log_probs, k=num_expansions)
  if append_refs:
    # [batch_size * beam_size, num_expansions, 2]
    refs = area_utils.area_to_refs(areas["starts"], areas["ends"],
                                   top_area_ids)
    if condition is not None:
      refs = refs * tf.expand_dims(tf.expand_dims(condition, 1), 2)
    refs = tf.reshape(refs, [batch_size, beam_size, num_expansions, 1, 2])
    if step > 0:
      previous_refs = tf.reshape(
          previous_refs, [batch_size, beam_size, 1, step, 2])
      previous_refs = tf.tile(previous_refs, [1, 1, num_expansions, 1, 1])
      new_refs = tf.concat([previous_refs, refs], axis=3)
    else:
      new_refs = refs
    new_refs = tf.reshape(
        new_refs, [batch_size * beam_size * num_expansions, step + 1, 2])
  # [batch_size, beam_size * num_expansions]
  log_probs = tf.reshape(tf.expand_dims(beam_log_probs, 1) + top_area_log_probs,
                         [batch_size, beam_size * num_expansions])
  # [batch_size, beam_size]
  beam_log_probs, beam_indices = tf.nn.top_k(log_probs, k=beam_size)
  beam_indices = tf.reshape(beam_indices, [-1])
  beam_log_probs = tf.reshape(beam_log_probs, [batch_size * beam_size])
  indices = tf.reshape(
      tf.tile(tf.expand_dims(tf.range(batch_size) * beam_size * num_expansions,
                             axis=1), [1, beam_size]), [-1]) + beam_indices
  if append_refs:
    new_refs = tf.gather(new_refs, indices=indices)
  else:
    new_refs = tf.gather(tf.reshape(top_area_ids, [-1]), indices=indices)
  return beam_log_probs, new_refs
//  [python]
def _advance(step, beam_log_probs, previous_refs,
             area_logits, areas, batch_size, beam_size, append_refs=True,
             condition=None):
  """Advance one element in the tuple for a decoding step.

  Args:
    step: the current decoding step.
    beam_log_probs: [batch_size * beam_size]
    previous_refs: [batch_size * beam_size, input_length - 1, 2]
    area_logits: [batch_size * beam_size, num_areas]
    areas: the areas.
    batch_size: the batch size.
    beam_size: the beam_size.
    append_refs: returning references or ids.
    condition: conditional probability mask in shape [batch_size * beam_size].
  Returns:
    beam_log_probs: [batch_size * beam_size]
    references in shape of [batch_size * beam_size, input_length, 2] or
        ids in shape of [batch_size * beam_size]
  """
  with tf.control_dependencies([
      tf.equal(tf.shape(beam_log_probs), (batch_size * beam_size,))]):
    num_expansions = tf.minimum(beam_size, tf.shape(area_logits)[-1])
    # [batch_size * beam_size, num_expansions]
    area_log_probs = common_layers.log_prob_from_logits(area_logits)
    if condition is not None:
      area_log_probs = area_log_probs * tf.cast(
          tf.expand_dims(condition, 1), tf.float32)
    top_area_log_probs, top_area_ids = tf.nn.top_k(
        area_log_probs, k=num_expansions)
  if append_refs:
    # [batch_size * beam_size, num_expansions, 2]
    refs = area_utils.area_to_refs(areas["starts"], areas["ends"],
                                   top_area_ids)
    if condition is not None:
      refs = refs * tf.expand_dims(tf.expand_dims(condition, 1), 2)
    refs = tf.reshape(refs, [batch_size, beam_size, num_expansions, 1, 2])
    if step > 0:
      previous_refs = tf.reshape(
          previous_refs, [batch_size, beam_size, 1, step, 2])
      previous_refs = tf.tile(previous_refs, [1, 1, num_expansions, 1, 1])
      new_refs = tf.concat([previous_refs, refs], axis=3)
    else:
      new_refs = refs
    new_refs = tf.reshape(
        new_refs, [batch_size * beam_size * num_expansions, step + 1, 2])
  # [batch_size, beam_size * num_expansions]
  log_probs = tf.reshape(tf.expand_dims(beam_log_probs, 1) + top_area_log_probs,
                         [batch_size, beam_size * num_expansions])
  # [batch_size, beam_size]
  beam_log_probs, beam_indices = tf.nn.top_k(log_probs, k=beam_size)
  beam_indices = tf.reshape(beam_indices, [-1])
  beam_log_probs = tf.reshape(beam_log_probs, [batch_size * beam_size])
  indices = tf.reshape(
      tf.tile(tf.expand_dims(tf.range(batch_size) * beam_size * num_expansions,
                             axis=1), [1, beam_size]), [-1]) + beam_indices
  if append_refs:
    new_refs = tf.gather(new_refs, indices=indices)
  else:
    new_refs = tf.gather(tf.reshape(top_area_ids, [-1]), indices=indices)
  return beam_log_probs, new_refs
//  [python]
def argmax_with_score(logits, axis=None):
  """Argmax along with the value."""
  axis = axis or len(logits.get_shape()) - 1
  predictions = tf.argmax(logits, axis=axis)

  logits_shape = shape_list(logits)
  prefix_shape, vocab_size = logits_shape[:-1], logits_shape[-1]
  prefix_size = 1
  for d in prefix_shape:
    prefix_size *= d

  # Flatten to extract scores
  flat_logits = tf.reshape(logits, [prefix_size, vocab_size])
  flat_predictions = tf.reshape(predictions, [prefix_size])
  flat_indices = tf.stack(
      [tf.range(tf.to_int64(prefix_size)),
       tf.to_int64(flat_predictions)],
      axis=1)
  flat_scores = tf.gather_nd(flat_logits, flat_indices)

  # Unflatten
  scores = tf.reshape(flat_scores, prefix_shape)

  return predictions, scores
//  [python]
def argmax_with_score(logits, axis=None):
  """Argmax along with the value."""
  axis = axis or len(logits.get_shape()) - 1
  predictions = tf.argmax(logits, axis=axis)

  logits_shape = shape_list(logits)
  prefix_shape, vocab_size = logits_shape[:-1], logits_shape[-1]
  prefix_size = 1
  for d in prefix_shape:
    prefix_size *= d

  # Flatten to extract scores
  flat_logits = tf.reshape(logits, [prefix_size, vocab_size])
  flat_predictions = tf.reshape(predictions, [prefix_size])
  flat_indices = tf.stack(
      [tf.range(tf.cast(prefix_size, tf.int64)),
       tf.cast(flat_predictions, tf.int64)],
      axis=1)
  flat_scores = tf.gather_nd(flat_logits, flat_indices)

  # Unflatten
  scores = tf.reshape(flat_scores, prefix_shape)

  return predictions, scores
//  [python]
def select_fg_for_masks(class_targets, box_targets, boxes,
                        proposal_to_label_map, max_num_fg=128):
  """Selects the fore ground objects for mask branch during training.

  Args:
    class_targets: a tensor of shape [batch_size, num_boxes]  representing the
      class label for each box.
    box_targets: a tensor with a shape of [batch_size, num_boxes, 4]. The tensor
      contains the ground truth pixel coordinates of the scaled images for each
      roi.
    boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
      represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
    proposal_to_label_map: a tensor with a shape of [batch_size, num_boxes].
      This tensor keeps the mapping between proposal to labels.
      proposal_to_label_map[i] means the index of the ground truth instance for
      the i-th proposal.
    max_num_fg: an integer represents the number of masks per image.
  Returns:
    class_targets, boxes, proposal_to_label_map, box_targets that have
    foreground objects.
  """
  with tf.name_scope('select_fg_for_masks'):
    # Masks are for positive (fg) objects only. Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py  # pylint: disable=line-too-long
    batch_size = boxes.shape[0]
    _, fg_indices = tf.nn.top_k(
        tf.to_float(tf.greater(class_targets, 0)), k=max_num_fg)
    # Contructs indices for gather.
    indices = tf.reshape(
        fg_indices + tf.expand_dims(
            tf.range(batch_size) * tf.shape(class_targets)[1], 1), [-1])

    fg_class_targets = tf.reshape(
        tf.gather(tf.reshape(class_targets, [-1, 1]), indices),
        [batch_size, -1])
    fg_box_targets = tf.reshape(
        tf.gather(tf.reshape(box_targets, [-1, 4]), indices),
        [batch_size, -1, 4])
    fg_box_rois = tf.reshape(
        tf.gather(tf.reshape(boxes, [-1, 4]), indices), [batch_size, -1, 4])
    fg_proposal_to_label_map = tf.reshape(
        tf.gather(tf.reshape(proposal_to_label_map, [-1, 1]), indices),
        [batch_size, -1])

  return (fg_class_targets, fg_box_targets, fg_box_rois,
          fg_proposal_to_label_map)
//  [python]
def select_fg_for_masks(class_targets, box_targets, boxes,
                        proposal_to_label_map, max_num_fg=128):
  """Selects the fore ground objects for mask branch during training.

  Args:
    class_targets: a tensor of shape [batch_size, num_boxes]  representing the
      class label for each box.
    box_targets: a tensor with a shape of [batch_size, num_boxes, 4]. The tensor
      contains the ground truth pixel coordinates of the scaled images for each
      roi.
    boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
      represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
    proposal_to_label_map: a tensor with a shape of [batch_size, num_boxes].
      This tensor keeps the mapping between proposal to labels.
      proposal_to_label_map[i] means the index of the ground truth instance for
      the i-th proposal.
    max_num_fg: an integer represents the number of masks per image.
  Returns:
    class_targets, boxes, proposal_to_label_map, box_targets that have
    foreground objects.
  """
  with tf.name_scope('select_fg_for_masks'):
    # Masks are for positive (fg) objects only. Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py  # pylint: disable=line-too-long
    batch_size = boxes.shape[0]
    _, fg_indices = tf.nn.top_k(
        tf.cast(tf.greater(class_targets, 0), tf.float32), k=max_num_fg)
    # Contructs indices for gather.
    indices = tf.reshape(
        fg_indices + tf.expand_dims(
            tf.range(batch_size) * tf.shape(class_targets)[1], 1), [-1])

    fg_class_targets = tf.reshape(
        tf.gather(tf.reshape(class_targets, [-1, 1]), indices),
        [batch_size, -1])
    fg_box_targets = tf.reshape(
        tf.gather(tf.reshape(box_targets, [-1, 4]), indices),
        [batch_size, -1, 4])
    fg_box_rois = tf.reshape(
        tf.gather(tf.reshape(boxes, [-1, 4]), indices), [batch_size, -1, 4])
    fg_proposal_to_label_map = tf.reshape(
        tf.gather(tf.reshape(proposal_to_label_map, [-1, 1]), indices),
        [batch_size, -1])

  return (fg_class_targets, fg_box_targets, fg_box_rois,
          fg_proposal_to_label_map)
//  [python]
def top_1_tpu(inputs):
  """find max and argmax over the last dimension.

  Works well on TPU

  Args:
    inputs: A tensor with shape [..., depth]

  Returns:
    values: a Tensor with shape [...]
    indices: a Tensor with shape [...]
  """
  inputs_max = tf.reduce_max(inputs, axis=-1, keepdims=True)
  mask = tf.to_int32(tf.equal(inputs_max, inputs))
  index = tf.range(tf.shape(inputs)[-1]) * mask
  return tf.squeeze(inputs_max, -1), tf.reduce_max(index, axis=-1)
//  [python]
def top_1_tpu(inputs):
  """find max and argmax over the last dimension.

  Works well on TPU

  Args:
    inputs: A tensor with shape [..., depth]

  Returns:
    values: a Tensor with shape [...]
    indices: a Tensor with shape [...]
  """
  inputs_max = tf.reduce_max(inputs, axis=-1, keepdims=True)
  mask = tf.cast(tf.equal(inputs_max, inputs), tf.int32)
  index = tf.range(tf.shape(inputs)[-1]) * mask
  return tf.squeeze(inputs_max, -1), tf.reduce_max(index, axis=-1)
//  [python]
def weight_targeting(w, k):
  """Weight-level magnitude pruning."""
  k = tf.to_int32(k)
  w_shape = shape_list(w)
  size = tf.to_int32(tf.reduce_prod(w_shape[:-1]))
  w = tf.reshape(w, [size, w_shape[-1]])

  transpose_w = tf.transpose(w)
  thres = contrib.framework().sort(tf.abs(transpose_w), axis=1)[:, k]
  mask = to_float(thres[None, :] >= tf.abs(w))

  return tf.reshape(mask, w_shape)
//  [python]
def weight_targeting(w, k):
  """Weight-level magnitude pruning."""
  k = tf.cast(k, tf.int32)
  w_shape = shape_list(w)
  size = tf.cast(tf.reduce_prod(w_shape[:-1]), tf.int32)
  w = tf.reshape(w, [size, w_shape[-1]])

  transpose_w = tf.transpose(w)
  thres = contrib.framework().sort(tf.abs(transpose_w), axis=1)[:, k]
  mask = tf.cast(thres[None, :] >= tf.abs(w), tf.float32)

  return tf.reshape(mask, w_shape)
//  [python]
def unit_targeting(w, k):
  """Unit-level magnitude pruning."""
  k = tf.to_int32(k)
  w_shape = shape_list(w)
  size = tf.to_int32(tf.reduce_prod(w_shape[:-1]))
  w = tf.reshape(w, [size, w_shape[-1]])

  norm = tf.norm(w, axis=0)
  thres = contrib.framework().sort(norm, axis=0)[k]
  mask = to_float(thres >= norm)[None, :]
  mask = tf.tile(mask, [size, 1])

  return tf.reshape(mask, w_shape)
//  [python]
def unit_targeting(w, k):
  """Unit-level magnitude pruning."""
  k = tf.cast(k, tf.int32)
  w_shape = shape_list(w)
  size = tf.cast(tf.reduce_prod(w_shape[:-1]), tf.int32)
  w = tf.reshape(w, [size, w_shape[-1]])

  norm = tf.norm(w, axis=0)
  thres = contrib.framework().sort(norm, axis=0)[k]
  mask = tf.cast(thres >= norm, tf.float32)[None, :]
  mask = tf.tile(mask, [size, 1])

  return tf.reshape(mask, w_shape)
//  [python]
def multinomial_sample(x, vocab_size, temperature):
  """Multinomial sampling from a n-dimensional tensor."""
  if temperature > 0:
    samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
  else:
    samples = tf.argmax(x, axis=-1)
  reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
  return tf.to_int32(reshaped_samples)
//  [python]
def multinomial_sample(x, vocab_size, temperature):
  """Multinomial sampling from a n-dimensional tensor."""
  if temperature > 0:
    samples = tf.multinomial(tf.reshape(x, [-1, vocab_size]) / temperature, 1)
  else:
    samples = tf.argmax(x, axis=-1)
  reshaped_samples = tf.reshape(samples, common_layers.shape_list(x)[:-1])
  return tf.cast(reshaped_samples, tf.int32)
//  [python]
def index_last_dim_with_indices(x, indices):
  """Use indices to index into the last axis of x.

  This can be useful for recovering the actual probabilities of a sample from a
  probability distribution.

  Args:
    x: Tensor, n-d.
    indices: Tensor, (n-1)-d, where the dimension sizes match the first (n-1)
      dimensions of x. The values of indices will be used to index into the last
      axis of x.

  Returns:
    Tensor, (n-1)-d.
  """
  assert len(x.shape) == len(indices.shape) + 1

  x_shape = shape_list(x)
  vocab_size = x_shape[-1]

  flat_x = tf.reshape(x, [list_product(x_shape[:-1]), vocab_size])
  flat_indices = tf.reshape(indices, [list_product(x_shape[:-1])])

  idx = tf.stack(
      [
          tf.range(tf.to_int64(shape_list(flat_indices)[0])),
          tf.to_int64(flat_indices)
      ],
      axis=1)
  flat_x_idx = tf.gather_nd(flat_x, idx)

  x_idx = tf.reshape(flat_x_idx, x_shape[:-1])

  return x_idx
//  [python]
def index_last_dim_with_indices(x, indices):
  """Use indices to index into the last axis of x.

  This can be useful for recovering the actual probabilities of a sample from a
  probability distribution.

  Args:
    x: Tensor, n-d.
    indices: Tensor, (n-1)-d, where the dimension sizes match the first (n-1)
      dimensions of x. The values of indices will be used to index into the last
      axis of x.

  Returns:
    Tensor, (n-1)-d.
  """
  assert len(x.shape) == len(indices.shape) + 1

  x_shape = shape_list(x)
  vocab_size = x_shape[-1]

  flat_x = tf.reshape(x, [list_product(x_shape[:-1]), vocab_size])
  flat_indices = tf.reshape(indices, [list_product(x_shape[:-1])])

  idx = tf.stack(
      [
          tf.range(tf.cast(shape_list(flat_indices)[0], tf.int64)),
          tf.cast(flat_indices, tf.int64)
      ],
      axis=1)
  flat_x_idx = tf.gather_nd(flat_x, idx)

  x_idx = tf.reshape(flat_x_idx, x_shape[:-1])

  return x_idx
//  [python]
def blend(image1, image2, factor):
    """Blend image1 and image2 using 'factor'.

    Factor can be above 0.0.  A value of 0.0 means only image1 is used.
    A value of 1.0 means only image2 is used.  A value between 0.0 and
    1.0 means we linearly interpolate the pixel values between the two
    images.  A value greater than 1.0 "extrapolates" the difference
    between the two pixel values, and we clip the results to values
    between 0 and 255.

    Args:
        image1: An image Tensor of type uint8.
        image2: An image Tensor of type uint8.
        factor: A floating point value above 0.0.

    Returns:
        A blended image Tensor of type uint8.
    """
    if factor == 0.0:
        return tf.convert_to_tensor(image1)
    if factor == 1.0:
        return tf.convert_to_tensor(image2)

    image1 = tf.to_float(image1)
    image2 = tf.to_float(image2)

    difference = image2 - image1
    scaled = factor * difference

    # Do addition in float.
    temp = tf.to_float(image1) + scaled

    # Interpolate
    if 1.0 > factor > 0.0:
        # Interpolation means we always stay within 0 and 255.
        return tf.cast(temp, tf.uint8)

    # Extrapolate:
    #
    # We need to clip and then cast.
    return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def blend(image1, image2, factor):
    """Blend image1 and image2 using 'factor'.

    Factor can be above 0.0.  A value of 0.0 means only image1 is used.
    A value of 1.0 means only image2 is used.  A value between 0.0 and
    1.0 means we linearly interpolate the pixel values between the two
    images.  A value greater than 1.0 "extrapolates" the difference
    between the two pixel values, and we clip the results to values
    between 0 and 255.

    Args:
        image1: An image Tensor of type uint8.
        image2: An image Tensor of type uint8.
        factor: A floating point value above 0.0.

    Returns:
        A blended image Tensor of type uint8.
    """
    if factor == 0.0:
        return tf.convert_to_tensor(image1)
    if factor == 1.0:
        return tf.convert_to_tensor(image2)

    image1 = tf.cast(image1, tf.float32)
    image2 = tf.cast(image2, tf.float32)

    difference = image2 - image1
    scaled = factor * difference

    # Do addition in float.
    temp = tf.cast(image1, tf.float32) + scaled

    # Interpolate
    if 1.0 > factor > 0.0:
        # Interpolation means we always stay within 0 and 255.
        return tf.cast(temp, tf.uint8)

    # Extrapolate:
    #
    # We need to clip and then cast.
    return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def _segmentation_loss(logits, labels, params):
  """Compute segmentation loss. So far it's only for single scale.

  Args:
    logits: A tensor specifies the logits as returned from model function.
      The tensor size is [batch_size, height_l, width_l, num_classes].
      The height_l and width_l depends on the min_level feature resolution.
    labels: A tensor specifies the groundtruth targets "cls_targets",
      as returned from dataloader. The tensor has the same spatial resolution
      as input image with size [batch_size, height, width, 1].
    params: Dictionary including training parameters specified in
      default_hparams function in this file.
  Returns:
    A float tensor representing total classification loss. The loss is
      normalized by the total non-ignored pixels.
  """
  # Downsample labels by the min_level feature stride.
  stride = 2**params['min_level']
  scaled_labels = labels[:, 0::stride, 0::stride]

  scaled_labels = tf.cast(scaled_labels, tf.int32)
  scaled_labels = scaled_labels[:, :, :, 0]
  bit_mask = tf.not_equal(scaled_labels, params['ignore_label'])
  # Assign ignore label to background to avoid error when computing
  # Cross entropy loss.
  scaled_labels = tf.where(bit_mask, scaled_labels,
                           tf.zeros_like(scaled_labels))

  normalizer = tf.reduce_sum(tf.to_float(bit_mask))
  cross_entropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=scaled_labels, logits=logits)
  cross_entropy_loss *= tf.to_float(bit_mask)
  loss = tf.reduce_sum(cross_entropy_loss) / normalizer
  return loss
//  [python]
def _segmentation_loss(logits, labels, params):
  """Compute segmentation loss. So far it's only for single scale.

  Args:
    logits: A tensor specifies the logits as returned from model function.
      The tensor size is [batch_size, height_l, width_l, num_classes].
      The height_l and width_l depends on the min_level feature resolution.
    labels: A tensor specifies the groundtruth targets "cls_targets",
      as returned from dataloader. The tensor has the same spatial resolution
      as input image with size [batch_size, height, width, 1].
    params: Dictionary including training parameters specified in
      default_hparams function in this file.
  Returns:
    A float tensor representing total classification loss. The loss is
      normalized by the total non-ignored pixels.
  """
  # Downsample labels by the min_level feature stride.
  stride = 2**params['min_level']
  scaled_labels = labels[:, 0::stride, 0::stride]

  scaled_labels = tf.cast(scaled_labels, tf.int32)
  scaled_labels = scaled_labels[:, :, :, 0]
  bit_mask = tf.not_equal(scaled_labels, params['ignore_label'])
  # Assign ignore label to background to avoid error when computing
  # Cross entropy loss.
  scaled_labels = tf.where(bit_mask, scaled_labels,
                           tf.zeros_like(scaled_labels))

  normalizer = tf.reduce_sum(tf.cast(bit_mask, tf.float32))
  cross_entropy_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=scaled_labels, logits=logits)
  cross_entropy_loss *= tf.cast(bit_mask, tf.float32)
  loss = tf.reduce_sum(cross_entropy_loss) / normalizer
  return loss
//  [python]
def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
            alpha=0.0, use_tpu=False):
    """Produce predictions from the model."""
    if not self._hparams.do_mask:
      infer_out = super(TransformerAE, self).infer(
          features, decode_length, beam_size, top_beams, alpha, use_tpu=use_tpu)
      return infer_out["outputs"]
    if not features:
      features = {}
    inputs_old = None
    if "inputs" in features and len(features["inputs"].shape) < 4:
      inputs_old = features["inputs"]
      features["inputs"] = tf.expand_dims(features["inputs"], 2)

    # Create an initial targets tensor.
    if "partial_targets" in features:
      initial_output = tf.convert_to_tensor(features["partial_targets"])
    else:
      # inputs might not be present in features (e.g.: language modeling),
      # in which case we fallback to 'infer_targets' for calculating initial
      # input shape, type, etc.
      inputs_or_targets = features.get("inputs", features.get("infer_targets"))
      batch_size = common_layers.shape_list(inputs_or_targets)[0]
      length = common_layers.shape_list(inputs_or_targets)[1]
      hidden_dim = common_layers.shape_list(inputs_or_targets)[-1]
      target_length = tf.to_int32(2.0 * tf.to_float(length))
      initial_output = tf.zeros((batch_size, target_length, 1, hidden_dim),
                                dtype=inputs_or_targets.dtype)

    features["targets"] = initial_output
    logits, _ = self(features)  # pylint: disable=not-callable
    # this should only happen if we're doing target_modality not real
    if inputs_or_targets.dtype == tf.float32:
      samples = logits
    else:
      samples = tf.argmax(logits, axis=-1)

    # More steps.
    self.predict_mask = 0.0  # Use the provided targets this time.
    how_many_more_steps = 0  # Set to 1 or more for Gibbs-like sampling.
    for _ in range(how_many_more_steps):
      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
        features["targets"] = samples
        logits, _ = self(features)  # pylint: disable=not-callable
        if inputs_or_targets.dtype == tf.float32:
          # When target_modality is real, the last axis does not represent
          # classes, so it should not be argmax'ed
          samples = logits
        else:
          samples = tf.argmax(logits, axis=-1)

    self.predict_mask = 1.0
    if inputs_old is not None:  # Restore to not confuse Estimator.
      features["inputs"] = inputs_old
    return samples
//  [python]
def infer(self, features=None, decode_length=50, beam_size=1, top_beams=1,
            alpha=0.0, use_tpu=False):
    """Produce predictions from the model."""
    if not self._hparams.do_mask:
      infer_out = super(TransformerAE, self).infer(
          features, decode_length, beam_size, top_beams, alpha, use_tpu=use_tpu)
      return infer_out["outputs"]
    if not features:
      features = {}
    inputs_old = None
    if "inputs" in features and len(features["inputs"].shape) < 4:
      inputs_old = features["inputs"]
      features["inputs"] = tf.expand_dims(features["inputs"], 2)

    # Create an initial targets tensor.
    if "partial_targets" in features:
      initial_output = tf.convert_to_tensor(features["partial_targets"])
    else:
      # inputs might not be present in features (e.g.: language modeling),
      # in which case we fallback to 'infer_targets' for calculating initial
      # input shape, type, etc.
      inputs_or_targets = features.get("inputs", features.get("infer_targets"))
      batch_size = common_layers.shape_list(inputs_or_targets)[0]
      length = common_layers.shape_list(inputs_or_targets)[1]
      hidden_dim = common_layers.shape_list(inputs_or_targets)[-1]
      target_length = tf.cast(2.0 * tf.cast(length, tf.float32), tf.int32)
      initial_output = tf.zeros((batch_size, target_length, 1, hidden_dim),
                                dtype=inputs_or_targets.dtype)

    features["targets"] = initial_output
    logits, _ = self(features)  # pylint: disable=not-callable
    # this should only happen if we're doing target_modality not real
    if inputs_or_targets.dtype == tf.float32:
      samples = logits
    else:
      samples = tf.argmax(logits, axis=-1)

    # More steps.
    self.predict_mask = 0.0  # Use the provided targets this time.
    how_many_more_steps = 0  # Set to 1 or more for Gibbs-like sampling.
    for _ in range(how_many_more_steps):
      with tf.variable_scope(tf.get_variable_scope(), reuse=True):
        features["targets"] = samples
        logits, _ = self(features)  # pylint: disable=not-callable
        if inputs_or_targets.dtype == tf.float32:
          # When target_modality is real, the last axis does not represent
          # classes, so it should not be argmax'ed
          samples = logits
        else:
          samples = tf.argmax(logits, axis=-1)

    self.predict_mask = 1.0
    if inputs_old is not None:  # Restore to not confuse Estimator.
      features["inputs"] = inputs_old
    return samples
//  [python]
def _apply_bbox_augmentation(image, bbox, augmentation_func, *args):
    """Applies augmentation_func to the subsection of image indicated by bbox.

    Args:
        image: 3D uint8 Tensor.
        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
            of type float that represents the normalized coordinates between 0 and 1.
        augmentation_func: Augmentation function that will be applied to the
            subsection of image.
        *args: Additional parameters that will be passed into augmentation_func
            when it is called.

    Returns:
        A modified version of image, where the bbox location in the image will
        have `ugmentation_func applied to it.
    """
    image_height = tf.to_float(tf.shape(image)[0])
    image_width = tf.to_float(tf.shape(image)[1])
    min_y = tf.to_int32(image_height * bbox[0])
    min_x = tf.to_int32(image_width * bbox[1])
    max_y = tf.to_int32(image_height * bbox[2])
    max_x = tf.to_int32(image_width * bbox[3])
    image_height = tf.to_int32(image_height)
    image_width = tf.to_int32(image_width)

    # Clip to be sure the max values do not fall out of range.
    max_y = tf.minimum(max_y, image_height - 1)
    max_x = tf.minimum(max_x, image_width - 1)

    # Get the sub-tensor that is the image within the bounding box region.
    bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

    # Apply the augmentation function to the bbox portion of the image.
    augmented_bbox_content = augmentation_func(bbox_content, *args)

    # Pad the augmented_bbox_content and the mask to match the shape of original
    # image.
    augmented_bbox_content = tf.pad(
        augmented_bbox_content,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]])

    # Create a mask that will be used to zero out a part of the original image.
    mask_tensor = tf.zeros_like(bbox_content)

    mask_tensor = tf.pad(
        mask_tensor,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]],
        constant_values=1)
    # Replace the old bbox content with the new augmented content.
    image = image * mask_tensor + augmented_bbox_content
    return image
//  [python]
def _apply_bbox_augmentation(image, bbox, augmentation_func, *args):
    """Applies augmentation_func to the subsection of image indicated by bbox.

    Args:
        image: 3D uint8 Tensor.
        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
            of type float that represents the normalized coordinates between 0 and 1.
        augmentation_func: Augmentation function that will be applied to the
            subsection of image.
        *args: Additional parameters that will be passed into augmentation_func
            when it is called.

    Returns:
        A modified version of image, where the bbox location in the image will
        have `ugmentation_func applied to it.
    """
    image_height = tf.cast(tf.shape(image)[0], tf.float32)
    image_width = tf.cast(tf.shape(image)[1], tf.float32)
    min_y = tf.to_int32(image_height * bbox[0])
    min_x = tf.to_int32(image_width * bbox[1])
    max_y = tf.to_int32(image_height * bbox[2])
    max_x = tf.to_int32(image_width * bbox[3])
    image_height = tf.to_int32(image_height)
    image_width = tf.to_int32(image_width)

    # Clip to be sure the max values do not fall out of range.
    max_y = tf.minimum(max_y, image_height - 1)
    max_x = tf.minimum(max_x, image_width - 1)

    # Get the sub-tensor that is the image within the bounding box region.
    bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

    # Apply the augmentation function to the bbox portion of the image.
    augmented_bbox_content = augmentation_func(bbox_content, *args)

    # Pad the augmented_bbox_content and the mask to match the shape of original
    # image.
    augmented_bbox_content = tf.pad(
        augmented_bbox_content,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]])

    # Create a mask that will be used to zero out a part of the original image.
    mask_tensor = tf.zeros_like(bbox_content)

    mask_tensor = tf.pad(
        mask_tensor,
        [[min_y, (image_height - 1) - max_y],
         [min_x, (image_width - 1) - max_x],
         [0, 0]],
        constant_values=1)
    # Replace the old bbox content with the new augmented content.
    image = image * mask_tensor + augmented_bbox_content
    return image
//  [python]
def step_num():
  return tf.to_float(tf.train.get_or_create_global_step())
//  [python]
def step_num():
  return tf.cast(tf.train.get_or_create_global_step(), tf.float32)
//  [python]
def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
    """Shifts the bbox according to how the image was sheared.

    Args:
        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
            of type float that represents the normalized coordinates between 0 and 1.
        image_height: Int, height of the image.
        image_width: Int, height of the image.
        level: Float. How much to shear the image.
        shear_horizontal: If true then shear in X dimension else shear in
            the Y dimension.

    Returns:
        A tensor of the same shape as bbox, but now with the shifted coordinates.
    """
    image_height, image_width = (
            tf.to_float(image_height), tf.to_float(image_width))

    # Change bbox coordinates to be pixels.
    min_y = tf.to_int32(image_height * bbox[0])
    min_x = tf.to_int32(image_width * bbox[1])
    max_y = tf.to_int32(image_height * bbox[2])
    max_x = tf.to_int32(image_width * bbox[3])
    coordinates = tf.stack(
        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
    coordinates = tf.cast(coordinates, tf.float32)

    # Shear the coordinates according to the translation matrix.
    if shear_horizontal:
        translation_matrix = tf.stack(
            [[1, 0], [-level, 1]])
    else:
        translation_matrix = tf.stack(
            [[1, -level], [0, 1]])
    translation_matrix = tf.cast(translation_matrix, tf.float32)
    new_coords = tf.cast(
        tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

    # Find min/max values and convert them back to floats.
    min_y = tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height
    min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width
    max_y = tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height
    max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
    """Shifts the bbox according to how the image was sheared.

    Args:
        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
            of type float that represents the normalized coordinates between 0 and 1.
        image_height: Int, height of the image.
        image_width: Int, height of the image.
        level: Float. How much to shear the image.
        shear_horizontal: If true then shear in X dimension else shear in
            the Y dimension.

    Returns:
        A tensor of the same shape as bbox, but now with the shifted coordinates.
    """
    image_height, image_width = (
            tf.to_float(image_height), tf.to_float(image_width))

    # Change bbox coordinates to be pixels.
    min_y = tf.cast(image_height * bbox[0], tf.int32)
    min_x = tf.cast(image_width * bbox[1], tf.int32)
    max_y = tf.cast(image_height * bbox[2], tf.int32)
    max_x = tf.cast(image_width * bbox[3], tf.int32)
    coordinates = tf.stack(
        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
    coordinates = tf.cast(coordinates, tf.float32)

    # Shear the coordinates according to the translation matrix.
    if shear_horizontal:
        translation_matrix = tf.stack(
            [[1, 0], [-level, 1]])
    else:
        translation_matrix = tf.stack(
            [[1, -level], [0, 1]])
    translation_matrix = tf.cast(translation_matrix, tf.float32)
    new_coords = tf.cast(
        tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

    # Find min/max values and convert them back to floats.
    min_y = tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height
    min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width
    max_y = tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height
    max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width

    # Clip the bboxes to be sure the fall between [0, 1].
    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
    return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _batch_stitch(features, mean_length=4.0, stddev=2.0):
  """Stitches a batch of single-step data to a batch of multi-step data."""
  batch_size = common_layers.shape_list(features['task'])[0]
  num_sequences = tf.maximum(
      tf.to_int32(tf.to_float(batch_size) / mean_length), 1)
  lengths = tf.random.truncated_normal(shape=[num_sequences],
                                       mean=mean_length, stddev=stddev)
  max_length = tf.reduce_max(lengths) * (
      tf.to_float(batch_size) / tf.reduce_sum(lengths))
  max_length = tf.to_int32(tf.ceil(max_length))
  total_items = max_length * num_sequences
  num_paddings = total_items - batch_size
  indices = tf.random.shuffle(tf.range(total_items))
  for key in features:
    shape_list = common_layers.shape_list(features[key])
    assert len(shape_list) >= 1
    with tf.control_dependencies([
        tf.assert_greater_equal(num_paddings, 0,
                                name='num_paddings_positive')]):
      paddings = [[0, num_paddings]] + [[0, 0]] * (len(shape_list) - 1)
    features[key] = tf.pad(features[key], paddings,
                           constant_values=-1 if key == 'obj_type' else 0)
    features[key] = tf.gather(features[key], indices)
    shape = [num_sequences, max_length]
    if len(shape_list) >= 2:
      shape += shape_list[1:]
    features[key] = tf.reshape(features[key], shape)
  # Remove all-padding seqs
  step_mask = tf.reduce_any(tf.greater(features['task'], 1), axis=-1)
  mask = tf.reduce_any(step_mask, axis=-1)
  step_mask = tf.boolean_mask(step_mask, mask)
  for key in features:
    features[key] = tf.boolean_mask(features[key], mask=mask)
  num_sequences = tf.shape(features['task'])[0]
  # Sort steps within each seq
  _, step_indices = tf.math.top_k(tf.to_int32(step_mask), k=max_length)
  step_indices = step_indices + tf.expand_dims(
      tf.range(num_sequences) * max_length, 1)
  step_indices = tf.reshape(step_indices, [-1])
  for key in features:
    shape_list = common_layers.shape_list(features[key])
    features[key] = tf.gather(tf.reshape(features[key], [-1] + shape_list[2:]),
                              step_indices)
    features[key] = tf.reshape(features[key], shape_list)
  features = _stitch(features)
  return features
//  [python]
def _batch_stitch(features, mean_length=4.0, stddev=2.0):
  """Stitches a batch of single-step data to a batch of multi-step data."""
  batch_size = common_layers.shape_list(features['task'])[0]
  num_sequences = tf.maximum(
      tf.cast(tf.to_float(batch_size) / mean_length, tf.int32), 1)
  lengths = tf.random.truncated_normal(shape=[num_sequences],
                                       mean=mean_length, stddev=stddev)
  max_length = tf.reduce_max(lengths) * (
      tf.to_float(batch_size) / tf.reduce_sum(lengths))
  max_length = tf.cast(tf.ceil(max_length), tf.int32)
  total_items = max_length * num_sequences
  num_paddings = total_items - batch_size
  indices = tf.random.shuffle(tf.range(total_items))
  for key in features:
    shape_list = common_layers.shape_list(features[key])
    assert len(shape_list) >= 1
    with tf.control_dependencies([
        tf.assert_greater_equal(num_paddings, 0,
                                name='num_paddings_positive')]):
      paddings = [[0, num_paddings]] + [[0, 0]] * (len(shape_list) - 1)
    features[key] = tf.pad(features[key], paddings,
                           constant_values=-1 if key == 'obj_type' else 0)
    features[key] = tf.gather(features[key], indices)
    shape = [num_sequences, max_length]
    if len(shape_list) >= 2:
      shape += shape_list[1:]
    features[key] = tf.reshape(features[key], shape)
  # Remove all-padding seqs
  step_mask = tf.reduce_any(tf.greater(features['task'], 1), axis=-1)
  mask = tf.reduce_any(step_mask, axis=-1)
  step_mask = tf.boolean_mask(step_mask, mask)
  for key in features:
    features[key] = tf.boolean_mask(features[key], mask=mask)
  num_sequences = tf.shape(features['task'])[0]
  # Sort steps within each seq
  _, step_indices = tf.math.top_k(tf.cast(step_mask, tf.int32), k=max_length)
  step_indices = step_indices + tf.expand_dims(
      tf.range(num_sequences) * max_length, 1)
  step_indices = tf.reshape(step_indices, [-1])
  for key in features:
    shape_list = common_layers.shape_list(features[key])
    features[key] = tf.gather(tf.reshape(features[key], [-1] + shape_list[2:]),
                              step_indices)
    features[key] = tf.reshape(features[key], shape_list)
  features = _stitch(features)
  return features
//  [python]
def adafactor_decay_rate_adam(beta2):
  """Second-moment decay rate like Adam, subsuming the correction factor.

  Args:
    beta2: a float between 0 and 1
  Returns:
    a scalar
  """
  t = tf.to_float(tf.train.get_or_create_global_step()) + 1.0
  decay = beta2 * (1.0 - tf.pow(beta2, t - 1.0)) / (1.0 - tf.pow(beta2, t))
  # decay = tf.cond(tf.equal(t, 1.0), lambda: beta2, lambda: decay)
  return decay
//  [python]
def adafactor_decay_rate_adam(beta2):
  """Second-moment decay rate like Adam, subsuming the correction factor.

  Args:
    beta2: a float between 0 and 1
  Returns:
    a scalar
  """
  t = tf.cast(tf.train.get_or_create_global_step(), tf.float32) + 1.0
  decay = beta2 * (1.0 - tf.pow(beta2, t - 1.0)) / (1.0 - tf.pow(beta2, t))
  # decay = tf.cond(tf.equal(t, 1.0), lambda: beta2, lambda: decay)
  return decay
//  [python]
def specgrams_to_melspecgrams(self, specgrams):
    """Converts specgrams to melspecgrams.
    Args:
      specgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2].
    Returns:
      melspecgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2], mel scaling of frequencies.
    """
    if self._mel_downscale is None:
      return specgrams

    logmag = specgrams[:, :, :, 0]
    p = specgrams[:, :, :, 1]

    mag2 = tf.exp(2.0 * logmag)
    phase_angle = tf.cumsum(p * np.pi, axis=-2)

    l2mel = tf.to_float(self._linear_to_mel_matrix())
    logmelmag2 = self._safe_log(tf.tensordot(mag2, l2mel, 1))
    mel_phase_angle = tf.tensordot(phase_angle, l2mel, 1)
    mel_p = spectral_ops.instantaneous_frequency(mel_phase_angle)

    return tf.concat(
        [logmelmag2[:, :, :, tf.newaxis], mel_p[:, :, :, tf.newaxis]], axis=-1)
//  [python]
def specgrams_to_melspecgrams(self, specgrams):
    """Converts specgrams to melspecgrams.
    Args:
      specgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2].
    Returns:
      melspecgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2], mel scaling of frequencies.
    """
    if self._mel_downscale is None:
      return specgrams

    logmag = specgrams[:, :, :, 0]
    p = specgrams[:, :, :, 1]

    mag2 = tf.exp(2.0 * logmag)
    phase_angle = tf.cumsum(p * np.pi, axis=-2)

    l2mel = tf.cast(self._linear_to_mel_matrix(), tf.float32)
    logmelmag2 = self._safe_log(tf.tensordot(mag2, l2mel, 1))
    mel_phase_angle = tf.tensordot(phase_angle, l2mel, 1)
    mel_p = spectral_ops.instantaneous_frequency(mel_phase_angle)

    return tf.concat(
        [logmelmag2[:, :, :, tf.newaxis], mel_p[:, :, :, tf.newaxis]], axis=-1)
//  [python]
def melspecgrams_to_specgrams(self, melspecgrams):
    """Converts melspecgrams to specgrams.
    Args:
      melspecgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2], mel scaling of frequencies.
    Returns:
      specgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2].
    """
    if self._mel_downscale is None:
      return melspecgrams

    logmelmag2 = melspecgrams[:, :, :, 0]
    mel_p = melspecgrams[:, :, :, 1]

    mel2l = tf.to_float(self._mel_to_linear_matrix())
    mag2 = tf.tensordot(tf.exp(logmelmag2), mel2l, 1)
    logmag = 0.5 * self._safe_log(mag2)
    mel_phase_angle = tf.cumsum(mel_p * np.pi, axis=-2)
    phase_angle = tf.tensordot(mel_phase_angle, mel2l, 1)
    p = spectral_ops.instantaneous_frequency(phase_angle)

    return tf.concat(
        [logmag[:, :, :, tf.newaxis], p[:, :, :, tf.newaxis]], axis=-1)
//  [python]
def melspecgrams_to_specgrams(self, melspecgrams):
    """Converts melspecgrams to specgrams.
    Args:
      melspecgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2], mel scaling of frequencies.
    Returns:
      specgrams: Tensor of log magnitudes and instantaneous frequencies,
        shape [batch, time, freq, 2].
    """
    if self._mel_downscale is None:
      return melspecgrams

    logmelmag2 = melspecgrams[:, :, :, 0]
    mel_p = melspecgrams[:, :, :, 1]

    mel2l = tf.cast(self._mel_to_linear_matrix(), tf.float32)
    mag2 = tf.tensordot(tf.exp(logmelmag2), mel2l, 1)
    logmag = 0.5 * self._safe_log(mag2)
    mel_phase_angle = tf.cumsum(mel_p * np.pi, axis=-2)
    phase_angle = tf.tensordot(mel_phase_angle, mel2l, 1)
    p = spectral_ops.instantaneous_frequency(phase_angle)

    return tf.concat(
        [logmag[:, :, :, tf.newaxis], p[:, :, :, tf.newaxis]], axis=-1)
//  [python]
def train(self, features, labels):
        logits = self.__call__(features, is_training=True)
        log_probs = tf.log(tf.nn.softmax(logits, -1) + 1e-5)  # (bs*seqsLen, num_items)

        labels = tf.reshape(labels, [-1])
        label_ids = tf.one_hot(labels, depth=self.num_items, dtype=tf.float32)
        label_weights = tf.to_float(tf.not_equal(labels, 0))

        per_example_loss = -tf.reduce_sum(log_probs * label_ids, axis=[-1])
        numerator = tf.reduce_sum(label_weights * per_example_loss)
        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = numerator / denominator + tf.losses.get_regularization_loss()
        train_op = self.trainOp(loss)

        with tf.variable_scope("Sequential/TRAIN"):
            _, loss_op = tf.metrics.mean(loss, name='loss')

        loss_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope="main/Sequential/TRAIN")
        assert len(loss_vars) > 0, "(train)metric local variables should not be None."
        loss_init_op = tf.variables_initializer(loss_vars)

        return train_op, loss_op, loss_init_op
//  [python]
def train(self, features, labels):
        logits = self.__call__(features, is_training=True)
        log_probs = tf.log(tf.nn.softmax(logits, -1) + 1e-5)  # (bs*seqsLen, num_items)

        labels = tf.reshape(labels, [-1])
        label_ids = tf.one_hot(labels, depth=self.num_items, dtype=tf.float32)
        label_weights = tf.cast(tf.not_equal(labels, 0), tf.float32)

        per_example_loss = -tf.reduce_sum(log_probs * label_ids, axis=[-1])
        numerator = tf.reduce_sum(label_weights * per_example_loss)
        denominator = tf.reduce_sum(label_weights) + 1e-5
        loss = numerator / denominator + tf.losses.get_regularization_loss()
        train_op = self.trainOp(loss)

        with tf.variable_scope("Sequential/TRAIN"):
            _, loss_op = tf.metrics.mean(loss, name='loss')

        loss_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope="main/Sequential/TRAIN")
        assert len(loss_vars) > 0, "(train)metric local variables should not be None."
        loss_init_op = tf.variables_initializer(loss_vars)

        return train_op, loss_op, loss_init_op