jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gather_tree_from_array(
    t: TensorLike, parent_ids: TensorLike, sequence_length: TensorLike
) -> tf.Tensor:
    """Calculates the full beams for a `TensorArray`.

    Args:
      t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of
        shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]`
        where `s` is the depth shape.
      parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
      sequence_length: The sequence length of shape `[batch_size, beam_width]`.

    Returns:
      A `Tensor` which is a stacked `TensorArray` of the same size and type as
      `t` and where beams are sorted in each `Tensor` according to
      `parent_ids`.
    """
    max_time = parent_ids.shape[0] or tf.shape(parent_ids)[0]
    batch_size = parent_ids.shape[1] or tf.shape(parent_ids)[1]
    beam_width = parent_ids.shape[2] or tf.shape(parent_ids)[2]

    # Generate beam ids that will be reordered by gather_tree.
    beam_ids = tf.reshape(tf.range(beam_width), [1, 1, -1])
    beam_ids = tf.tile(beam_ids, [max_time, batch_size, 1])

    max_sequence_lengths = tf.reduce_max(sequence_length, axis=1)
    max_sequence_lengths = tf.to_int32(max_sequence_lengths)
    sorted_beam_ids = gather_tree(
        step_ids=beam_ids,
        parent_ids=parent_ids,
        max_sequence_lengths=max_sequence_lengths,
        end_token=beam_width + 1,
    )

    # For out of range steps, simply copy the same beam.
    in_bound_steps = tf.transpose(
        tf.sequence_mask(sequence_length, maxlen=max_time), perm=[2, 0, 1]
    )
    sorted_beam_ids = tf.where(in_bound_steps, x=sorted_beam_ids, y=beam_ids)

    # Gather from a tensor with collapsed additional dimensions.
    final_shape = tf.shape(t)
    gather_from = tf.reshape(t, [max_time, batch_size, beam_width, -1])
    ordered = tf.gather(gather_from, sorted_beam_ids, axis=2, batch_dims=2)
    ordered = tf.reshape(ordered, final_shape)

    return ordered
//  [python]
def gather_tree_from_array(
    t: TensorLike, parent_ids: TensorLike, sequence_length: TensorLike
) -> tf.Tensor:
    """Calculates the full beams for a `TensorArray`.

    Args:
      t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of
        shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]`
        where `s` is the depth shape.
      parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
      sequence_length: The sequence length of shape `[batch_size, beam_width]`.

    Returns:
      A `Tensor` which is a stacked `TensorArray` of the same size and type as
      `t` and where beams are sorted in each `Tensor` according to
      `parent_ids`.
    """
    max_time = parent_ids.shape[0] or tf.shape(parent_ids)[0]
    batch_size = parent_ids.shape[1] or tf.shape(parent_ids)[1]
    beam_width = parent_ids.shape[2] or tf.shape(parent_ids)[2]

    # Generate beam ids that will be reordered by gather_tree.
    beam_ids = tf.reshape(tf.range(beam_width), [1, 1, -1])
    beam_ids = tf.tile(beam_ids, [max_time, batch_size, 1])

    max_sequence_lengths = tf.cast(tf.reduce_max(sequence_length, axis=1), tf.int32)
    sorted_beam_ids = gather_tree(
        step_ids=beam_ids,
        parent_ids=parent_ids,
        max_sequence_lengths=max_sequence_lengths,
        end_token=beam_width + 1,
    )

    # For out of range steps, simply copy the same beam.
    in_bound_steps = tf.transpose(
        tf.sequence_mask(sequence_length, maxlen=max_time), perm=[2, 0, 1]
    )
    sorted_beam_ids = tf.where(in_bound_steps, x=sorted_beam_ids, y=beam_ids)

    # Gather from a tensor with collapsed additional dimensions.
    final_shape = tf.shape(t)
    gather_from = tf.reshape(t, [max_time, batch_size, beam_width, -1])
    ordered = tf.gather(gather_from, sorted_beam_ids, axis=2, batch_dims=2)
    ordered = tf.reshape(ordered, final_shape)

    return ordered
//  [python]
def compute_padded_size(desired_size, stride):
  """Compute the padded size given the desired size and the stride.

  The padded size will be the smallest rectangle, such that each dimension is
  the smallest multiple of the stride which is larger than the desired
  dimension. For example, if desired_size = (100, 200) and stride = 32,
  the output padded_size = (128, 224).

  Args:
    desired_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the target output image size.
    stride: an integer, the stride of the backbone network.

  Returns:
    padded_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the padded output image size.
  """
  if isinstance(desired_size, list) or isinstance(desired_size, tuple):
    padded_size = [
        int(math.ceil(d * 1.0 / stride) * stride) for d in desired_size
    ]
  else:
    padded_size = t.cast(
        t.ceil(t.cast(desired_size, dtype=t.float32) / stride) * stride,
        t.int32,
    )
  return padded_size
//  [python]
def compute_padded_size(desired_size, stride):
  """Compute the padded size given the desired size and the stride.

  The padded size will be the smallest rectangle, such that each dimension is
  the smallest multiple of the stride which is larger than the desired
  dimension. For example, if desired_size = (100, 200) and stride = 32,
  the output padded_size = (128, 224).

  Args:
    desired_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the target output image size.
    stride: an integer, the stride of the backbone network.

  Returns:
    padded_size: a `Tensor` or `int` list/tuple of two elements representing
      [height, width] of the padded output image size.
  """
  if isinstance(desired_size, list) or isinstance(desired_size, tuple):
    padded_size = [
        int(math.ceil(d * 1.0 / stride) * stride) for d in desired_size
    ]
  else:
    padded_size = tf.cast(
        tf.math.ceil(tf.cast(desired_size, dtype=tf.float32) / stride) * stride,
        tf.int32,
    )
  return padded_size
