jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sparse_softmax_cross_entropy_with_logits_alpha_balanced(
        _sentinel=None,  # pylint: disable=invalid-name
        labels=None,
        logits=None,
        alpha="auto",
        max_alpha_scale=10.0,
        name=None):
    if alpha is None:
        return tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels,logits=logits,name=name)
    elif alpha != "auto":
        raise ValueError(f"alphpa shold be None or auto, not {alpha}.")
    with tf.variable_scope(name,default_name="sparse_softmax_cross_entropy_with_logits_alpha_balanced"):
        probability = tf.nn.softmax(logits)
        labels = tf.expand_dims(labels,axis=-1)
        r_probability = wmlt.batch_gather(probability,labels)
        r_probability = tf.squeeze(r_probability,axis=-1)
        r_probability = tf.maximum(1e-10*(1+r_probability),r_probability)
        num_classes = logits.get_shape().as_list()[-1]
        labels = tf.squeeze(labels,axis=-1)
        one_hot_labels = tf.one_hot(indices=labels,depth=num_classes)
        count = tf.reduce_sum(one_hot_labels,axis=list(range(one_hot_labels.get_shape().ndims-1)),keepdims=True)
        count = tf.cast(tf.maximum(count,1),tf.float32)
        ratio = tf.cast(tf.reduce_prod(tf.shape(labels)),tf.float32)/(count*float(num_classes))
        alpha_t = tf.minimum(max_alpha_scale,1.0/ratio)
        alpha_t = alpha_t*tf.ones_like(one_hot_labels,dtype=tf.float32)
        labels = tf.expand_dims(labels,axis=-1)
        alpha_t = wmlt.batch_gather(params=alpha_t,indices=labels)
        alpha_t = tf.squeeze(alpha_t,axis=-1)
        loss = -alpha_t*tf.math.log(r_probability)
        return loss
//  [python]
def random_blur_3d(frames, height, width, blur_probability=0.5):
  """Apply efficient batch data transformations.

  Args:
    frames: `Tensor` of shape [timesteps, height, width, 3].
    height: the height of image.
    width: the width of image.
    blur_probability: the probaility to apply the blur operator.

  Returns:
    Preprocessed feature list.
  """
  def generate_selector(p, bsz):
    shape = [bsz, 1, 1, 1]
    selector = tf.to_float(
        tf.less(tf.random.uniform(shape, 0, 1, dtype=tf.float32), p))
    return selector

  frames_new = random_blur(frames, height, width, p=1.)
  selector = generate_selector(blur_probability, 1)
  frames = frames_new * selector + frames * (1 - selector)
  frames = tf.clip_by_value(frames, 0., 1.)

  return frames
//  [python]
def random_blur_3d(frames, height, width, blur_probability=0.5):
  """Apply efficient batch data transformations.

  Args:
    frames: `Tensor` of shape [timesteps, height, width, 3].
    height: the height of image.
    width: the width of image.
    blur_probability: the probaility to apply the blur operator.

  Returns:
    Preprocessed feature list.
  """
  def generate_selector(p, bsz):
    shape = [bsz, 1, 1, 1]
    selector = tf.cast(
        tf.less(tf.random.uniform(shape, 0, 1, dtype=tf.float32), p),
        tf.float32)
    return selector

  frames_new = random_blur(frames, height, width, p=1.)
  selector = generate_selector(blur_probability, 1)
  frames = frames_new * selector + frames * (1 - selector)
  frames = tf.clip_by_value(frames, 0., 1.)

  return frames
//  [python]
def transform_targets(y_train, anchors, anchor_masks, classes):
    y_outs = []
    grid_size = 13

    # calculate anchor index for true boxes
    anchors = t.cast(anchors, t.float32)
    anchor_area = anchors[..., 0] * anchors[..., 1]
    box_wh = y_train[..., 2:4] - y_train[..., 0:2]
    box_wh = t.tile(t.expand_dims(box_wh, -2),
                     (1, 1, t.shape(anchors)[0], 1))
    box_area = box_wh[..., 0] * box_wh[..., 1]
    intersection = t.minimum(box_wh[..., 0], anchors[..., 0]) * \
        t.minimum(box_wh[..., 1], anchors[..., 1])
    iou = intersection / (box_area + anchor_area - intersection)
    anchor_idx = t.cast(t.argmax(iou, axis=-1), t.float32)
    anchor_idx = t.expand_dims(anchor_idx, axis=-1)

    y_train = t.concat([y_train, anchor_idx], axis=-1)

    for anchor_idxs in anchor_masks:
        y_outs.append(transform_targets_for_output(
            y_train, grid_size, anchor_idxs, classes))
        grid_size *= 2

    return tuple(y_outs)
//  [python]
def transform_targets(y_train, anchors, anchor_masks, classes):
    y_outs = []
    grid_size = 13

    # calculate anchor index for true boxes
    anchors = tf.cast(anchors, tf.float32)
    anchor_area = anchors[..., 0] * anchors[..., 1]
    box_wh = y_train[..., 2:4] - y_train[..., 0:2]
    box_wh = tf.tile(tf.expand_dims(box_wh, -2),
                     (1, 1, tf.shape(anchors)[0], 1))
    box_area = box_wh[..., 0] * box_wh[..., 1]
    intersection = tf.minimum(box_wh[..., 0], anchors[..., 0]) * \
        tf.minimum(box_wh[..., 1], anchors[..., 1])
    iou = intersection / (box_area + anchor_area - intersection)
    anchor_idx = tf.cast(tf.argmax(iou, axis=-1), tf.float32)
    anchor_idx = tf.expand_dims(anchor_idx, axis=-1)

    y_train = tf.concat([y_train, anchor_idx], axis=-1)

    for anchor_idxs in anchor_masks:
        y_outs.append(transform_targets_for_output(
            y_train, grid_size, anchor_idxs, classes))
        grid_size *= 2

    return tuple(y_outs)