//  [python]
def clip_or_pad_to_fixed_size(input_tensor, size, constant_values=0):
  """Pads data to a fixed length at the first dimension.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.
    constant_values: `int` or `str` value assigned to the paddings.

  Returns:
    `Tensor` with the first dimension padded to `size`.
  """
  input_shape = input_tensor.get_shape().as_list()
  padding_shape = []

  # Computes the padding length on the first dimension, clip input tensor if it
  # is longer than `size`.
  input_length = tf.shape(input_tensor)[0]
  input_length = tf.clip_by_value(input_length, 0, size)
  input_tensor = input_tensor[:input_length]

  padding_length = tf.maximum(0, size - input_length)
  padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = tf.fill(dims=padding_shape, value=constant_values)
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  output_shape = input_shape
  output_shape[0] = size
  padded_tensor.set_shape(output_shape)
  return padded_tensor
//  [python]
def clip_or_pad_to_fixed_size(input_tensor, size, constant_values=0):
  """Pads data to a fixed length at the first dimension.

  Args:
    input_tensor: `Tensor` with any dimension.
    size: `int` number for the first dimension of output Tensor.
    constant_values: `int` or `str` value assigned to the paddings.

  Returns:
    `Tensor` with the first dimension padded to `size`.
  """
  input_shape = input_tensor.get_shape().as_list()
  padding_shape = []

  # Computes the padding length on the first dimension, clip input tensor if it
  # is longer than `size`.
  input_length = tf.shape(input_tensor)[0]
  input_length = tf.clip_by_value(input_length, 0, size)
  input_tensor = input_tensor[:input_length]

  padding_length = tf.maximum(0, size - input_length)
  padding_shape.append(padding_length)

  # Copies shapes of the rest of input shape dimensions.
  for i in range(1, len(input_shape)):
    padding_shape.append(tf.shape(input_tensor)[i])

  # Pads input tensor to the fixed first dimension.
  paddings = tf.cast(
      tf.fill(dims=padding_shape, value=constant_values), input_tensor.dtype
  )
  padded_tensor = tf.concat([input_tensor, paddings], axis=0)
  output_shape = input_shape
  output_shape[0] = size
  padded_tensor.set_shape(output_shape)
  return padded_tensor
//  [python]
def center_crop_image(
    image, center_crop_fraction: float = CENTER_CROP_FRACTION
):
  """Center crop a square shape slice from the input image.

  It crops a square shape slice from the image. The side of the actual crop
  is 224 / 256 = 0.875 of the short side of the original image. References:
  [1] Very Deep Convolutional Networks for Large-Scale Image Recognition
      https://arxiv.org/abs/1409.1556
  [2] Deep Residual Learning for Image Recognition
      https://arxiv.org/abs/1512.03385

  Args:
    image: a Tensor of shape [height, width, 3] representing the input image.
    center_crop_fraction: a float of ratio between the side of the cropped image
      and the short side of the original image

  Returns:
    cropped_image: a Tensor representing the center cropped image.
  """
  with tf.name_scope('center_crop_image'):
    image_size = tf.to_int32(tf.shape(image)[:2])
    crop_size = center_crop_fraction * tf.math.minimum(
        image_size[0], image_size[1]
    )
    crop_offset = tf.to_int32((tf.to_float(image_size) - crop_size) / 2.0)
    crop_size = tf.to_int32(crop_size)
    cropped_image = image[
        crop_offset[0] : crop_offset[0] + crop_size,
        crop_offset[1] : crop_offset[1] + crop_size,
        :,
    ]
    return cropped_image
//  [python]
def center_crop_image(
    image, center_crop_fraction: float = CENTER_CROP_FRACTION
):
  """Center crop a square shape slice from the input image.

  It crops a square shape slice from the image. The side of the actual crop
  is 224 / 256 = 0.875 of the short side of the original image. References:
  [1] Very Deep Convolutional Networks for Large-Scale Image Recognition
      https://arxiv.org/abs/1409.1556
  [2] Deep Residual Learning for Image Recognition
      https://arxiv.org/abs/1512.03385

  Args:
    image: a Tensor of shape [height, width, 3] representing the input image.
    center_crop_fraction: a float of ratio between the side of the cropped image
      and the short side of the original image

  Returns:
    cropped_image: a Tensor representing the center cropped image.
  """
  with tf.name_scope('center_crop_image'):
    image_size = tf.cast(tf.shape(image)[:2], dtype=tf.float32)
    crop_size = center_crop_fraction * tf.math.minimum(
        image_size[0], image_size[1]
    )
    crop_offset = tf.cast((image_size - crop_size) / 2.0, dtype=tf.int32)
    crop_size = tf.cast(crop_size, dtype=tf.int32)
    cropped_image = image[
        crop_offset[0] : crop_offset[0] + crop_size,
        crop_offset[1] : crop_offset[1] + crop_size,
        :,
    ]
    return cropped_image
//  [python]
def center_crop_image_v2(
    image_bytes, image_shape, center_crop_fraction: float = CENTER_CROP_FRACTION
):
  """Center crop a square shape slice from the input image.

  It crops a square shape slice from the image. The side of the actual crop
  is 224 / 256 = 0.875 of the short side of the original image. References:
  [1] Very Deep Convolutional Networks for Large-Scale Image Recognition
      https://arxiv.org/abs/1409.1556
  [2] Deep Residual Learning for Image Recognition
      https://arxiv.org/abs/1512.03385

  This is a faster version of `center_crop_image` which takes the original
  image bytes and image size as the inputs, and partially decode the JPEG
  bytes according to the center crop.

  Args:
    image_bytes: a Tensor of type string representing the raw image bytes.
    image_shape: a Tensor specifying the shape of the raw image.
    center_crop_fraction: a float of ratio between the side of the cropped image
      and the short side of the original image

  Returns:
    cropped_image: a Tensor representing the center cropped image.
  """
  with tf.name_scope('center_image_crop_v2'):
    image_shape = tf.to_float(image_shape)
    crop_size = center_crop_fraction * tf.math.minimum(
        image_shape[0], image_shape[1]
    )
    crop_offset = tf.cast((image_shape - crop_size) / 2.0, dtype=tf.int32)
    crop_size = tf.cast(crop_size, dtype=tf.int32)
    crop_window = tf.stack(
        [crop_offset[0], crop_offset[1], crop_size, crop_size]
    )
    cropped_image = tf.image.decode_and_crop_jpeg(
        image_bytes, crop_window, channels=3
    )
    return cropped_image
//  [python]
def center_crop_image_v2(
    image_bytes, image_shape, center_crop_fraction: float = CENTER_CROP_FRACTION
):
  """Center crop a square shape slice from the input image.

  It crops a square shape slice from the image. The side of the actual crop
  is 224 / 256 = 0.875 of the short side of the original image. References:
  [1] Very Deep Convolutional Networks for Large-Scale Image Recognition
      https://arxiv.org/abs/1409.1556
  [2] Deep Residual Learning for Image Recognition
      https://arxiv.org/abs/1512.03385

  This is a faster version of `center_crop_image` which takes the original
  image bytes and image size as the inputs, and partially decode the JPEG
  bytes according to the center crop.

  Args:
    image_bytes: a Tensor of type string representing the raw image bytes.
    image_shape: a Tensor specifying the shape of the raw image.
    center_crop_fraction: a float of ratio between the side of the cropped image
      and the short side of the original image

  Returns:
    cropped_image: a Tensor representing the center cropped image.
  """
  with tf.name_scope('center_image_crop_v2'):
    image_shape = tf.cast(image_shape, tf.float32)
    crop_size = center_crop_fraction * tf.math.minimum(
        image_shape[0], image_shape[1]
    )
    crop_offset = tf.cast((image_shape - crop_size) / 2.0, dtype=tf.int32)
    crop_size = tf.cast(crop_size, dtype=tf.int32)
    crop_window = tf.stack(
        [crop_offset[0], crop_offset[1], crop_size, crop_size]
    )
    cropped_image = tf.image.decode_and_crop_jpeg(
        image_bytes, crop_window, channels=3
    )
    return cropped_image
//  [python]
def _get_scores(
    log_probs,
    sequence_lengths,
    length_penalty_weight,
    coverage_penalty_weight,
    finished,
    accumulated_attention_probs,
):
    """Calculates scores for beam search hypotheses.

    Args:
      log_probs: The log probabilities with shape
        `[batch_size, beam_width, vocab_size]`.
      sequence_lengths: The array of sequence lengths.
      length_penalty_weight: Float weight to penalize length. Disabled with
        0.0.
      coverage_penalty_weight: Float weight to penalize the coverage of source
        sentence. Disabled with 0.0.
      finished: A boolean tensor of shape `[batch_size, beam_width]` that
        specifies which elements in the beam are finished already.
      accumulated_attention_probs: Accumulated attention probabilities up to
        the current time step, with shape `[batch_size, beam_width, max_time]`
        if coverage_penalty_weight is not 0.0.

    Returns:
      The scores normalized by the length_penalty and coverage_penalty.

    Raises:
      ValueError: accumulated_attention_probs is None when coverage penalty is
        enabled.
    """
    length_penalty_ = _length_penalty(
        sequence_lengths=sequence_lengths, penalty_factor=length_penalty_weight
    )
    length_penalty_ = t.cast(length_penalty_, dtype=log_probs.dtype)
    scores = log_probs / length_penalty_

    coverage_penalty_weight = t.convert_to_tensor(
        coverage_penalty_weight, name="coverage_penalty_weight"
    )
    if coverage_penalty_weight.shape.ndims != 0:
        raise ValueError(
            "coverage_penalty_weight should be a scalar, "
            "but saw shape: %s" % coverage_penalty_weight.shape
        )

    if t.get_static_value(coverage_penalty_weight) == 0.0:
        return scores

    if accumulated_attention_probs is None:
        raise ValueError(
            "accumulated_attention_probs can be None only if coverage penalty "
            "is disabled."
        )

    # Add source sequence length mask before computing coverage penalty.
    accumulated_attention_probs = t.where(
        t.equal(accumulated_attention_probs, 0.0),
        t.ones_like(accumulated_attention_probs),
        accumulated_attention_probs,
    )

    # coverage penalty =
    #     sum over `max_time` {log(min(accumulated_attention_probs, 1.0))}
    coverage_penalty = t.reduce_sum(
        t.math.log(t.minimum(accumulated_attention_probs, 1.0)), 2
    )
    # Apply coverage penalty to finished predictions.
    coverage_penalty *= t.cast(finished, t.float32)
    weighted_coverage_penalty = coverage_penalty * coverage_penalty_weight
    # Reshape from [batch_size, beam_width] to [batch_size, beam_width, 1]
    weighted_coverage_penalty = t.expand_dims(weighted_coverage_penalty, 2)
    return scores + weighted_coverage_penalty