//  [python]
def _covariance_loss(representations: tf.Tensor) -> tf.Tensor:
  """Covariance loss component of VicReg loss.

  Computes normalized square of the off-diagonal elements of the covariance
  matrix of input representations.

  Args:
    representations: Input representations.

  Returns:
    Loss value as scalar `tf.Tensor`.
  """
  batch_size, feature_dim = tf.unstack(
      tf.cast(tf.shape(representations), tf.float32)
  )
  representations = _normalize(representations, scale=False)
  covariance = (
      tf.matmul(representations, representations, transpose_a=True) / batch_size
  )
  covariance = tf.pow(covariance, 2)
  return (
      tf.reduce_sum(
          tf.linalg.set_diag(
              covariance, tf.zeros(tf.cast(feature_dim, tf.int32))
          )
      )
      / feature_dim
  )
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.

  Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.

  Returns:
    A Tensor representing the blurred image.
  """
  radius = t.cast(kernel_size / 2, dtype=t.int32)
  kernel_size = radius * 2 + 1
  x = t.cast(t.range(-radius, radius + 1), dtype=t.float32)
  blur_filter = t.exp(
      -t.pow(x, 2.0) / (2.0 * t.pow(t.cast(sigma, dtype=t.float32), 2.0)))
  blur_filter /= t.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = t.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = t.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = t.shape(image)[-1]
  blur_h = t.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = t.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = t.expand_dims(image, axis=0)
  blurred = t.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = t.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = t.squeeze(blurred, axis=0)
  return blurred
//  [python]
def gaussian_blur(image, kernel_size, sigma, padding='SAME'):
  """Blurs the given image with separable convolution.

  Args:
    image: Tensor of shape [height, width, channels] and dtype float to blur.
    kernel_size: Integer Tensor for the size of the blur kernel. This is should
      be an odd number. If it is an even number, the actual kernel size will be
      size + 1.
    sigma: Sigma value for gaussian operator.
    padding: Padding to use for the convolution. Typically 'SAME' or 'VALID'.

  Returns:
    A Tensor representing the blurred image.
  """
  radius = tf.cast(kernel_size / 2, dtype=tf.int32)
  kernel_size = radius * 2 + 1
  x = tf.cast(tf.range(-radius, radius + 1), dtype=tf.float32)
  blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, dtype=tf.float32), 2.0)))
  blur_filter /= tf.reduce_sum(blur_filter)
  # One vertical and one horizontal filter.
  blur_v = tf.reshape(blur_filter, [kernel_size, 1, 1, 1])
  blur_h = tf.reshape(blur_filter, [1, kernel_size, 1, 1])
  num_channels = tf.shape(image)[-1]
  blur_h = tf.tile(blur_h, [1, 1, num_channels, 1])
  blur_v = tf.tile(blur_v, [1, 1, num_channels, 1])
  expand_batch_dim = image.shape.ndims == 3
  if expand_batch_dim:
    # Tensorflow requires batched input to convolutions, which we can fake with
    # an extra dimension.
    image = tf.expand_dims(image, axis=0)
  blurred = tf.nn.depthwise_conv2d(
      image, blur_h, strides=[1, 1, 1, 1], padding=padding)
  blurred = tf.nn.depthwise_conv2d(
      blurred, blur_v, strides=[1, 1, 1, 1], padding=padding)
  if expand_batch_dim:
    blurred = tf.squeeze(blurred, axis=0)
  return blurred
//  [python]
def lstsq(a, b, rcond=None):
    a = convert_to_tensor(a)
    b = convert_to_tensor(b)
    if a.shape[0] != b.shape[0]:
        raise ValueError("Leading dimensions of input arrays must match")
    b_orig_ndim = b.ndim
    if b_orig_ndim == 1:
        b = b[:, None]
    if a.ndim != 2:
        raise TypeError(
            f"{a.ndim}-dimensional array given. "
            "Array must be two-dimensional"
        )
    if b.ndim != 2:
        raise TypeError(
            f"{b.ndim}-dimensional array given. "
            "Array must be one or two-dimensional"
        )
    m, n = a.shape
    dtype = a.dtype
    eps = tf.experimental.numpy.finfo(dtype).eps
    if a.shape == ():
        s = tf.zeros(0, dtype=a.dtype)
        x = tf.zeros((n, *b.shape[1:]), dtype=a.dtype)
    else:
        if rcond is None:
            rcond = eps * max(n, m)
        else:
            rcond = tf.where(rcond < 0, eps, rcond)
        u, s, vt = svd(a, full_matrices=False)
        mask = s >= tf.convert_to_tensor(rcond, dtype=s.dtype) * s[0]
        safe_s = tf.to_float(tf.where(mask, s, 1))
        s_inv = tf.where(mask, 1 / safe_s, 0)[:, tf.newaxis]
        u_t_b = tf.matmul(tf.transpose(tf.math.conj(u)), b)
        x = tf.matmul(tf.transpose(tf.math.conj(vt)), s_inv * u_t_b)

    if b_orig_ndim == 1:
        x = tf.reshape(x, [-1])
    return x
//  [python]
def lstsq(a, b, rcond=None):
    a = convert_to_tensor(a)
    b = convert_to_tensor(b)
    if a.shape[0] != b.shape[0]:
        raise ValueError("Leading dimensions of input arrays must match")
    b_orig_ndim = b.ndim
    if b_orig_ndim == 1:
        b = b[:, None]
    if a.ndim != 2:
        raise TypeError(
            f"{a.ndim}-dimensional array given. "
            "Array must be two-dimensional"
        )
    if b.ndim != 2:
        raise TypeError(
            f"{b.ndim}-dimensional array given. "
            "Array must be one or two-dimensional"
        )
    m, n = a.shape
    dtype = a.dtype
    eps = tf.experimental.numpy.finfo(dtype).eps
    if a.shape == ():
        s = tf.zeros(0, dtype=a.dtype)
        x = tf.zeros((n, *b.shape[1:]), dtype=a.dtype)
    else:
        if rcond is None:
            rcond = eps * max(n, m)
        else:
            rcond = tf.where(rcond < 0, eps, rcond)
        u, s, vt = svd(a, full_matrices=False)
        mask = s >= tf.convert_to_tensor(rcond, dtype=s.dtype) * s[0]
        safe_s = tf.cast(tf.where(mask, s, 1), dtype=a.dtype)
        s_inv = tf.where(mask, 1 / safe_s, 0)[:, tf.newaxis]
        u_t_b = tf.matmul(tf.transpose(tf.math.conj(u)), b)
        x = tf.matmul(tf.transpose(tf.math.conj(vt)), s_inv * u_t_b)

    if b_orig_ndim == 1:
        x = tf.reshape(x, [-1])
    return x
//  [python]
def transform_targets_for_output(y_true, grid_size, anchor_idxs, classes):
    # y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
    N = tf.shape(y_true)[0]

    # y_true_out: (N, grid, grid, anchors, [x, y, w, h, obj, class])
    y_true_out = tf.zeros(
        (N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))

    anchor_idxs = tf.to_int32(anchor_idxs)

    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    for i in tf.range(N):
        for j in tf.range(tf.shape(y_true)[1]):
            if tf.equal(y_true[i][j][2], 0):
                continue
            anchor_eq = tf.equal(
                anchor_idxs, tf.to_int32(y_true[i][j][5]))

            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][0:4]
                box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2

                anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)
                grid_xy = tf.to_int32(box_xy // (1/grid_size))

                # grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
                indexes = indexes.write(
                    idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
                updates = updates.write(
                    idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
                idx += 1

    # tf.print(indexes.stack())
    # tf.print(updates.stack())

    return tf.tensor_scatter_nd_update(
        y_true_out, indexes.stack(), updates.stack())
//  [python]
def transform_targets_for_output(y_true, grid_size, anchor_idxs, classes):
    # y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
    N = tf.shape(y_true)[0]

    # y_true_out: (N, grid, grid, anchors, [x, y, w, h, obj, class])
    y_true_out = tf.zeros(
        (N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))

    anchor_idxs = tf.cast(anchor_idxs, tf.int32)

    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    for i in tf.range(N):
        for j in tf.range(tf.shape(y_true)[1]):
            if tf.equal(y_true[i][j][2], 0):
                continue
            anchor_eq = tf.equal(
                anchor_idxs, tf.cast(y_true[i][j][5], tf.int32))

            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][0:4]
                box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2

                anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)
                grid_xy = tf.cast(box_xy // (1/grid_size), tf.int32)

                # grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
                indexes = indexes.write(
                    idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
                updates = updates.write(
                    idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
                idx += 1

    # tf.print(indexes.stack())
    # tf.print(updates.stack())

    return tf.tensor_scatter_nd_update(
        y_true_out, indexes.stack(), updates.stack())
//  [python]
def contrastive_loss(hidden,
                     num_replicas,
                     normalize_hidden,
                     temperature,
                     model,
                     weight_decay):
  """Computes contrastive loss.

  Args:
    hidden: embedding of video clips after projection head.
    num_replicas: number of distributed replicas.
    normalize_hidden: whether or not to l2 normalize the hidden vector.
    temperature: temperature in the InfoNCE contrastive loss.
    model: keras model for calculating weight decay.
    weight_decay: weight decay parameter.

  Returns:
    A loss scalar.
    The logits for contrastive prediction task.
    The labels for contrastive prediction task.
  """
  large_num = 1e9

  hidden1, hidden2 = tf.split(hidden, num_or_size_splits=2, axis=0)
  if normalize_hidden:
    hidden1 = tf.math.l2_normalize(hidden1, -1)
    hidden2 = tf.math.l2_normalize(hidden2, -1)
  batch_size = tf.shape(hidden1)[0]

  if num_replicas == 1:
    # This is the local version
    hidden1_large = hidden1
    hidden2_large = hidden2
    labels = tf.one_hot(tf.range(batch_size), batch_size * 2)
    masks = tf.one_hot(tf.range(batch_size), batch_size)

  else:
    # This is the cross-tpu version.
    hidden1_large = tpu_cross_replica_concat(hidden1, num_replicas)
    hidden2_large = tpu_cross_replica_concat(hidden2, num_replicas)
    enlarged_batch_size = tf.shape(hidden1_large)[0]
    replica_id = tf.cast(tf.cast(xla.replica_id(), tf.uint32), tf.int32)
    labels_idx = tf.range(batch_size) + replica_id * batch_size
    labels = tf.one_hot(labels_idx, enlarged_batch_size * 2)
    masks = tf.one_hot(labels_idx, enlarged_batch_size)

  logits_aa = tf.matmul(hidden1, hidden1_large, transpose_b=True) / temperature
  logits_aa = logits_aa - tf.cast(masks, logits_aa.dtype) * large_num
  logits_bb = tf.matmul(hidden2, hidden2_large, transpose_b=True) / temperature
  logits_bb = logits_bb - tf.cast(masks, logits_bb.dtype) * large_num
  logits_ab = tf.matmul(hidden1, hidden2_large, transpose_b=True) / temperature
  logits_ba = tf.matmul(hidden2, hidden1_large, transpose_b=True) / temperature

  loss_a = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
      labels, tf.concat([logits_ab, logits_aa], 1)))
  loss_b = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
      labels, tf.concat([logits_ba, logits_bb], 1)))
  loss = loss_a + loss_b

  l2_loss = weight_decay * tf.add_n([
      tf.nn.l2_loss(v)
      for v in model.trainable_variables
      if 'kernel' in v.name
    ])

  total_loss = loss +  tf.cast(l2_loss, loss.dtype)

  contrast_prob = tf.nn.softmax(logits_ab)
  contrast_entropy = - tf.reduce_mean(
      tf.reduce_sum(contrast_prob * tf.math.log(contrast_prob + 1e-8), -1))

  contrast_acc = tf.equal(tf.argmax(labels, 1), tf.argmax(logits_ab, axis=1))
  contrast_acc = tf.reduce_mean(tf.cast(contrast_acc, tf.float32))

  return {
      'total_loss': total_loss,
      'contrastive_loss': loss,
      'reg_loss': l2_loss,
      'contrast_acc': contrast_acc,
      'contrast_entropy': contrast_entropy,
  }
//  [python]
def barlow_twins_loss(
    representations_clean: tf.Tensor,
    representations_corrupted: tf.Tensor,
    *,
    lambda_: Optional[Union[tf.Tensor, float]] = None,
    normalize_batch: bool = True,
) -> tf.Tensor:
  """Barlow Twins loss implementation.

  Implements BarlowTwins loss from the paper https://arxiv.org/abs/2103.03230.
  Note that BarlowTwins loss scale is sensitive to the batch size.

  Args:
    representations_clean: Representations from the clean view of the data.
    representations_corrupted: Representations from the corrupted view of the
      data.
    lambda_: Parameter lambda of the BarlowTwins model. If None (default), uses
      `1 / feature_dim`.
    normalize_batch: If `True` (default), normalizes representations per-batch.

  Returns:
    BarlowTwins loss value as `tf.Tensor`.
  """
  if normalize_batch:
    representations_clean = _normalize(representations_clean)
    representations_corrupted = _normalize(representations_corrupted)
  batch_size, feature_dim = tf.unstack(
      tf.to_float(tf.shape(representations_clean))
  )
  lambda_ = 1 / feature_dim if lambda_ is None else lambda_
  correlation = (
      tf.linalg.matmul(
          representations_clean, representations_corrupted, transpose_a=True
      )
      / batch_size
  )

  loss_matrix = tf.pow(correlation - tf.eye(feature_dim), 2)
  loss_diagonal_sum = tf.linalg.trace(loss_matrix)
  loss_sum = tf.reduce_sum(loss_matrix)
  return (1.0 - lambda_) * loss_diagonal_sum + lambda_ * loss_sum
//  [python]
def barlow_twins_loss(
    representations_clean: tf.Tensor,
    representations_corrupted: tf.Tensor,
    *,
    lambda_: Optional[Union[tf.Tensor, float]] = None,
    normalize_batch: bool = True,
) -> tf.Tensor:
  """Barlow Twins loss implementation.

  Implements BarlowTwins loss from the paper https://arxiv.org/abs/2103.03230.
  Note that BarlowTwins loss scale is sensitive to the batch size.

  Args:
    representations_clean: Representations from the clean view of the data.
    representations_corrupted: Representations from the corrupted view of the
      data.
    lambda_: Parameter lambda of the BarlowTwins model. If None (default), uses
      `1 / feature_dim`.
    normalize_batch: If `True` (default), normalizes representations per-batch.

  Returns:
    BarlowTwins loss value as `tf.Tensor`.
  """
  if normalize_batch:
    representations_clean = _normalize(representations_clean)
    representations_corrupted = _normalize(representations_corrupted)
  batch_size, feature_dim = tf.unstack(
      tf.cast(tf.shape(representations_clean), tf.float32)
  )
  lambda_ = 1 / feature_dim if lambda_ is None else lambda_
  correlation = (
      tf.linalg.matmul(
          representations_clean, representations_corrupted, transpose_a=True
      )
      / batch_size
  )

  loss_matrix = tf.pow(correlation - tf.eye(feature_dim), 2)
  loss_diagonal_sum = tf.linalg.trace(loss_matrix)
  loss_sum = tf.reduce_sum(loss_matrix)
  return (1.0 - lambda_) * loss_diagonal_sum + lambda_ * loss_sum
//  [python]
def encode(self, inputs, attention_bias, training):
    """Generate continuous representation for inputs.

    Args:
      inputs: int tensor with shape [batch_size, input_length].
      attention_bias: float tensor with shape [batch_size, 1, 1, input_length].
      training: boolean, whether in training mode or not.

    Returns:
      float tensor with shape [batch_size, input_length, hidden_size]
    """
    with tf.name_scope("encode"):
      # Prepare inputs to the layer stack by adding positional encodings and
      # applying dropout.
      embedded_inputs = self.embedding_softmax_layer(inputs)
      embedded_inputs = tf.to_float(embedded_inputs)
      inputs_padding = model_utils.get_padding(inputs)
      attention_bias = tf.to_float(attention_bias)

      with tf.name_scope("add_pos_encoding"):
        pos_encoding = self.position_embedding(inputs=embedded_inputs)
        pos_encoding = tf.to_float(pos_encoding)
        encoder_inputs = embedded_inputs + pos_encoding

      if training:
        encoder_inputs = tf.nn.dropout(
            encoder_inputs, rate=self.params["layer_postprocess_dropout"])

      return self.encoder_stack(
          encoder_inputs, attention_bias, inputs_padding, training=training)
//  [python]
def encode(self, inputs, attention_bias, training):
    """Generate continuous representation for inputs.

    Args:
      inputs: int tensor with shape [batch_size, input_length].
      attention_bias: float tensor with shape [batch_size, 1, 1, input_length].
      training: boolean, whether in training mode or not.

    Returns:
      float tensor with shape [batch_size, input_length, hidden_size]
    """
    with tf.name_scope("encode"):
      # Prepare inputs to the layer stack by adding positional encodings and
      # applying dropout.
      embedded_inputs = self.embedding_softmax_layer(inputs)
      embedded_inputs = tf.cast(embedded_inputs, self.params["dtype"])
      inputs_padding = model_utils.get_padding(inputs)
      attention_bias = tf.cast(attention_bias, self.params["dtype"])

      with tf.name_scope("add_pos_encoding"):
        pos_encoding = self.position_embedding(inputs=embedded_inputs)
        pos_encoding = tf.cast(pos_encoding, self.params["dtype"])
        encoder_inputs = embedded_inputs + pos_encoding

      if training:
        encoder_inputs = tf.nn.dropout(
            encoder_inputs, rate=self.params["layer_postprocess_dropout"])

      return self.encoder_stack(
          encoder_inputs, attention_bias, inputs_padding, training=training)
//  [python]
def decode(self, targets, encoder_outputs, attention_bias, training):
    """Generate logits for each value in the target sequence.

    Args:
      targets: target values for the output sequence. int tensor with shape
        [batch_size, target_length]
      encoder_outputs: continuous representation of input sequence. float tensor
        with shape [batch_size, input_length, hidden_size]
      attention_bias: float tensor with shape [batch_size, 1, 1, input_length]
      training: boolean, whether in training mode or not.

    Returns:
      float32 tensor with shape [batch_size, target_length, vocab_size]
    """
    with tf.name_scope("decode"):
      # Prepare inputs to decoder layers by shifting targets, adding positional
      # encoding and applying dropout.
      with tf.name_scope("shift_targets"):
        # Shift targets to the right, and remove the last element
        targets = tf.pad(targets, [[0, 0], [1, 0]])[:, :-1]
      decoder_inputs = self.embedding_softmax_layer(targets)
      decoder_inputs = tf.to_float(decoder_inputs)
      attention_bias = tf.to_float(attention_bias)
      with tf.name_scope("add_pos_encoding"):
        length = tf.shape(decoder_inputs)[1]
        pos_encoding = self.position_embedding(decoder_inputs)
        pos_encoding = tf.to_float(pos_encoding)
        decoder_inputs += pos_encoding
      if training:
        decoder_inputs = tf.nn.dropout(
            decoder_inputs, rate=self.params["layer_postprocess_dropout"])

      # Run values
      decoder_self_attention_bias = model_utils.get_decoder_self_attention_bias(
          length, dtype=self.params["dtype"])
      outputs = self.decoder_stack(
          decoder_inputs,
          encoder_outputs,
          decoder_self_attention_bias,
          attention_bias,
          training=training)
      logits = self.embedding_softmax_layer(outputs, mode="linear")
      logits = tf.to_float(logits)
      return logits
//  [python]
def decode(self, targets, encoder_outputs, attention_bias, training):
    """Generate logits for each value in the target sequence.

    Args:
      targets: target values for the output sequence. int tensor with shape
        [batch_size, target_length]
      encoder_outputs: continuous representation of input sequence. float tensor
        with shape [batch_size, input_length, hidden_size]
      attention_bias: float tensor with shape [batch_size, 1, 1, input_length]
      training: boolean, whether in training mode or not.

    Returns:
      float32 tensor with shape [batch_size, target_length, vocab_size]
    """
    with tf.name_scope("decode"):
      # Prepare inputs to decoder layers by shifting targets, adding positional
      # encoding and applying dropout.
      with tf.name_scope("shift_targets"):
        # Shift targets to the right, and remove the last element
        targets = tf.pad(targets, [[0, 0], [1, 0]])[:, :-1]
      decoder_inputs = self.embedding_softmax_layer(targets)
      decoder_inputs = tf.cast(decoder_inputs, self.params["dtype"])
      attention_bias = tf.cast(attention_bias, self.params["dtype"])
      with tf.name_scope("add_pos_encoding"):
        length = tf.shape(decoder_inputs)[1]
        pos_encoding = self.position_embedding(decoder_inputs)
        pos_encoding = tf.cast(pos_encoding, self.params["dtype"])
        decoder_inputs += pos_encoding
      if training:
        decoder_inputs = tf.nn.dropout(
            decoder_inputs, rate=self.params["layer_postprocess_dropout"])

      # Run values
      decoder_self_attention_bias = model_utils.get_decoder_self_attention_bias(
          length, dtype=self.params["dtype"])
      outputs = self.decoder_stack(
          decoder_inputs,
          encoder_outputs,
          decoder_self_attention_bias,
          attention_bias,
          training=training)
      logits = self.embedding_softmax_layer(outputs, mode="linear")
      logits = tf.cast(logits, tf.float32)
      return logits
//  [python]
def act(self, states, data_aug=False):
    if data_aug and self.num_augmentations > 0:
      states = states[0]
    if self.num_augmentations > 0:
      # use pad of 2 to bump 64 to 68 with 2 + 64 + 2 on each side
      img_pad = 2
      paddings = tf.constant(
          [[0, 0], [img_pad, img_pad], [img_pad, img_pad], [0, 0]],
          dtype=tf.int32)
      states = tf.pad(tf.cast(states * 255., tf.int32), paddings, 'SYMMETRIC')
      states = tf.cast(states, tf.float32) / 255.

    q1, q2 = self.critic(states, stop_grad_features=True, actions=None)
    q = tf.minimum(q1, q2)
    actions = tf.argmax(q, -1)
    return actions
//  [python]
def act(self, states, data_aug=False):
    if data_aug and self.num_augmentations > 0:
      states = states[0]
    if self.num_augmentations > 0:
      # use pad of 2 to bump 64 to 68 with 2 + 64 + 2 on each side
      img_pad = 2
      paddings = tf.constant(
          [[0, 0], [img_pad, img_pad], [img_pad, img_pad], [0, 0]],
          dtype=tf.int32)
      states = tf.cast(
          tf.pad(tf.cast(states * 255., tf.int32), paddings, 'SYMMETRIC'),
          tf.float32) / 255.

    q1, q2 = self.critic(states, stop_grad_features=True, actions=None)
    q = tf.minimum(q1, q2)
    actions = tf.argmax(q, -1)
    return actions
//  [python]
def _get_symbols_to_logits_fn(self, max_decode_length, training):
    """Returns a decoding function that calculates logits of the next tokens."""
    timing_signal = self.position_embedding(
        inputs=None, length=max_decode_length + 1)
    timing_signal = t.cast(timing_signal, self.params["dtype"])
    decoder_self_attention_bias = model_utils.get_decoder_self_attention_bias(
        max_decode_length, dtype=self.params["dtype"])

    def symbols_to_logits_fn(ids, i, cache):
      """Generate logits for next potential IDs.

      Args:
        ids: Current decoded sequences. int tensor with shape [batch_size *
          beam_size, i + 1].
        i: Loop index.
        cache: dictionary of values storing the encoder output, encoder-decoder
          attention bias, and previous decoder attention values.

      Returns:
        Tuple of
          (logits with shape [batch_size * beam_size, vocab_size],
           updated cache values)
      """
      # Set decoder input to the last generated IDs
      decoder_input = ids[:, -1:]

      # Preprocess decoder input by getting embeddings and adding timing signal.
      decoder_input = self.embedding_softmax_layer(decoder_input)
      decoder_input += timing_signal[i]
      if self.params["padded_decode"]:
        bias_shape = decoder_self_attention_bias.shape.as_list()
        self_attention_bias = tf.slice(
            decoder_self_attention_bias, [0, 0, i, 0],
            [bias_shape[0], bias_shape[1], 1, bias_shape[3]])
      else:
        self_attention_bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]

      decoder_outputs = self.decoder_stack(
          decoder_input,
          cache.get("encoder_outputs"),
          self_attention_bias,
          cache.get("encoder_decoder_attention_bias"),
          training=training,
          cache=cache,
          decode_loop_step=i if self.params["padded_decode"] else None)
      logits = self.embedding_softmax_layer(decoder_outputs, mode="linear")
      logits = tf.squeeze(logits, axis=[1])
      return logits, cache

    return symbols_to_logits_fn
//  [python]
def _get_symbols_to_logits_fn(self, max_decode_length, training):
    """Returns a decoding function that calculates logits of the next tokens."""
    timing_signal = self.position_embedding(
        inputs=None, length=max_decode_length + 1)
    timing_signal = tf.cast(timing_signal, self.params["dtype"])
    decoder_self_attention_bias = model_utils.get_decoder_self_attention_bias(
        max_decode_length, dtype=self.params["dtype"])

    def symbols_to_logits_fn(ids, i, cache):
      """Generate logits for next potential IDs.

      Args:
        ids: Current decoded sequences. int tensor with shape [batch_size *
          beam_size, i + 1].
        i: Loop index.
        cache: dictionary of values storing the encoder output, encoder-decoder
          attention bias, and previous decoder attention values.

      Returns:
        Tuple of
          (logits with shape [batch_size * beam_size, vocab_size],
           updated cache values)
      """
      # Set decoder input to the last generated IDs
      decoder_input = ids[:, -1:]

      # Preprocess decoder input by getting embeddings and adding timing signal.
      decoder_input = self.embedding_softmax_layer(decoder_input)
      decoder_input += timing_signal[i]
      if self.params["padded_decode"]:
        bias_shape = decoder_self_attention_bias.shape.as_list()
        self_attention_bias = tf.slice(
            decoder_self_attention_bias, [0, 0, i, 0],
            [bias_shape[0], bias_shape[1], 1, bias_shape[3]])
      else:
        self_attention_bias = decoder_self_attention_bias[:, :, i:i + 1, :i + 1]

      decoder_outputs = self.decoder_stack(
          decoder_input,
          cache.get("encoder_outputs"),
          self_attention_bias,
          cache.get("encoder_decoder_attention_bias"),
          training=training,
          cache=cache,
          decode_loop_step=i if self.params["padded_decode"] else None)
      logits = self.embedding_softmax_layer(decoder_outputs, mode="linear")
      logits = tf.squeeze(logits, axis=[1])
      return logits, cache

    return symbols_to_logits_fn
//  [python]
def update_step(self,
                  replay_buffer_iter,
                  train_target='both'):
    """Performs a single training step for critic and embedding.

    Args:
      replay_buffer_iter: A tensorflow graph iteratable object.
      train_target: string specifying whether update RL and or representation

    Returns:
      Dictionary with losses to track.
    """
    transition = next(replay_buffer_iter)
    numpy_dataset = isinstance(replay_buffer_iter, np.ndarray)
    # observation: n_batch x n_timesteps x 1 x H*W*3*n_frames x 1 ->
    # n_batch x H x W x 3*n_frames
    if not numpy_dataset:
      states = transition.observation[:, 0]
      next_states = transition.observation[:, 1]
      actions = transition.action[:, 0]
      rewards = transition.reward[:, 0]
      discounts = transition.discount[:, 0]

      if transition.observation.dtype == tf.uint8:
        states = tf.to_float(states) / 255.
        next_states = tf.to_float(next_states) / 255.
    else:
      states, actions, rewards, next_states, discounts = transition

    if self.num_augmentations > 0:
      states, next_states = tf_utils.image_aug(
          states,
          next_states,
          img_pad=4,
          num_augmentations=self.num_augmentations,
          obs_dim=64,
          channels=3,
          cropped_shape=[self.batch_size, 68, 68, 3])

    next_actions = self.act(next_states, data_aug=True)

    if train_target == 'both':
      ssl_dict = self.fit_embedding(states, actions, next_states, next_actions,
                                    rewards, discounts)
      critic_dict = self.fit_critic(states, actions, next_states, next_actions,
                                    rewards, discounts)
    elif train_target == 'encoder':
      ssl_dict = self.fit_embedding(states, actions, next_states, next_actions,
                                    rewards, discounts)
      critic_dict = {}
    elif train_target == 'rl':
      ssl_dict = {}
      critic_dict = self.fit_critic(states, actions, next_states, next_actions,
                                    rewards, discounts)

    return {**ssl_dict, **critic_dict}
//  [python]
def update_step(self,
                  replay_buffer_iter,
                  train_target='both'):
    """Performs a single training step for critic and embedding.

    Args:
      replay_buffer_iter: A tensorflow graph iteratable object.
      train_target: string specifying whether update RL and or representation

    Returns:
      Dictionary with losses to track.
    """
    transition = next(replay_buffer_iter)
    numpy_dataset = isinstance(replay_buffer_iter, np.ndarray)
    # observation: n_batch x n_timesteps x 1 x H*W*3*n_frames x 1 ->
    # n_batch x H x W x 3*n_frames
    if not numpy_dataset:
      states = transition.observation[:, 0]
      next_states = transition.observation[:, 1]
      actions = transition.action[:, 0]
      rewards = transition.reward[:, 0]
      discounts = transition.discount[:, 0]

      if transition.observation.dtype == tf.uint8:
        states = tf.cast(states, tf.float32) / 255.
        next_states = tf.cast(next_states, tf.float32) / 255.
    else:
      states, actions, rewards, next_states, discounts = transition

    if self.num_augmentations > 0:
      states, next_states = tf_utils.image_aug(
          states,
          next_states,
          img_pad=4,
          num_augmentations=self.num_augmentations,
          obs_dim=64,
          channels=3,
          cropped_shape=[self.batch_size, 68, 68, 3])

    next_actions = self.act(next_states, data_aug=True)

    if train_target == 'both':
      ssl_dict = self.fit_embedding(states, actions, next_states, next_actions,
                                    rewards, discounts)
      critic_dict = self.fit_critic(states, actions, next_states, next_actions,
                                    rewards, discounts)
    elif train_target == 'encoder':
      ssl_dict = self.fit_embedding(states, actions, next_states, next_actions,
                                    rewards, discounts)
      critic_dict = {}
    elif train_target == 'rl':
      ssl_dict = {}
      critic_dict = self.fit_critic(states, actions, next_states, next_actions,
                                    rewards, discounts)

    return {**ssl_dict, **critic_dict}
//  [python]
def make_source_and_target_mask(self, source: tf.Tensor, target: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        # [1, 1, 1, ..., 0, 0, 0] -> 0 is masked
        # (N, 1, 1, max_length)
        target_pad_mask = tf.to_int32(tf.math.not_equal(target, self.vocab_size + 2))
        target_pad_mask = target_pad_mask[:, tf.newaxis, tf.newaxis, :]
        target_length = target.shape[1]
        # sub mask filled diagonal with 1 = see 0 = masked (max_length, max_length)
        target_sub_mask = tf.linalg.band_part(tf.ones((target_length, target_length)), -1, 0)
        # source mask filled with ones (max_length, positional_encoded_seq_len)
        source_mask = tf.ones((target_length, source.shape[1]))
        # combine the two masks into one boolean mask where False is masked (N, 1, max_length, max_length)
        target_mask = tf.math.logical_and(
            tf.to_bool(target_sub_mask), tf.to_bool(target_pad_mask)
        )
        return source_mask, target_mask
//  [python]
def make_source_and_target_mask(self, source: tf.Tensor, target: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        # [1, 1, 1, ..., 0, 0, 0] -> 0 is masked
        # (N, 1, 1, max_length)
        target_pad_mask = tf.cast(tf.math.not_equal(target, self.vocab_size + 2), dtype=tf.uint8)
        target_pad_mask = target_pad_mask[:, tf.newaxis, tf.newaxis, :]
        target_length = target.shape[1]
        # sub mask filled diagonal with 1 = see 0 = masked (max_length, max_length)
        target_sub_mask = tf.linalg.band_part(tf.ones((target_length, target_length)), -1, 0)
        # source mask filled with ones (max_length, positional_encoded_seq_len)
        source_mask = tf.ones((target_length, source.shape[1]))
        # combine the two masks into one boolean mask where False is masked (N, 1, max_length, max_length)
        target_mask = tf.math.logical_and(
            tf.cast(target_sub_mask, dtype=tf.bool), tf.cast(target_pad_mask, dtype=tf.bool)
        )
        return source_mask, target_mask
//  [python]
def tower_loss(num_boxes, np_boxes_5, np_images_features, np_idx_class_labels, np_len_class_labels,
               np_idx_query_, len_query_, label_list_, segment_ids, is_training, label_query, weight_label_query, reuse_variables):
    # Build inference graph
    with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
        model_loss, probs, loss_list, probs_all = model.model_attention_channel_e(num_boxes, np_boxes_5, np_images_features, np_idx_class_labels, np_len_class_labels,
               np_idx_query_, len_query_, label_list_, segment_ids, label_query, weight_label_query, is_training=is_training, reuse=reuse_variables)
    #acc = tf.reduce_mean(tf.cast(probs[:,0] > 0.5, tf.float32))
    preds = tf.argmax(probs, axis=1)
    #num_labels = tf.cast(tf.size(weight_label_query), tf.float32) / tf.reduce_sum(weight_label_query)
    #acc = tf.reduce_mean(tf.cast(tf.equal(preds, label_query), dtype=tf.float32) * weight_label_query) * num_labels
    #acc = tf.reduce_mean(tf.cast(tf.equal(preds, label_list_), dtype=tf.float32))
    acc = tf.to_float(tf.equal(preds, label_list_))
    total_loss = tf.add_n([model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))

    if reuse_variables is None:
        tf.summary.scalar('total_loss', total_loss)

    return total_loss, model_loss, acc, probs, loss_list, probs_all
//  [python]
def tower_loss(num_boxes, np_boxes_5, np_images_features, np_idx_class_labels, np_len_class_labels,
               np_idx_query_, label_list_, len_query_, segment_ids, is_training, label_query, weight_label_query, reuse_variables):
    # Build inference graph
    with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
        model_loss, probs, loss_list, probs_all = model.model_attention_channel_e(num_boxes, np_boxes_5, np_images_features, np_idx_class_labels, np_len_class_labels,
               np_idx_query_, len_query_, label_list_, segment_ids, label_query, weight_label_query, is_training=is_training, reuse=reuse_variables)
    #acc = tf.reduce_mean(tf.cast(probs[:,0] > 0.5, tf.float32))
    preds = tf.argmax(probs, axis=1)
    #num_labels = tf.cast(tf.size(weight_label_query), tf.float32) / tf.reduce_sum(weight_label_query)
    #acc = tf.reduce_mean(tf.cast(tf.equal(preds, label_query), dtype=tf.float32) * weight_label_query) * num_labels
    #acc = tf.reduce_mean(tf.cast(tf.equal(preds, label_list_), dtype=tf.float32))
    acc = tf.cast(tf.equal(preds, label_list_), dtype=tf.float32)
    total_loss = tf.add_n([model_loss] + tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))

    if reuse_variables is None:
        tf.summary.scalar('total_loss', total_loss)

    return total_loss, model_loss, acc, probs, loss_list, probs_all
//  [python]
def predict(self, encoder_outputs, encoder_decoder_attention_bias, training):
    """Return predicted sequence."""
    encoder_outputs = t.cast(encoder_outputs, self.params["dtype"])
    if self.params["padded_decode"]:
      batch_size = encoder_outputs.shape.as_list()[0]
      input_length = encoder_outputs.shape.as_list()[1]
    else:
      batch_size = tf.shape(encoder_outputs)[0]
      input_length = tf.shape(encoder_outputs)[1]
    max_decode_length = input_length + self.params["extra_decode_length"]
    encoder_decoder_attention_bias = t.cast(encoder_decoder_attention_bias,
                                             self.params["dtype"])

    symbols_to_logits_fn = self._get_symbols_to_logits_fn(
        max_decode_length, training)

    # Create initial set of IDs that will be passed into symbols_to_logits_fn.
    initial_ids = tf.zeros([batch_size], dtype=tf.int32)

    # Create cache storing decoder attention values for each layer.
    # pylint: disable=g-complex-comprehension
    init_decode_length = (
        max_decode_length if self.params["padded_decode"] else 0)
    num_heads = self.params["num_heads"]
    dim_per_head = self.params["hidden_size"] // num_heads
    cache = {
        "layer_%d" % layer: {
            "k":
                t.zeros(
                    [batch_size, init_decode_length, num_heads, dim_per_head],
                    dtype=self.params["dtype"]),
            "v":
                t.zeros(
                    [batch_size, init_decode_length, num_heads, dim_per_head],
                    dtype=self.params["dtype"])
        } for layer in range(self.params["num_hidden_layers"])
    }
    # pylint: enable=g-complex-comprehension

    # Add encoder output and attention bias to the cache.
    cache["encoder_outputs"] = encoder_outputs
    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias

    # Use beam search to find the top beam_size sequences and scores.
    decoded_ids, scores = beam_search.sequence_beam_search(
        symbols_to_logits_fn=symbols_to_logits_fn,
        initial_ids=initial_ids,
        initial_cache=cache,
        vocab_size=self.params["vocab_size"],
        beam_size=self.params["beam_size"],
        alpha=self.params["alpha"],
        max_decode_length=max_decode_length,
        eos_id=EOS_ID,
        padded_decode=self.params["padded_decode"],
        dtype=self.params["dtype"])

    # Get the top sequence for each batch element
    top_decoded_ids = decoded_ids[:, 0, 1:]
    top_scores = scores[:, 0]

    return {"outputs": top_decoded_ids, "scores": top_scores}
//  [python]
def predict(self, encoder_outputs, encoder_decoder_attention_bias, training):
    """Return predicted sequence."""
    encoder_outputs = tf.cast(encoder_outputs, self.params["dtype"])
    if self.params["padded_decode"]:
      batch_size = encoder_outputs.shape.as_list()[0]
      input_length = encoder_outputs.shape.as_list()[1]
    else:
      batch_size = tf.shape(encoder_outputs)[0]
      input_length = tf.shape(encoder_outputs)[1]
    max_decode_length = input_length + self.params["extra_decode_length"]
    encoder_decoder_attention_bias = tf.cast(encoder_decoder_attention_bias,
                                             self.params["dtype"])

    symbols_to_logits_fn = self._get_symbols_to_logits_fn(
        max_decode_length, training)

    # Create initial set of IDs that will be passed into symbols_to_logits_fn.
    initial_ids = tf.zeros([batch_size], dtype=tf.int32)

    # Create cache storing decoder attention values for each layer.
    # pylint: disable=g-complex-comprehension
    init_decode_length = (
        max_decode_length if self.params["padded_decode"] else 0)
    num_heads = self.params["num_heads"]
    dim_per_head = self.params["hidden_size"] // num_heads
    cache = {
        "layer_%d" % layer: {
            "k":
                tf.zeros(
                    [batch_size, init_decode_length, num_heads, dim_per_head],
                    dtype=self.params["dtype"]),
            "v":
                tf.zeros(
                    [batch_size, init_decode_length, num_heads, dim_per_head],
                    dtype=self.params["dtype"])
        } for layer in range(self.params["num_hidden_layers"])
    }
    # pylint: enable=g-complex-comprehension

    # Add encoder output and attention bias to the cache.
    cache["encoder_outputs"] = encoder_outputs
    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias

    # Use beam search to find the top beam_size sequences and scores.
    decoded_ids, scores = beam_search.sequence_beam_search(
        symbols_to_logits_fn=symbols_to_logits_fn,
        initial_ids=initial_ids,
        initial_cache=cache,
        vocab_size=self.params["vocab_size"],
        beam_size=self.params["beam_size"],
        alpha=self.params["alpha"],
        max_decode_length=max_decode_length,
        eos_id=EOS_ID,
        padded_decode=self.params["padded_decode"],
        dtype=self.params["dtype"])

    # Get the top sequence for each batch element
    top_decoded_ids = decoded_ids[:, 0, 1:]
    top_scores = scores[:, 0]

    return {"outputs": top_decoded_ids, "scores": top_scores}