//  [python]
def _get_scores(
    log_probs,
    sequence_lengths,
    length_penalty_weight,
    coverage_penalty_weight,
    finished,
    accumulated_attention_probs,
):
    """Calculates scores for beam search hypotheses.

    Args:
      log_probs: The log probabilities with shape
        `[batch_size, beam_width, vocab_size]`.
      sequence_lengths: The array of sequence lengths.
      length_penalty_weight: Float weight to penalize length. Disabled with
        0.0.
      coverage_penalty_weight: Float weight to penalize the coverage of source
        sentence. Disabled with 0.0.
      finished: A boolean tensor of shape `[batch_size, beam_width]` that
        specifies which elements in the beam are finished already.
      accumulated_attention_probs: Accumulated attention probabilities up to
        the current time step, with shape `[batch_size, beam_width, max_time]`
        if coverage_penalty_weight is not 0.0.

    Returns:
      The scores normalized by the length_penalty and coverage_penalty.

    Raises:
      ValueError: accumulated_attention_probs is None when coverage penalty is
        enabled.
    """
    length_penalty_ = _length_penalty(
        sequence_lengths=sequence_lengths, penalty_factor=length_penalty_weight
    )
    length_penalty_ = tf.cast(length_penalty_, dtype=log_probs.dtype)
    scores = log_probs / length_penalty_

    coverage_penalty_weight = tf.convert_to_tensor(
        coverage_penalty_weight, name="coverage_penalty_weight"
    )
    if coverage_penalty_weight.shape.ndims != 0:
        raise ValueError(
            "coverage_penalty_weight should be a scalar, "
            "but saw shape: %s" % coverage_penalty_weight.shape
        )

    if tf.get_static_value(coverage_penalty_weight) == 0.0:
        return scores

    if accumulated_attention_probs is None:
        raise ValueError(
            "accumulated_attention_probs can be None only if coverage penalty "
            "is disabled."
        )

    # Add source sequence length mask before computing coverage penalty.
    accumulated_attention_probs = tf.where(
        tf.equal(accumulated_attention_probs, 0.0),
        tf.ones_like(accumulated_attention_probs),
        accumulated_attention_probs,
    )

    # coverage penalty =
    #     sum over `max_time` {log(min(accumulated_attention_probs, 1.0))}
    coverage_penalty = tf.reduce_sum(
        tf.math.log(tf.minimum(accumulated_attention_probs, 1.0)), 2
    )
    # Apply coverage penalty to finished predictions.
    coverage_penalty *= tf.cast(finished, tf.float32)
    weighted_coverage_penalty = coverage_penalty * coverage_penalty_weight
    # Reshape from [batch_size, beam_width] to [batch_size, beam_width, 1]
    weighted_coverage_penalty = tf.expand_dims(weighted_coverage_penalty, 2)
    return scores + weighted_coverage_penalty
//  [python]
def color_jitter(
    image: tf.Tensor,
    brightness: Optional[float] = 0.0,
    contrast: Optional[float] = 0.0,
    saturation: Optional[float] = 0.0,
    seed: Optional[int] = None,
) -> tf.Tensor:
  """Applies color jitter to an image, similarly to torchvision`s ColorJitter.

  Args:
    image (tf.Tensor): Of shape [height, width, 3] and type uint8.
    brightness (float, optional): Magnitude for brightness jitter. Defaults to
      0.
    contrast (float, optional): Magnitude for contrast jitter. Defaults to 0.
    saturation (float, optional): Magnitude for saturation jitter. Defaults to
      0.
    seed (int, optional): Random seed. Defaults to None.

  Returns:
    tf.Tensor: The augmented `image` of type uint8.
  """
  image = t.cast(image, dtype=tf.uint8)
  image = random_brightness(image, brightness, seed=seed)
  image = random_contrast(image, contrast, seed=seed)
  image = random_saturation(image, saturation, seed=seed)
  return image
//  [python]
def color_jitter(
    image: tf.Tensor,
    brightness: Optional[float] = 0.0,
    contrast: Optional[float] = 0.0,
    saturation: Optional[float] = 0.0,
    seed: Optional[int] = None,
) -> tf.Tensor:
  """Applies color jitter to an image, similarly to torchvision`s ColorJitter.

  Args:
    image (tf.Tensor): Of shape [height, width, 3] and type uint8.
    brightness (float, optional): Magnitude for brightness jitter. Defaults to
      0.
    contrast (float, optional): Magnitude for contrast jitter. Defaults to 0.
    saturation (float, optional): Magnitude for saturation jitter. Defaults to
      0.
    seed (int, optional): Random seed. Defaults to None.

  Returns:
    tf.Tensor: The augmented `image` of type uint8.
  """
  image = tf.cast(image, dtype=tf.uint8)
  image = random_brightness(image, brightness, seed=seed)
  image = random_contrast(image, contrast, seed=seed)
  image = random_saturation(image, saturation, seed=seed)
  return image
//  [python]
def LogitAdjust(y_true, y_pred, sample_per_cls, tau):
  """Implementation of logit adjustment loss.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    tau: Temperature scaling parameter for the base probabilities.

  Returns:
    A loss function with signature loss(y_true, y_pred).
  """
  spc = t.cast(sample_per_cls, dtype=t.float32)
  # Generate class prior (a list of probabilities: P(Y=i))
  spc_norm = spc / t.reduce_sum(spc)
  y_pred = t.cast(y_pred, dtype=t.float32)
  y_pred = y_pred + tau * t.math.log(
      t.cast(spc_norm + 1e-12, dtype=t.float32))
  loss = t.nn.sparse_softmax_cross_entropy_with_logits(
      labels=y_true, logits=y_pred)
  return t.reduce_mean(loss, axis=0)
//  [python]
def LogitAdjust(y_true, y_pred, sample_per_cls, tau):
  """Implementation of logit adjustment loss.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    tau: Temperature scaling parameter for the base probabilities.

  Returns:
    A loss function with signature loss(y_true, y_pred).
  """
  spc = tf.cast(sample_per_cls, dtype=tf.float32)
  # Generate class prior (a list of probabilities: P(Y=i))
  spc_norm = spc / tf.reduce_sum(spc)
  y_pred = tf.cast(y_pred, dtype=tf.float32)
  y_pred = y_pred + tau * tf.math.log(
      tf.cast(spc_norm + 1e-12, dtype=tf.float32))
  loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=y_true, logits=y_pred)
  return tf.reduce_mean(loss, axis=0)
//  [python]
def accuracy(logits, labels, weights):
  """Estimate accuracy of predictions
  Args:
    logits: Logits from inference(). [FLAGS.batch_size, FLAGS.max_doc_length, FLAGS.target_label_size]
    labels: Sentence extraction gold levels [FLAGS.batch_size, FLAGS.max_doc_length, FLAGS.target_label_size]
    weights: Weights to avoid padded part [FLAGS.batch_size, FLAGS.max_doc_length]
  Returns:
    Accuracy: Estimates average of accuracy for each sentence
  """
  with tf.variable_scope('Accuracy') as scope:
    logits = tf.reshape(logits, [-1, FLAGS.target_label_size]) # [FLAGS.batch_size*FLAGS.max_doc_length, FLAGS.target_label_size]
    labels = tf.reshape(labels, [-1, FLAGS.target_label_size]) # [FLAGS.batch_size*FLAGS.max_doc_length, FLAGS.target_label_size]
    correct_pred = tf.equal(tf.argmax(logits,1), tf.argmax(labels,1)) # [FLAGS.batch_size*FLAGS.max_doc_length]
    correct_pred =  tf.reshape(correct_pred, [-1, FLAGS.max_doc_length])  # [FLAGS.batch_size, FLAGS.max_doc_length]
    correct_pred = tf.to_float(correct_pred)
    # Get Accuracy
    accuracy = tf.reduce_mean(correct_pred, name='accuracy')
    if FLAGS.weighted_loss:
      correct_pred = tf.mul(correct_pred, weights)      
      correct_pred = tf.reduce_sum(correct_pred, reduction_indices=1) # [FLAGS.batch_size]
      doc_lengths = tf.reduce_sum(weights, reduction_indices=1) # [FLAGS.batch_size]
      correct_pred_avg = tf.div(correct_pred, doc_lengths)
      accuracy = tf.reduce_mean(correct_pred_avg, name='accuracy')
  return accuracy
//  [python]
def accuracy(logits, labels, weights):
  """Estimate accuracy of predictions
  Args:
    logits: Logits from inference(). [FLAGS.batch_size, FLAGS.max_doc_length, FLAGS.target_label_size]
    labels: Sentence extraction gold levels [FLAGS.batch_size, FLAGS.max_doc_length, FLAGS.target_label_size]
    weights: Weights to avoid padded part [FLAGS.batch_size, FLAGS.max_doc_length]
  Returns:
    Accuracy: Estimates average of accuracy for each sentence
  """
  with tf.variable_scope('Accuracy') as scope:
    logits = tf.reshape(logits, [-1, FLAGS.target_label_size]) # [FLAGS.batch_size*FLAGS.max_doc_length, FLAGS.target_label_size]
    labels = tf.reshape(labels, [-1, FLAGS.target_label_size]) # [FLAGS.batch_size*FLAGS.max_doc_length, FLAGS.target_label_size]
    correct_pred = tf.equal(tf.argmax(logits,1), tf.argmax(labels,1)) # [FLAGS.batch_size*FLAGS.max_doc_length]
    correct_pred =  tf.reshape(correct_pred, [-1, FLAGS.max_doc_length])  # [FLAGS.batch_size, FLAGS.max_doc_length]
    correct_pred = tf.cast(correct_pred, tf.float32)
    # Get Accuracy
    accuracy = tf.reduce_mean(correct_pred, name='accuracy')
    if FLAGS.weighted_loss:
      correct_pred = tf.mul(correct_pred, weights)      
      correct_pred = tf.reduce_sum(correct_pred, reduction_indices=1) # [FLAGS.batch_size]
      doc_lengths = tf.reduce_sum(weights, reduction_indices=1) # [FLAGS.batch_size]
      correct_pred_avg = tf.div(correct_pred, doc_lengths)
      accuracy = tf.reduce_mean(correct_pred_avg, name='accuracy')
  return accuracy
//  [python]
def resize_and_crop_masks(
    masks, image_scale, output_size, offset, centered_crop: bool = False
):
  """Resizes boxes to output size with scale and offset.

  Args:
    masks: `Tensor` of shape [N, H, W, C] representing ground truth masks.
    image_scale: 2D float `Tensor` representing scale factors that apply to
      [height, width] of input image.
    output_size: 2D `Tensor` or `int` representing [height, width] of target
      output image size.
    offset: 2D `Tensor` representing top-left corner [y0, x0] to crop scaled
      boxes.
    centered_crop: If `centered_crop` is set to True, then resized crop (if
      smaller than padded size) is place in the center of the image. Default
      behaviour is to place it at left top corner.

  Returns:
    masks: `Tensor` of shape [N, H, W, C] representing the scaled masks.
  """
  with tf.name_scope('resize_and_crop_masks'):
    mask_size = tf.to_int32(tf.shape(masks)[1:3])
    num_channels = tf.shape(masks)[3]
    # Pad masks to avoid empty mask annotations.
    masks = tf.concat(
        [
            tf.zeros(
                [1, mask_size[0], mask_size[1], num_channels], dtype=masks.dtype
            ),
            masks,
        ],
        axis=0,
    )

    scaled_size = tf.to_int32(image_scale * tf.to_float(mask_size))
    scaled_masks = tf.image.resize(
        masks, scaled_size, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR
    )

    offset = tf.to_int32(offset)
    scaled_masks = scaled_masks[
        :,
        offset[0] : offset[0] + output_size[0],
        offset[1] : offset[1] + output_size[1],
        :,
    ]

    if centered_crop:
      scaled_mask_size = tf.to_int32(tf.shape(scaled_masks)[1:3])
      output_masks = tf.image.pad_to_bounding_box(
          scaled_masks,
          tf.maximum((output_size[0] - scaled_mask_size[0]) // 2, 0),
          tf.maximum((output_size[1] - scaled_mask_size[1]) // 2, 0),
          output_size[0],
          output_size[1],
      )
    else:
      output_masks = tf.image.pad_to_bounding_box(
          scaled_masks, 0, 0, output_size[0], output_size[1]
      )

    # Remove padding.
    output_masks = output_masks[1::]
    return output_masks
//  [python]
def resize_and_crop_masks(
    masks, image_scale, output_size, offset, centered_crop: bool = False
):
  """Resizes boxes to output size with scale and offset.

  Args:
    masks: `Tensor` of shape [N, H, W, C] representing ground truth masks.
    image_scale: 2D float `Tensor` representing scale factors that apply to
      [height, width] of input image.
    output_size: 2D `Tensor` or `int` representing [height, width] of target
      output image size.
    offset: 2D `Tensor` representing top-left corner [y0, x0] to crop scaled
      boxes.
    centered_crop: If `centered_crop` is set to True, then resized crop (if
      smaller than padded size) is place in the center of the image. Default
      behaviour is to place it at left top corner.

  Returns:
    masks: `Tensor` of shape [N, H, W, C] representing the scaled masks.
  """
  with tf.name_scope('resize_and_crop_masks'):
    mask_size = tf.cast(tf.shape(masks)[1:3], tf.float32)
    num_channels = tf.shape(masks)[3]
    # Pad masks to avoid empty mask annotations.
    masks = tf.concat(
        [
            tf.zeros(
                [1, mask_size[0], mask_size[1], num_channels], dtype=masks.dtype
            ),
            masks,
        ],
        axis=0,
    )

    scaled_size = tf.cast(image_scale * mask_size, tf.int32)
    scaled_masks = tf.image.resize(
        masks, scaled_size, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR
    )

    offset = tf.cast(offset, tf.int32)
    scaled_masks = scaled_masks[
        :,
        offset[0] : offset[0] + output_size[0],
        offset[1] : offset[1] + output_size[1],
        :,
    ]

    if centered_crop:
      scaled_mask_size = tf.cast(tf.shape(scaled_masks)[1:3], tf.int32)
      output_masks = tf.image.pad_to_bounding_box(
          scaled_masks,
          tf.maximum((output_size[0] - scaled_mask_size[0]) // 2, 0),
          tf.maximum((output_size[1] - scaled_mask_size[1]) // 2, 0),
          output_size[0],
          output_size[1],
      )
    else:
      output_masks = tf.image.pad_to_bounding_box(
          scaled_masks, 0, 0, output_size[0], output_size[1]
      )

    # Remove padding.
    output_masks = output_masks[1::]
    return output_masks
//  [python]
def LDAMLoss(y_true,
             y_pred,
             sample_per_cls,
             gamma,
             s=30):
  """Computer LDAM loss for MULTICLASS classification.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    gamma: A scalar for the re-weighting of hyper-parameter.
    s: hyper-parameter.

  Returns:
    loss: A tensor of shape [batch_size, num_classes].
  """
  num_classes = y_pred.shape[1]
  class_weight = 1.0 / np.sqrt(np.sqrt(sample_per_cls))
  class_weight = class_weight * (gamma / np.max(class_weight))
  class_weight = tf.convert_to_tensor(class_weight, dtype=tf.float32)
  y_true_oh = ConvertToOneHot(y_true, depth=num_classes)
  y_pred = tf.convert_to_tensor(y_pred)
  y_true = tf.to_int64(y_true)
  batch_w = tf.matmul(class_weight[None, :], tf.transpose(y_true_oh))
  batch_w = tf.reshape(batch_w, (-1, 1))
  y_pred_m = y_pred - batch_w
  # if condition is true, return y_pred_m[index], otherwise return y_pred[index]
  index_bool = tf.cast(y_true_oh, tf.bool)
  output = tf.where(index_bool, y_pred_m, y_pred)
  logits = output
  loss = tf.nn.softmax_cross_entropy_with_logits(
      labels=y_true_oh, logits=logits*s)
  return tf.reduce_mean(loss)
//  [python]
def LDAMLoss(y_true,
             y_pred,
             sample_per_cls,
             gamma,
             s=30):
  """Computer LDAM loss for MULTICLASS classification.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    gamma: A scalar for the re-weighting of hyper-parameter.
    s: hyper-parameter.

  Returns:
    loss: A tensor of shape [batch_size, num_classes].
  """
  num_classes = y_pred.shape[1]
  class_weight = 1.0 / np.sqrt(np.sqrt(sample_per_cls))
  class_weight = class_weight * (gamma / np.max(class_weight))
  class_weight = tf.convert_to_tensor(class_weight, dtype=tf.float32)
  y_true_oh = ConvertToOneHot(y_true, depth=num_classes)
  y_pred = tf.convert_to_tensor(y_pred)
  y_true = tf.dtypes.cast(y_true, dtype=tf.dtypes.int64)
  batch_w = tf.matmul(class_weight[None, :], tf.transpose(y_true_oh))
  batch_w = tf.reshape(batch_w, (-1, 1))
  y_pred_m = y_pred - batch_w
  # if condition is true, return y_pred_m[index], otherwise return y_pred[index]
  index_bool = tf.cast(y_true_oh, tf.bool)
  output = tf.where(index_bool, y_pred_m, y_pred)
  logits = output
  loss = tf.nn.softmax_cross_entropy_with_logits(
      labels=y_true_oh, logits=logits*s)
  return tf.reduce_mean(loss)
//  [python]
def padded_accuracy(logits, labels):
  """Percentage of times that predictions matches labels on non-0s."""
  with tf.name_scope("padded_accuracy"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = t.cast(t.not_equal(labels, 0), t.float32)
    outputs = t.cast(t.argmax(logits, axis=-1), t.int32)
    padded_labels = t.cast(labels, t.int32)
    return t.cast(t.equal(outputs, padded_labels), t.float32), weights
//  [python]
def padded_accuracy(logits, labels):
  """Percentage of times that predictions matches labels on non-0s."""
  with tf.name_scope("padded_accuracy"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    outputs = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
    padded_labels = tf.cast(labels, tf.int32)
    return tf.cast(tf.equal(outputs, padded_labels), tf.float32), weights
//  [python]
def FocalLoss(y_true,
              y_pred,
              sample_per_cls,
              gamma,
              from_logits=False):
  """Computer focal loss for MULTICLASS classification.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    gamma: A scalar for focal loss gamma hyper-parameter.
    from_logits: True if y_pred is not soft-maxed.

  Returns:
    loss: A tensor of shape [batch_size, num_classes].
  """
  batch_size = y_pred.shape[0]
  recip_spc = [1/i for i in sample_per_cls]
  class_weight = tf.cast(recip_spc, dtype=tf.float32)
  # Normalizer to ensure that sum of class weights is equal to batch_size (like
  # in ERM)
  class_weight_norm = float(batch_size) * class_weight
  class_weight_norm /= tf.reduce_sum(class_weight)
  y_pred = tf.convert_to_tensor(y_pred)
  y_true = tf.dtypes.cast(y_true, dtype=tf.dtypes.int64)
  if from_logits:
    y_pred_prob = tf.nn.softmax(y_pred, axis=-1)
  else:
    y_pred_prob = y_pred
   # Equations:
   # loss = -alpha_y * (1-prob_y)^gamma * log(prob_y)
  xent_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=y_true, logits=y_pred)
  probs = tf.gather(y_pred_prob, y_true, axis=-1, batch_dims=y_true.shape.rank)
  focal_modulation = (1 - probs) ** gamma
  loss = focal_modulation * xent_loss
  class_weight_norm = tf.gather(class_weight_norm, y_true, axis=0,
                                batch_dims=y_true.shape.rank)

  loss *= class_weight_norm
  return tf.reduce_sum(loss)/batch_size
//  [python]
def BalancedSoftmax(y_true,
                    y_pred,
                    sample_per_cls,
                    from_logits=False):
  """Implementation of balacned softmax https://arxiv.org/abs/2007.10740.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    from_logits: True if y_pred is not soft-maxed.

  Returns:
    loss: A tensor of shape [batch_size, num_classes].
  """
  num_classes = y_pred.shape[1]
  # batch_size = y_pred.shape[0]
  y_true_oh = ConvertToOneHot(y_true, depth=num_classes)
  y_pred = t.cast(y_pred, dtype=t.float32)
  spc = t.cast(sample_per_cls, dtype=t.float32)
  # Generate class prior (a list of probabilities: P(Y=i))
  spc_norm = spc / t.reduce_sum(spc)
  if from_logits:
    # reweight exponential of logits with class prior
    y_b_exp = t.exp(y_pred) * spc_norm
    logits_modified = y_b_exp / t.reduce_sum(y_b_exp, 1, keepdims=True)
  else:
    raise ValueError(
        'please give me logits inputs'
    )
  cce = tf.keras.losses.CategoricalCrossentropy()
  loss = cce(y_true_oh, logits_modified)
  return loss
//  [python]
def BalancedSoftmax(y_true,
                    y_pred,
                    sample_per_cls,
                    from_logits=False):
  """Implementation of balacned softmax https://arxiv.org/abs/2007.10740.

  Args:
    y_true: True labels, categorical of shape (batch_size,).
    y_pred: logits, a float32 tensor of shape [batch_size,num_classes].
    sample_per_cls: number of samples per class [num_classes].
    from_logits: True if y_pred is not soft-maxed.

  Returns:
    loss: A tensor of shape [batch_size, num_classes].
  """
  num_classes = y_pred.shape[1]
  # batch_size = y_pred.shape[0]
  y_true_oh = ConvertToOneHot(y_true, depth=num_classes)
  y_pred = tf.cast(y_pred, dtype=tf.float32)
  spc = tf.cast(sample_per_cls, dtype=tf.float32)
  # Generate class prior (a list of probabilities: P(Y=i))
  spc_norm = spc / tf.reduce_sum(spc)
  if from_logits:
    # reweight exponential of logits with class prior
    y_b_exp = tf.exp(y_pred) * spc_norm
    logits_modified = y_b_exp / tf.reduce_sum(y_b_exp, 1, keepdims=True)
  else:
    raise ValueError(
        'please give me logits inputs'
    )
  cce = tf.keras.losses.CategoricalCrossentropy()
  loss = cce(y_true_oh, logits_modified)
  return loss
//  [python]
def padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary

  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  with tf.name_scope("loss"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)

    # Calculate smoothing cross entropy
    with tf.name_scope("smoothing_cross_entropy"):
      confidence = 1.0 - smoothing
      low_confidence = (1.0 - confidence) / t.cast(vocab_size - 1, t.float32)
      soft_targets = tf.one_hot(
          t.cast(labels, t.int32),
          depth=vocab_size,
          on_value=confidence,
          off_value=low_confidence)
      xentropy = tf.nn.softmax_cross_entropy_with_logits(
          logits=logits, labels=soft_targets)

      # Calculate the best (lowest) possible value of cross entropy, and
      # subtract from the cross entropy loss.
      normalizing_constant = -(
          confidence * tf.math.log(confidence) +
          t.cast(vocab_size - 1, t.float32) * low_confidence *
          tf.math.log(low_confidence + 1e-20))
      xentropy -= normalizing_constant

    weights = t.cast(tf.not_equal(labels, 0), t.float32)
    return xentropy * weights, weights
//  [python]
def padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary

  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  with tf.name_scope("loss"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)

    # Calculate smoothing cross entropy
    with tf.name_scope("smoothing_cross_entropy"):
      confidence = 1.0 - smoothing
      low_confidence = (1.0 - confidence) / tf.cast(vocab_size - 1, tf.float32)
      soft_targets = tf.one_hot(
          tf.cast(labels, tf.int32),
          depth=vocab_size,
          on_value=confidence,
          off_value=low_confidence)
      xentropy = tf.nn.softmax_cross_entropy_with_logits(
          logits=logits, labels=soft_targets)

      # Calculate the best (lowest) possible value of cross entropy, and
      # subtract from the cross entropy loss.
      normalizing_constant = -(
          confidence * tf.math.log(confidence) +
          tf.cast(vocab_size - 1, tf.float32) * low_confidence *
          tf.math.log(low_confidence + 1e-20))
      xentropy -= normalizing_constant

    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    return xentropy * weights, weights
//  [python]
def padded_sequence_accuracy(logits, labels):
  """Percentage of times that predictions matches labels everywhere (non-0)."""
  with tf.name_scope("padded_sequence_accuracy"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = t.cast(t.not_equal(labels, 0), t.float32)
    outputs = t.cast(t.argmax(logits, axis=-1), t.int32)
    padded_labels = t.cast(labels, t.int32)
    not_correct = t.cast(t.not_equal(outputs, padded_labels),
                          t.float32) * weights
    axis = list(range(1, len(outputs.get_shape())))
    correct_seq = 1.0 - t.minimum(1.0, t.reduce_sum(not_correct, axis=axis))
    return correct_seq, t.constant(1.0)
//  [python]
def padded_sequence_accuracy(logits, labels):
  """Percentage of times that predictions matches labels everywhere (non-0)."""
  with tf.name_scope("padded_sequence_accuracy"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    outputs = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
    padded_labels = tf.cast(labels, tf.int32)
    not_correct = tf.cast(tf.not_equal(outputs, padded_labels),
                          tf.float32) * weights
    axis = list(range(1, len(outputs.get_shape())))
    correct_seq = 1.0 - tf.minimum(1.0, tf.reduce_sum(not_correct, axis=axis))
    return correct_seq, tf.constant(1.0)
//  [python]
def padded_accuracy_topk(logits, labels, k):
  """Percentage of times that top-k predictions matches labels on non-0s."""
  with tf.name_scope("padded_accuracy_topk"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = t.cast(t.not_equal(labels, 0), t.float32)
    effective_k = tf.minimum(k, tf.shape(logits)[-1])
    _, outputs = tf.nn.top_k(logits, k=effective_k)
    outputs = t.cast(outputs, t.int32)
    padded_labels = t.cast(labels, t.int32)
    padded_labels = tf.expand_dims(padded_labels, axis=-1)
    padded_labels += tf.zeros_like(outputs)  # Pad to same shape.
    same = t.cast(t.equal(outputs, padded_labels), t.float32)
    same_topk = tf.reduce_sum(same, axis=-1)
    return same_topk, weights
//  [python]
def padded_accuracy_topk(logits, labels, k):
  """Percentage of times that top-k predictions matches labels on non-0s."""
  with tf.name_scope("padded_accuracy_topk"):
    logits, labels = _pad_tensors_to_same_length(logits, labels)
    weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
    effective_k = tf.minimum(k, tf.shape(logits)[-1])
    _, outputs = tf.nn.top_k(logits, k=effective_k)
    outputs = tf.cast(outputs, tf.int32)
    padded_labels = tf.cast(labels, tf.int32)
    padded_labels = tf.expand_dims(padded_labels, axis=-1)
    padded_labels += tf.zeros_like(outputs)  # Pad to same shape.
    same = tf.cast(tf.equal(outputs, padded_labels), tf.float32)
    same_topk = tf.reduce_sum(same, axis=-1)
    return same_topk, weights
//  [python]
def encode_rtg(self, rtg):
    if self.min_rtg and self.max_rtg:
      rtg = (rtg - self.min_rtg) / (self.max_rtg - self.min_rtg) * self.num_rtg
    return tf.one_hot(t.cast(rtg, tf.int32), self.num_rtg)
//  [python]
def encode_rtg(self, rtg):
    if self.min_rtg and self.max_rtg:
      rtg = (rtg - self.min_rtg) / (self.max_rtg - self.min_rtg) * self.num_rtg
    return tf.one_hot(tf.cast(rtg, tf.int32), self.num_rtg)
//  [python]
def prepare_dataset(self, dataset):
    episodes, valid_steps = dataset.get_all_episodes()
    for episode_num in range(tf.shape(valid_steps)[0]):
      for step_num in range(tf.shape(valid_steps)[1] - 1):
        this_step = tf.nest.map_structure(lambda t: t[episode_num, step_num],
                                          episodes)
        if this_step.is_last() or not valid_steps[episode_num, step_num]:
          continue
        self.rtg_prior[tf.to_int32(this_step.reward)] += 1.
    self.rtg_prior = self.rtg_prior / np.sum(self.rtg_prior)
//  [python]
def prepare_dataset(self, dataset):
    episodes, valid_steps = dataset.get_all_episodes()
    for episode_num in range(tf.shape(valid_steps)[0]):
      for step_num in range(tf.shape(valid_steps)[1] - 1):
        this_step = tf.nest.map_structure(lambda t: t[episode_num, step_num],
                                          episodes)
        if this_step.is_last() or not valid_steps[episode_num, step_num]:
          continue
        self.rtg_prior[tf.cast(this_step.reward, this_step.action.dtype)] += 1.
    self.rtg_prior = self.rtg_prior / np.sum(self.rtg_prior)
//  [python]
def _metric_at_k(self, k=20):
        prediction = self.prediction
        prediction_transposed = tf.transpose(prediction)
        labels = tf.reshape(self.Y, shape=(-1,))
        pred_values = tf.expand_dims(tf.diag_part(tf.nn.embedding_lookup(prediction_transposed, labels)), -1)
        tile_pred_values = tf.tile(pred_values, [1, self.args.n_items])
        ranks = tf.reduce_sum(tf.to_float(prediction[:,1:] > tile_pred_values), -1) + 1

        ndcg = 1. / (log2(1.0 + ranks))
        hit_at_k = tf.nn.in_top_k(prediction, labels, k=k) # also known as Recall@k
        hit_at_k = tf.to_float(hit_at_k)
        istarget = tf.reshape(self.mask, shape=(-1,))
        hit_at_k *= istarget
        ndcg_at_k = ndcg * istarget

        return (tf.reduce_sum(hit_at_k), tf.reduce_sum(ndcg_at_k), tf.reduce_sum(istarget))
//  [python]
def _metric_at_k(self, k=20):
        prediction = self.prediction
        prediction_transposed = tf.transpose(prediction)
        labels = tf.reshape(self.Y, shape=(-1,))
        pred_values = tf.expand_dims(tf.diag_part(tf.nn.embedding_lookup(prediction_transposed, labels)), -1)
        tile_pred_values = tf.tile(pred_values, [1, self.args.n_items])
        ranks = tf.reduce_sum(tf.cast(prediction[:,1:] > tile_pred_values, dtype=tf.float32), -1) + 1

        ndcg = 1. / (log2(1.0 + ranks))
        hit_at_k = tf.nn.in_top_k(prediction, labels, k=k) # also known as Recall@k
        hit_at_k = tf.cast(hit_at_k, dtype=tf.float32)
        istarget = tf.reshape(self.mask, shape=(-1,))
        hit_at_k *= istarget
        ndcg_at_k = ndcg * istarget

        return (tf.reduce_sum(hit_at_k), tf.reduce_sum(ndcg_at_k), tf.reduce_sum(istarget))