jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def reshape(vector: tf.Tensor,
            shapes: List[Tuple],
            dtypes: Optional[List[Type]] = None) -> List[tf.Tensor]:
    """
    Split and reshape a vector into a list of TF tensors.

    :param vector:
        A 1-dimensional tensor to split and reshape.
    :param shapes:
        A list of tuples of integers, representing the shapes of multiple
        target weights to construct.
    :param dtypes:
        A list of types for the new weights.
    :return:
        A list of TF tensors constructed from the inputs.
    """
    separating_index = 0
    weights = []
    for i, shape in enumerate(shapes):
        weight_len = np.prod(shape)
        new_flat_weight = vector[separating_index:separating_index +
                                 weight_len]
        if dtypes is not None:
            new_flat_weight = tf.cast(new_flat_weight, dtypes[i])
        weights.append(tf.reshape(new_flat_weight, shape))
        separating_index += weight_len
    return weights
//  [python]
def one_hot(indices: tf.Tensor, depth: int) -> np.ndarray:
    """
    One-hot encode indices to vector with depth dimension.

    :param indices:
        A vector of indices to be one-hot encoded.
    :param depth:
        The dimension of one-hot encoding.
    :return:
        One-hot encoded vectors.
    """
    return tf.one_hot(tf.to_int32(indices), depth, dtype=tf.float32)
//  [python]
def one_hot(indices: tf.Tensor, depth: int) -> np.ndarray:
    """
    One-hot encode indices to vector with depth dimension.

    :param indices:
        A vector of indices to be one-hot encoded.
    :param depth:
        The dimension of one-hot encoding.
    :return:
        One-hot encoded vectors.
    """
    return tf.one_hot(tf.cast(indices, tf.int32), depth, dtype=tf.float32)
//  [python]
def gumbel_softmax(logits, temperature, hard=True):
    """Sample from the Gumbel-Softmax distribution and optionally discretize.
    Args:
      logits: [batch_size, n_class] unnormalized log-probs
      temperature: non-negative scalar
      hard: if True, take argmax, but differentiate w.r.t. soft sample y
    Returns:
      [batch_size, n_class] sample from the Gumbel-Softmax distribution.
      If hard=True, then the returned sample will be one-hot, otherwise it will
      be a probabilitiy distribution that sums to 1 across classes
    """
    y = gumbel_softmax_sample(logits, temperature)
    if hard:
        k = tf.shape(logits)[-1]
        #y_hard = tf.cast(tf.one_hot(tf.argmax(y,1),k), y.dtype)
        y_hard = t.cast(tf.equal(y, tf.reduce_max(y, 1, keepdims=True)), y.dtype)
        y = tf.stop_gradient(y_hard - y) + y
    return y
//  [python]
def gumbel_softmax(logits, temperature, hard=True):
    """Sample from the Gumbel-Softmax distribution and optionally discretize.
    Args:
      logits: [batch_size, n_class] unnormalized log-probs
      temperature: non-negative scalar
      hard: if True, take argmax, but differentiate w.r.t. soft sample y
    Returns:
      [batch_size, n_class] sample from the Gumbel-Softmax distribution.
      If hard=True, then the returned sample will be one-hot, otherwise it will
      be a probabilitiy distribution that sums to 1 across classes
    """
    y = gumbel_softmax_sample(logits, temperature)
    if hard:
        k = tf.shape(logits)[-1]
        #y_hard = tf.cast(tf.one_hot(tf.argmax(y,1),k), y.dtype)
        y_hard = tf.cast(tf.equal(y, tf.reduce_max(y, 1, keepdims=True)), y.dtype)
        y = tf.stop_gradient(y_hard - y) + y
    return y
//  [python]
def call(self, za: FloatTensor, zb: FloatTensor) -> FloatTensor:
        """Compute the lost.

        Args:
            za: Embedding A
            zb: Embedding B

        Returns:
            loss
        """
        # compute the diagonal
        batch_size = tf.shape(za)[0]

        za = self.standardize_columns(za)
        zb = self.standardize_columns(zb)

        # compute pairwise
        c = tf.matmul(za, zb, transpose_a=True)
        c = c / tf.cast(batch_size, dtype=c.dtype)

        on_diag = 1.0 - tf.linalg.diag_part(c)
        on_diag = tf.math.pow(on_diag, 2)
        on_diag = tf.math.reduce_sum(on_diag)

        off_diag = self.off_diagonal(c)
        off_diag = tf.math.pow(off_diag, 2)
        off_diag = tf.math.reduce_sum(off_diag)

        # 1D Tensor
        loss: FloatTensor = off_diag * self.lambda_ + on_diag + self.margin

        return loss
//  [python]
def discounted_rewards(reward, done, gae_gamma, end_values):
  """Discounted rewards."""
  not_done = tf.expand_dims(1 - t.cast(done, tf.float32), axis=2)
  end_values = end_values * not_done[-1, :, :]
  return_ = tf.scan(
      lambda agg, cur: cur + gae_gamma * agg,
      tf.expand_dims(reward, axis=2) * not_done,
      initializer=end_values,
      reverse=True,
      back_prop=False,
      parallel_iterations=2)
  return tf.check_numerics(return_, "return")
//  [python]
def discounted_rewards(reward, done, gae_gamma, end_values):
  """Discounted rewards."""
  not_done = tf.expand_dims(1 - tf.cast(done, tf.float32), axis=2)
  end_values = end_values * not_done[-1, :, :]
  return_ = tf.scan(
      lambda agg, cur: cur + gae_gamma * agg,
      tf.expand_dims(reward, axis=2) * not_done,
      initializer=end_values,
      reverse=True,
      back_prop=False,
      parallel_iterations=2)
  return tf.check_numerics(return_, "return")
//  [python]
def _tiled_global_box_search(self,
                               pred_boxes,
                               pred_classes,
                               boxes,
                               classes,
                               true_conf,
                               smoothed,
                               scale=None):
    """Search of all groundtruths to associate groundtruths to predictions."""

    boxes = box_ops.yxyx_to_xcycwh(boxes)

    if scale is not None:
      boxes = boxes * t.cast(t.stop_gradient(scale), boxes.dtype)

    # Search all predictions against ground truths to find mathcing boxes for
    # each pixel.
    _, _, iou_max, _ = self._search_pairs(pred_boxes, pred_classes, boxes,
                                          classes)

    if iou_max is None:
      return true_conf, t.ones_like(true_conf)

    # Find the exact indexes to ignore and keep.
    ignore_mask = t.cast(iou_max < self._ignore_thresh, pred_boxes.dtype)
    iou_mask = iou_max > self._ignore_thresh

    if not smoothed:
      # Ignore all pixels where a box was not supposed to be predicted but a
      # high confidence box was predicted.
      obj_mask = true_conf + (1 - true_conf) * ignore_mask
    else:
      # Replace pixels in the tre confidence map with the max iou predicted
      # with in that cell.
      obj_mask = t.ones_like(true_conf)
      iou_ = (1 - self._objectness_smooth) + self._objectness_smooth * iou_max
      iou_ = t.where(iou_max > 0, iou_, t.zeros_like(iou_))
      true_conf = t.where(iou_mask, iou_, true_conf)

    # Stop gradient so while loop is not tracked.
    obj_mask = t.stop_gradient(obj_mask)
    true_conf = t.stop_gradient(true_conf)
    return true_conf, obj_mask
//  [python]
def _tiled_global_box_search(self,
                               pred_boxes,
                               pred_classes,
                               boxes,
                               classes,
                               true_conf,
                               smoothed,
                               scale=None):
    """Search of all groundtruths to associate groundtruths to predictions."""

    boxes = box_ops.yxyx_to_xcycwh(boxes)

    if scale is not None:
      boxes = boxes * tf.cast(tf.stop_gradient(scale), boxes.dtype)

    # Search all predictions against ground truths to find mathcing boxes for
    # each pixel.
    _, _, iou_max, _ = self._search_pairs(pred_boxes, pred_classes, boxes,
                                          classes)

    if iou_max is None:
      return true_conf, tf.ones_like(true_conf)

    # Find the exact indexes to ignore and keep.
    ignore_mask = tf.cast(iou_max < self._ignore_thresh, pred_boxes.dtype)
    iou_mask = iou_max > self._ignore_thresh

    if not smoothed:
      # Ignore all pixels where a box was not supposed to be predicted but a
      # high confidence box was predicted.
      obj_mask = true_conf + (1 - true_conf) * ignore_mask
    else:
      # Replace pixels in the tre confidence map with the max iou predicted
      # with in that cell.
      obj_mask = tf.ones_like(true_conf)
      iou_ = (1 - self._objectness_smooth) + self._objectness_smooth * iou_max
      iou_ = tf.where(iou_max > 0, iou_, tf.zeros_like(iou_))
      true_conf = tf.where(iou_mask, iou_, true_conf)

    # Stop gradient so while loop is not tracked.
    obj_mask = tf.stop_gradient(obj_mask)
    true_conf = tf.stop_gradient(true_conf)
    return true_conf, obj_mask
//  [python]
def normalize_embeds(embeds):
    return embeds * (tf.to_float(tf.shape(embeds)[-1]) ** -0.5)
//  [python]
def normalize_embeds(embeds):
    return embeds * (tf.cast(tf.shape(embeds)[-1], dtype=tf.float32) ** -0.5)
//  [python]
def calculate_generalized_advantage_estimator(
    reward, value, done, gae_gamma, gae_lambda):
  # pylint: disable=g-doc-args
  """Generalized advantage estimator.

  Returns:
    GAE estimator. It will be one element shorter than the input; this is
    because to compute GAE for [0, ..., N-1] one needs V for [1, ..., N].
  """
  # pylint: enable=g-doc-args

  next_value = value[1:, :]
  next_not_done = 1 - t.cast(done[1:, :], t.float32)
  delta = (reward[:-1, :] + gae_gamma * next_value * next_not_done
           - value[:-1, :])

  return_ = tf.reverse(tf.scan(
      lambda agg, cur: cur[0] + cur[1] * gae_gamma * gae_lambda * agg,
      [tf.reverse(delta, [0]), tf.reverse(next_not_done, [0])],
      tf.zeros_like(delta[0, :]),
      parallel_iterations=1), [0])
  return tf.check_numerics(return_, "return")
//  [python]
def calculate_generalized_advantage_estimator(
    reward, value, done, gae_gamma, gae_lambda):
  # pylint: disable=g-doc-args
  """Generalized advantage estimator.

  Returns:
    GAE estimator. It will be one element shorter than the input; this is
    because to compute GAE for [0, ..., N-1] one needs V for [1, ..., N].
  """
  # pylint: enable=g-doc-args

  next_value = value[1:, :]
  next_not_done = 1 - tf.cast(done[1:, :], tf.float32)
  delta = (reward[:-1, :] + gae_gamma * next_value * next_not_done
           - value[:-1, :])

  return_ = tf.reverse(tf.scan(
      lambda agg, cur: cur[0] + cur[1] * gae_gamma * gae_lambda * agg,
      [tf.reverse(delta, [0]), tf.reverse(next_not_done, [0])],
      tf.zeros_like(delta[0, :]),
      parallel_iterations=1), [0])
  return tf.check_numerics(return_, "return")
//  [python]
def goemotions(features, runner_config):
  labels = runner_config["model_config"]["labels"]
  label_tensor = tf.stack([features[label] for label in labels], axis=1)
  return features["comment_text"], t.cast(label_tensor, tf.float32)
//  [python]
def goemotions(features, runner_config):
  labels = runner_config["model_config"]["labels"]
  label_tensor = tf.stack([features[label] for label in labels], axis=1)
  return features["comment_text"], tf.cast(label_tensor, tf.float32)
//  [python]
def objective(self, params, data, labels):
    """Computes the softmax cross entropy.

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      loss: Softmax cross entropy loss averaged over the samples in the batch

    Raises:
      ValueError: If the objective is to be computed over >2 classes, because
        this operation is broken in tensorflow at the moment.
    """
    # Forward pass.
    # pdb.set_trace()
    logits = self.inference(params, data)
    print(logits.get_shape)
    # Compute the loss.
    # l2reg = [tf.reduce_sum(param ** 2) for param in params]
    if int(logits.get_shape()[1]) == 2:
      labels = tf.to_float(labels)
      losses = tf.nn.sigmoid_cross_entropy_with_logits(
        labels=labels, logits=logits[:, 0])
    else:
      # labels = tf.squeeze(labels)
      losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)

    return tf.reduce_mean(losses)
//  [python]
def objective(self, params, data, labels):
    """Computes the softmax cross entropy.

    Args:
      params: List of parameter tensors or variables
      data: Batch of features with samples along the first dimension
      labels: Vector of labels with the same number of samples as the data

    Returns:
      loss: Softmax cross entropy loss averaged over the samples in the batch

    Raises:
      ValueError: If the objective is to be computed over >2 classes, because
        this operation is broken in tensorflow at the moment.
    """
    # Forward pass.
    # pdb.set_trace()
    logits = self.inference(params, data)
    print(logits.get_shape)
    # Compute the loss.
    # l2reg = [tf.reduce_sum(param ** 2) for param in params]
    if int(logits.get_shape()[1]) == 2:
      labels = tf.cast(labels, tf.float32)
      losses = tf.nn.sigmoid_cross_entropy_with_logits(
        labels=labels, logits=logits[:, 0])
    else:
      # labels = tf.squeeze(labels)
      losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels)

    return tf.reduce_mean(losses)
//  [python]
def cross_entropy_loss(prediction_logits, labels):
    prediction_logits = tf.reshape(prediction_logits, shape=[-1, 10])
    labels = tf.reshape(labels, shape=[-1, 1])
    cro_ent = tf.reduce_sum(tf.keras.losses.sparse_categorical_crossentropy(
        labels, prediction_logits, from_logits=True))
    return cro_ent
//  [python]
def cross_entropy_loss(prediction_logits, labels):
    prediction_logits = tf.cast(tf.reshape(prediction_logits, shape=[-1, 10]),
                                dtype=tf.float32)
    labels = tf.cast(tf.reshape(labels, shape=[-1, 1]), dtype=tf.float32)
    cro_ent = tf.reduce_sum(tf.keras.losses.sparse_categorical_crossentropy(
        labels, prediction_logits, from_logits=True))
    return cro_ent
//  [python]
def get_keypoint_v0(image, predictions, heatmap_size, height, width, category, clip_at_zero=True, data_format='channels_last', name=None):
    predictions = tf.reshape(predictions, [1, -1, heatmap_size*heatmap_size])

    pred_max = tf.reduce_max(predictions, axis=-1)
    pred_indices = tf.argmax(predictions, axis=-1)
    pred_x, pred_y = tf.to_int32(tf.floormod(pred_indices, heatmap_size)), tf.to_int32(tf.floordiv(pred_indices, heatmap_size))

    width, height = tf.to_float(width), tf.to_float(height)
    pred_x, pred_y = tf.to_float(pred_x) * width / tf.to_float(heatmap_size), tf.to_float(pred_y) * height / tf.to_float(heatmap_size)

    if clip_at_zero:
      pred_x, pred_y =  pred_x * tf.to_float(pred_max>0), pred_y * tf.to_float(pred_max>0)
      pred_x = pred_x * tf.to_float(pred_max>0) + tf.to_float(pred_max<=0) * (width / 2.)
      pred_y = pred_y * tf.to_float(pred_max>0) + tf.to_float(pred_max<=0) * (height / 2.)

    if config.PRED_DEBUG:
      pred_indices_ = tf.squeeze(pred_indices)
      image_ = tf.squeeze(image) * 255.
      pred_heatmap = tf.one_hot(pred_indices_, heatmap_size*heatmap_size, on_value=255, off_value=0, axis=-1, dtype=tf.int32)

      pred_heatmap = tf.reshape(pred_heatmap, [-1, heatmap_size, heatmap_size])
      if data_format == 'channels_first':
        image_ = tf.transpose(image_, perm=(1, 2, 0))
      save_image_op = tf.py_func(save_image_with_heatmap,
                                  [image_, height, width,
                                  heatmap_size,
                                  pred_heatmap,
                                  tf.reshape(predictions, [-1, heatmap_size, heatmap_size]),
                                  config.left_right_group_map[category][0],
                                  config.left_right_group_map[category][1],
                                  config.left_right_group_map[category][2]],
                                  tf.int64, stateful=True)
      with tf.control_dependencies([save_image_op]):
        pred_x, pred_y = pred_x * 1., pred_y * 1.
    return pred_x, pred_y
//  [python]
def get_keypoint_v0(image, predictions, heatmap_size, height, width, category, clip_at_zero=True, data_format='channels_last', name=None):
    predictions = tf.reshape(predictions, [1, -1, heatmap_size*heatmap_size])

    pred_max = tf.reduce_max(predictions, axis=-1)
    pred_indices = tf.argmax(predictions, axis=-1)
    pred_x, pred_y = tf.cast(tf.floormod(pred_indices, heatmap_size), tf.float32), tf.cast(tf.floordiv(pred_indices, heatmap_size), tf.float32)

    width, height = tf.cast(width, tf.float32), tf.cast(height, tf.float32)
    pred_x, pred_y = pred_x * width / tf.cast(heatmap_size, tf.float32), pred_y * height / tf.cast(heatmap_size, tf.float32)

    if clip_at_zero:
      pred_x, pred_y =  pred_x * tf.cast(pred_max>0, tf.float32), pred_y * tf.cast(pred_max>0, tf.float32)
      pred_x = pred_x * tf.cast(pred_max>0, tf.float32) + tf.cast(pred_max<=0, tf.float32) * (width / 2.)
      pred_y = pred_y * tf.cast(pred_max>0, tf.float32) + tf.cast(pred_max<=0, tf.float32) * (height / 2.)

    if config.PRED_DEBUG:
      pred_indices_ = tf.squeeze(pred_indices)
      image_ = tf.squeeze(image) * 255.
      pred_heatmap = tf.one_hot(pred_indices_, heatmap_size*heatmap_size, on_value=255, off_value=0, axis=-1, dtype=tf.int32)

      pred_heatmap = tf.reshape(pred_heatmap, [-1, heatmap_size, heatmap_size])
      if data_format == 'channels_first':
        image_ = tf.transpose(image_, perm=(1, 2, 0))
      save_image_op = tf.py_func(save_image_with_heatmap,
                                  [image_, height, width,
                                  heatmap_size,
                                  pred_heatmap,
                                  tf.reshape(predictions, [-1, heatmap_size, heatmap_size]),
                                  config.left_right_group_map[category][0],
                                  config.left_right_group_map[category][1],
                                  config.left_right_group_map[category][2]],
                                  tf.int64, stateful=True)
      with tf.control_dependencies([save_image_op]):
        pred_x, pred_y = pred_x * 1., pred_y * 1.
    return pred_x, pred_y
//  [python]
def create_input_fn(runner_config, mode, drop_remainder):
  """Returns an input function to use in the instantiation of tf.estimator.*."""

  def _post_processor(features, batch_size):
    """Post process the data to a form expected by model_fn."""
    data_processor = getattr(sys.modules[__name__], runner_config["dataset"])
    text, label = data_processor(features, runner_config)
    model_config = runner_config["model_config"]
    if "max_seq_len" in model_config:
      max_seq_len = model_config["max_seq_len"]
      logging.info("Truncating text to have at most %d tokens", max_seq_len)
      text = misc_utils.random_substr(text, max_seq_len)
    text = tf.reshape(text, [batch_size])
    num_classes = len(model_config["labels"])
    label = tf.reshape(label, [batch_size, num_classes])
    prxlayer = projection_layers.ProjectionLayer(model_config, mode)
    projection, seq_length = prxlayer(text)
    gbst_max_token_len = max_seq_len
    if "gbst_max_token_len" in model_config:
      gbst_max_token_len = model_config["gbst_max_token_len"]
    byte_int = tftext.ByteSplitter().split(text).to_tensor(
        default_value=0, shape=[batch_size, gbst_max_token_len])
    token_ids = tf.to_int32(byte_int)
    token_len = tf.strings.length(text)
    mask = tf.to_int32(
        tf.sequence_mask(token_len, maxlen=gbst_max_token_len))
    mask *= 3
    token_ids += mask
    return {
        "projection": projection,
        "seq_length": seq_length,
        "token_ids": token_ids,
        "token_len": token_len,
        "label": label
    }

  def _input_fn(params):
    """Method to be used for reading the data."""
    assert mode != tf_estimator.ModeKeys.PREDICT
    split = "train" if mode == tf_estimator.ModeKeys.TRAIN else "test"
    ds = tfds.load(runner_config["dataset"], split=split)
    ds = ds.batch(params["batch_size"], drop_remainder=drop_remainder)
    ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    ds = ds.shuffle(buffer_size=100)
    ds = ds.repeat(count=1 if mode == tf_estimator.ModeKeys.EVAL else None)
    ds = ds.map(
        functools.partial(_post_processor, batch_size=params["batch_size"]),
        num_parallel_calls=tf.data.experimental.AUTOTUNE,
        deterministic=False)
    return ds

  return _input_fn
//  [python]
def create_input_fn(runner_config, mode, drop_remainder):
  """Returns an input function to use in the instantiation of tf.estimator.*."""

  def _post_processor(features, batch_size):
    """Post process the data to a form expected by model_fn."""
    data_processor = getattr(sys.modules[__name__], runner_config["dataset"])
    text, label = data_processor(features, runner_config)
    model_config = runner_config["model_config"]
    if "max_seq_len" in model_config:
      max_seq_len = model_config["max_seq_len"]
      logging.info("Truncating text to have at most %d tokens", max_seq_len)
      text = misc_utils.random_substr(text, max_seq_len)
    text = tf.reshape(text, [batch_size])
    num_classes = len(model_config["labels"])
    label = tf.reshape(label, [batch_size, num_classes])
    prxlayer = projection_layers.ProjectionLayer(model_config, mode)
    projection, seq_length = prxlayer(text)
    gbst_max_token_len = max_seq_len
    if "gbst_max_token_len" in model_config:
      gbst_max_token_len = model_config["gbst_max_token_len"]
    byte_int = tftext.ByteSplitter().split(text).to_tensor(
        default_value=0, shape=[batch_size, gbst_max_token_len])
    token_ids = tf.cast(byte_int, tf.int32)
    token_len = tf.strings.length(text)
    mask = tf.cast(
        tf.sequence_mask(token_len, maxlen=gbst_max_token_len), tf.int32)
    mask *= 3
    token_ids += mask
    return {
        "projection": projection,
        "seq_length": seq_length,
        "token_ids": token_ids,
        "token_len": token_len,
        "label": label
    }

  def _input_fn(params):
    """Method to be used for reading the data."""
    assert mode != tf_estimator.ModeKeys.PREDICT
    split = "train" if mode == tf_estimator.ModeKeys.TRAIN else "test"
    ds = tfds.load(runner_config["dataset"], split=split)
    ds = ds.batch(params["batch_size"], drop_remainder=drop_remainder)
    ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    ds = ds.shuffle(buffer_size=100)
    ds = ds.repeat(count=1 if mode == tf_estimator.ModeKeys.EVAL else None)
    ds = ds.map(
        functools.partial(_post_processor, batch_size=params["batch_size"]),
        num_parallel_calls=tf.data.experimental.AUTOTUNE,
        deterministic=False)
    return ds

  return _input_fn
//  [python]
def tf_distance_cov_cor(input1, input2, debug=False):
    start = time.time()

    input1 = tf.debugging.check_numerics(input1, "input1 contains nan/inf")
    input2 = tf.debugging.check_numerics(input2, "input2 contains nan/inf")

    n = tf.cast(tf.shape(input1)[0], tf.float32)
    a = pairwise_dist(input1, input1)
    b = pairwise_dist(input2, input2)

    A = a - tf.reduce_mean(a,
                           axis=1) - tf.expand_dims(tf.reduce_mean(a,
                                                                   axis=0),
                                                    axis=1) + tf.reduce_mean(a)
    B = b - tf.reduce_mean(b,
                           axis=1) - tf.expand_dims(tf.reduce_mean(b,
                                                                   axis=0),
                                                    axis=1) + tf.reduce_mean(b)

    dCovXY = tf.sqrt(tf.reduce_sum(A * B) / (n ** 2))
    dVarXX = tf.sqrt(tf.reduce_sum(A * A) / (n ** 2))
    dVarYY = tf.sqrt(tf.reduce_sum(B * B) / (n ** 2))

    dCorXY = dCovXY / tf.sqrt(dVarXX * dVarYY)
    end = time.time()
    if debug:
        print(("tf distance cov: {} and cor: {}, dVarXX: {}, "
               "dVarYY:{} uses: {}").format(
            dCovXY, dCorXY,
            dVarXX, dVarYY,
            end - start))
    return dCovXY, dCorXY
//  [python]
def features_to_int32(features: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
  """Converts tf.int64 features to tf.int32, keep other features the same.

  tf.Example only supports tf.int64, but the TPU only supports tf.int32.

  Args:
    features: Input tensor dictionary.

  Returns:
    Features with tf.int64 converted to tf.int32.
  """
  converted_features = {}
  for name, tensor in features.items():
    if tensor.dtype == tf.int64:
      converted_features[name] = tf.to_int32(tensor)
    else:
      converted_features[name] = tensor
  return converted_features
//  [python]
def features_to_int32(features: Dict[str, tf.Tensor]) -> Dict[str, tf.Tensor]:
  """Converts tf.int64 features to tf.int32, keep other features the same.

  tf.Example only supports tf.int64, but the TPU only supports tf.int32.

  Args:
    features: Input tensor dictionary.

  Returns:
    Features with tf.int64 converted to tf.int32.
  """
  converted_features = {}
  for name, tensor in features.items():
    if tensor.dtype == tf.int64:
      converted_features[name] = tf.cast(tensor, tf.int32)
    else:
      converted_features[name] = tensor
  return converted_features
//  [python]
def siamese_cosine_loss(left, right, y, scope="cosine_loss"):
    r"""Loss for Siamese networks (cosine version).
    Same as :func:`contrastive_loss` but with different similarity measurement.

    .. math::
        [\frac{l \cdot r}{\lVert l\rVert \lVert r\rVert} - (2y-1)]^2

    Args:
        left (tf.Tensor): left feature vectors of shape [Batch, N].
        right (tf.Tensor): right feature vectors of shape [Batch, N].
        y (tf.Tensor): binary labels of shape [Batch]. 1: similar, 0: not similar.

    Returns:
        tf.Tensor: cosine-loss as a scalar tensor.
    """

    def l2_norm(t, eps=1e-12):
        """
        Returns:
            tf.Tensor: norm of 2D input tensor on axis 1
        """
        with tf.name_scope("l2_norm"):
            return tf.sqrt(tf.reduce_sum(tf.square(t), 1) + eps)

    with tf.name_scope(scope):
        y = 2 * tf.to_float(y) - 1
        pred = tf.reduce_sum(left * right, 1) / (l2_norm(left) * l2_norm(right) + 1e-10)

        return tf.nn.l2_loss(y - pred) / tf.to_float(tf.shape(left)[0])
//  [python]
def siamese_cosine_loss(left, right, y, scope="cosine_loss"):
    r"""Loss for Siamese networks (cosine version).
    Same as :func:`contrastive_loss` but with different similarity measurement.

    .. math::
        [\frac{l \cdot r}{\lVert l\rVert \lVert r\rVert} - (2y-1)]^2

    Args:
        left (tf.Tensor): left feature vectors of shape [Batch, N].
        right (tf.Tensor): right feature vectors of shape [Batch, N].
        y (tf.Tensor): binary labels of shape [Batch]. 1: similar, 0: not similar.

    Returns:
        tf.Tensor: cosine-loss as a scalar tensor.
    """

    def l2_norm(t, eps=1e-12):
        """
        Returns:
            tf.Tensor: norm of 2D input tensor on axis 1
        """
        with tf.name_scope("l2_norm"):
            return tf.sqrt(tf.reduce_sum(tf.square(t), 1) + eps)

    with tf.name_scope(scope):
        y = 2 * tf.cast(y, tf.float32) - 1
        pred = tf.reduce_sum(left * right, 1) / (l2_norm(left) * l2_norm(right) + 1e-10)

        return tf.nn.l2_loss(y - pred) / tf.cast(tf.shape(left)[0], tf.float32)
//  [python]
def contrastive_loss(left, right, y, margin, extra=False, scope="constrastive_loss"):
    r"""Loss for Siamese networks as described in the paper:
    `Learning a Similarity Metric Discriminatively, with Application to Face
    Verification <http://yann.lecun.com/exdb/publis/pdf/chopra-05.pdf>`_ by Chopra et al.

    .. math::
        \frac{1}{2} [y \cdot d^2 + (1-y) \cdot \max(0, m - d)^2], d = \Vert l - r \Vert_2

    Args:
        left (tf.Tensor): left feature vectors of shape [Batch, N].
        right (tf.Tensor): right feature vectors of shape [Batch, N].
        y (tf.Tensor): binary labels of shape [Batch]. 1: similar, 0: not similar.
        margin (float): horizon for negative examples (y==0).
        extra (bool): also return distances for pos and neg.

    Returns:
        tf.Tensor: constrastive_loss (averaged over the batch), (and optionally average_pos_dist, average_neg_dist)
    """
    with tf.name_scope(scope):
        y = tf.to_float(y)

        delta = tf.reduce_sum(tf.square(left - right), 1)
        delta_sqrt = tf.sqrt(delta + 1e-10)

        match_loss = delta
        missmatch_loss = tf.square(tf.nn.relu(margin - delta_sqrt))

        loss = tf.reduce_mean(0.5 * (y * match_loss + (1 - y) * missmatch_loss))

        if extra:
            num_pos = tf.count_nonzero(y)
            num_neg = tf.count_nonzero(1 - y)
            pos_dist = tf.where(tf.equal(num_pos, 0), 0.,
                                tf.reduce_sum(y * delta_sqrt) / tf.cast(num_pos, tf.float32),
                                name="pos-dist")
            neg_dist = tf.where(tf.equal(num_neg, 0), 0.,
                                tf.reduce_sum((1 - y) * delta_sqrt) / tf.cast(num_neg, tf.float32),
                                name="neg-dist")
            return loss, pos_dist, neg_dist
        else:
            return loss
//  [python]
def contrastive_loss(left, right, y, margin, extra=False, scope="constrastive_loss"):
    r"""Loss for Siamese networks as described in the paper:
    `Learning a Similarity Metric Discriminatively, with Application to Face
    Verification <http://yann.lecun.com/exdb/publis/pdf/chopra-05.pdf>`_ by Chopra et al.

    .. math::
        \frac{1}{2} [y \cdot d^2 + (1-y) \cdot \max(0, m - d)^2], d = \Vert l - r \Vert_2

    Args:
        left (tf.Tensor): left feature vectors of shape [Batch, N].
        right (tf.Tensor): right feature vectors of shape [Batch, N].
        y (tf.Tensor): binary labels of shape [Batch]. 1: similar, 0: not similar.
        margin (float): horizon for negative examples (y==0).
        extra (bool): also return distances for pos and neg.

    Returns:
        tf.Tensor: constrastive_loss (averaged over the batch), (and optionally average_pos_dist, average_neg_dist)
    """
    with tf.name_scope(scope):
        y = tf.cast(y, tf.float32)

        delta = tf.reduce_sum(tf.square(left - right), 1)
        delta_sqrt = tf.sqrt(delta + 1e-10)

        match_loss = delta
        missmatch_loss = tf.square(tf.nn.relu(margin - delta_sqrt))

        loss = tf.reduce_mean(0.5 * (y * match_loss + (1 - y) * missmatch_loss))

        if extra:
            num_pos = tf.count_nonzero(y)
            num_neg = tf.count_nonzero(1 - y)
            pos_dist = tf.where(tf.equal(num_pos, 0), 0.,
                                tf.reduce_sum(y * delta_sqrt) / tf.cast(num_pos, tf.float32),
                                name="pos-dist")
            neg_dist = tf.where(tf.equal(num_neg, 0), 0.,
                                tf.reduce_sum((1 - y) * delta_sqrt) / tf.cast(num_neg, tf.float32),
                                name="neg-dist")
            return loss, pos_dist, neg_dist
        else:
            return loss
//  [python]
def random_grayscale(p):
  """Randomly converts imageto gray."""
  def _to_grayscale(image):
    return tf.cond(
        tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                tf.cast(p, tf.float32)),
        lambda: tf.tile(tf.image.rgb_to_grayscale(image), [1, 1, 3]),
        lambda: image)
  return _to_grayscale
//  [python]
def random_blur(height=224, p=1.0):
  """Randomly blurs the image."""
  def _random_blur(image):
    sig = tf.random.uniform([], 0.1, 2.0, dtype=tf.float32)
    ks = height // 10
    return tf.cond(
        tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                t.cast(p, tf.float32)),
        lambda: gaussian_blur(image, kernel_size=ks, sigma=sig, padding="SAME"),
        lambda: image)
  return _random_blur
//  [python]
def random_blur(height=224, p=1.0):
  """Randomly blurs the image."""
  def _random_blur(image):
    sig = tf.random.uniform([], 0.1, 2.0, dtype=tf.float32)
    ks = height // 10
    return tf.cond(
        tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                tf.cast(p, tf.float32)),
        lambda: gaussian_blur(image, kernel_size=ks, sigma=sig, padding="SAME"),
        lambda: image)
  return _random_blur
//  [python]
def random_solarization(p=0.1):
  """Randomly solarizes the images."""
  def _solarize(image):
    image = image * t.cast(t.less(image, 0.5), t.float32) + (
        1.0 - image) * t.cast(t.greater_equal(image, 0.5), t.float32)
    return image
  def _random_solarize(image):
    return t.cond(
        t.less(t.random.uniform([], minval=0, maxval=1, dtype=t.float32),
                t.cast(p, t.float32)),
        lambda: _solarize(image),
        lambda: image)
  return _random_solarize
//  [python]
def random_solarization(p=0.1):
  """Randomly solarizes the images."""
  def _solarize(image):
    image = image * tf.cast(tf.less(image, 0.5), tf.float32) + (
        1.0 - image) * tf.cast(tf.greater_equal(image, 0.5), tf.float32)
    return image
  def _random_solarize(image):
    return tf.cond(
        tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),
                tf.cast(p, tf.float32)),
        lambda: _solarize(image),
        lambda: image)
  return _random_solarize
//  [python]
def sg_len(x): # 3D: batch x seqlen x dim
    non_zero = tf.not_equal(tf.reduce_sum(x, axis=2), 0.)
    return tf.reduce_sum(tf.to_int32(non_zero), axis=1)
//  [python]
def sg_len(x): # 3D: batch x seqlen x dim
    non_zero = tf.not_equal(tf.reduce_sum(x, axis=2), 0.)
    return tf.reduce_sum(tf.cast(non_zero, tf.int32), axis=1)
//  [python]
def atari_model(img_input, num_actions, scope, reuse=False):
    with tf.variable_scope(scope, reuse=reuse):
        out = tf.to_float(img_input) / 255.0
        with tf.variable_scope("convnet"):
            # original architecture
            out = layers.convolution2d(out, num_outputs=32, kernel_size=8, stride=4, activation_fn=tf.nn.relu)
            out = layers.convolution2d(out, num_outputs=64, kernel_size=4, stride=2, activation_fn=tf.nn.relu)
            out = layers.convolution2d(out, num_outputs=64, kernel_size=3, stride=1, activation_fn=tf.nn.relu)
        out = layers.flatten(out)
        with tf.variable_scope("action_value"):
            out = layers.fully_connected(out, num_outputs=512, activation_fn=tf.nn.relu)
            out = layers.fully_connected(out, num_outputs=num_actions, activation_fn=None)

        return out
//  [python]
def atari_model(img_input, num_actions, scope, reuse=False):
    with tf.variable_scope(scope, reuse=reuse):
        out = tf.cast(img_input, tf.float32) / 255.0
        with tf.variable_scope("convnet"):
            # original architecture
            out = layers.convolution2d(out, num_outputs=32, kernel_size=8, stride=4, activation_fn=tf.nn.relu)
            out = layers.convolution2d(out, num_outputs=64, kernel_size=4, stride=2, activation_fn=tf.nn.relu)
            out = layers.convolution2d(out, num_outputs=64, kernel_size=3, stride=1, activation_fn=tf.nn.relu)
        out = layers.flatten(out)
        with tf.variable_scope("action_value"):
            out = layers.fully_connected(out, num_outputs=512, activation_fn=tf.nn.relu)
            out = layers.fully_connected(out, num_outputs=num_actions, activation_fn=None)

        return out
//  [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, t.int32)
  kernel_size = radius * 2 + 1
  x = t.cast(t.range(-radius, radius + 1), t.float32)
  blur_filter = t.exp(
      -t.pow(x, 2.0) / (2.0 * t.pow(t.cast(sigma, 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, tf.int32)
  kernel_size = radius * 2 + 1
  x = tf.cast(tf.range(-radius, radius + 1), tf.float32)
  blur_filter = tf.exp(
      -tf.pow(x, 2.0) / (2.0 * tf.pow(tf.cast(sigma, 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 _expand_local_ids_to_blocks(local_ids: tf.Tensor,
                                mask_padding_ids: bool = True) -> tf.Tensor:
  """Helper to expand local ids to blocked format.

  Args:
    local_ids: [batch_size, seq_len, 2*local_radius + 1] shaped Tensor. This is
      the shape of the `att_mask` and `relative_att_id` Tensors in the
      `QkvRelativeLocalAttention` layer.
    mask_padding_ids: If True (the default), zero out ids representing attention
      from the first tokens to padding tokens before the start of the sequence
      and from the last tokens to padding tokens after the end of the sequence.
      When the ids are actually attention masks (such that 0 means no attention)
      this removes attention to any padding tokens beyond the boundaries of the
      sequence.

  Returns:
    A Tensor of shape [batch_size, num_blocks, block_len, 3*block_len],
    where block_len = local_radius + 1, and
    num_blocks = ceiling(seq_len / block_len).
  """
  batch_size = tf.shape(local_ids)[0]
  local_window_size = local_ids.shape.as_list()[-1]
  if local_window_size is None:
    raise ValueError('`local_ids.shape[-1]` must be known statically.')
  block_len = local_window_size // 2 + 1

  if mask_padding_ids:
    padding_mask = feature_utils.make_local_segmented_att_mask(
        tf.ones(tf.shape(local_ids)[:-1]), local_radius=block_len - 1)
    local_ids *= t.cast(padding_mask, local_ids.dtype)

  # [batch_size, num_blocks, block_len, 2*local_radius + 1]
  blocked_ids = tensor_utils.split_into_blocks(
      local_ids, block_len=block_len, axis=1)

  # [batch_size, num_blocks, block_len, 3*block_len - 2]
  skewed_ids = tensor_utils.skew_elements_right(blocked_ids, axis=-1)

  # [batch_size, num_blocks, block_len, 3*block_len]
  result = tf.pad(skewed_ids, [[0, 0], [0, 0], [0, 0], [1, 1]])

  if not mask_padding_ids:
    return result

  ones = tf.ones_like(result, dtype=result.dtype)

  # [batch_size, 1, block_len, 3*block_len]
  leftmost_mask = tf.concat([
      tf.zeros([batch_size, 1, block_len, block_len], dtype=ones.dtype),
      tf.ones([batch_size, 1, block_len, 2 * block_len], dtype=ones.dtype)
  ], -1)

  # [batch_size, 1, block_len, 3*block_len]
  rightmost_mask = tf.concat([
      tf.ones([batch_size, 1, block_len, 2 * block_len], dtype=ones.dtype),
      tf.zeros([batch_size, 1, block_len, block_len], dtype=ones.dtype)
  ], -1)

  result *= tf.concat([leftmost_mask, ones[:, 1:, :, :]], 1)
  result *= tf.concat([ones[:, :-1, :, :], rightmost_mask], 1)
  return result
//  [python]
def _expand_local_ids_to_blocks(local_ids: tf.Tensor,
                                mask_padding_ids: bool = True) -> tf.Tensor:
  """Helper to expand local ids to blocked format.

  Args:
    local_ids: [batch_size, seq_len, 2*local_radius + 1] shaped Tensor. This is
      the shape of the `att_mask` and `relative_att_id` Tensors in the
      `QkvRelativeLocalAttention` layer.
    mask_padding_ids: If True (the default), zero out ids representing attention
      from the first tokens to padding tokens before the start of the sequence
      and from the last tokens to padding tokens after the end of the sequence.
      When the ids are actually attention masks (such that 0 means no attention)
      this removes attention to any padding tokens beyond the boundaries of the
      sequence.

  Returns:
    A Tensor of shape [batch_size, num_blocks, block_len, 3*block_len],
    where block_len = local_radius + 1, and
    num_blocks = ceiling(seq_len / block_len).
  """
  batch_size = tf.shape(local_ids)[0]
  local_window_size = local_ids.shape.as_list()[-1]
  if local_window_size is None:
    raise ValueError('`local_ids.shape[-1]` must be known statically.')
  block_len = local_window_size // 2 + 1

  if mask_padding_ids:
    padding_mask = feature_utils.make_local_segmented_att_mask(
        tf.ones(tf.shape(local_ids)[:-1]), local_radius=block_len - 1)
    local_ids *= tf.cast(padding_mask, local_ids.dtype)

  # [batch_size, num_blocks, block_len, 2*local_radius + 1]
  blocked_ids = tensor_utils.split_into_blocks(
      local_ids, block_len=block_len, axis=1)

  # [batch_size, num_blocks, block_len, 3*block_len - 2]
  skewed_ids = tensor_utils.skew_elements_right(blocked_ids, axis=-1)

  # [batch_size, num_blocks, block_len, 3*block_len]
  result = tf.pad(skewed_ids, [[0, 0], [0, 0], [0, 0], [1, 1]])

  if not mask_padding_ids:
    return result

  ones = tf.ones_like(result, dtype=result.dtype)

  # [batch_size, 1, block_len, 3*block_len]
  leftmost_mask = tf.concat([
      tf.zeros([batch_size, 1, block_len, block_len], dtype=ones.dtype),
      tf.ones([batch_size, 1, block_len, 2 * block_len], dtype=ones.dtype)
  ], -1)

  # [batch_size, 1, block_len, 3*block_len]
  rightmost_mask = tf.concat([
      tf.ones([batch_size, 1, block_len, 2 * block_len], dtype=ones.dtype),
      tf.zeros([batch_size, 1, block_len, block_len], dtype=ones.dtype)
  ], -1)

  result *= tf.concat([leftmost_mask, ones[:, 1:, :, :]], 1)
  result *= tf.concat([ones[:, :-1, :, :], rightmost_mask], 1)
  return result
//  [python]
def loss(self, prediction_dict):

        # these should include mini-batch values only
        objectness_gt = prediction_dict[self.PRED_MB_OBJECTNESS_GT]
        offsets_gt = prediction_dict[self.PRED_MB_OFFSETS_GT]

        # Predictions
        with tf.variable_scope('rpn_prediction_mini_batch'):
            objectness = prediction_dict[self.PRED_MB_OBJECTNESS]
            offsets = prediction_dict[self.PRED_MB_OFFSETS]

        with tf.variable_scope('rpn_losses'):
            with tf.variable_scope('objectness'):
                cls_loss = losses.WeightedSoftmaxLoss()
                cls_loss_weight = self._config.loss_config.cls_loss_weight
                objectness_loss = cls_loss(objectness,
                                           objectness_gt,
                                           weight=cls_loss_weight)

                with tf.variable_scope('obj_norm'):
                    # normalize by the number of anchor mini-batches
                    objectness_loss = objectness_loss / tf.to_float(
                        tf.shape(objectness_gt)[0])
                    tf.summary.scalar('objectness', objectness_loss)

            with tf.variable_scope('regression'):
                reg_loss = losses.WeightedSmoothL1Loss()
                reg_loss_weight = self._config.loss_config.reg_loss_weight
                anchorwise_localization_loss = reg_loss(offsets,
                                                        offsets_gt,
                                                        weight=reg_loss_weight)
                masked_localization_loss = \
                    anchorwise_localization_loss * objectness_gt[:, 1]
                localization_loss = tf.reduce_sum(masked_localization_loss)

                with tf.variable_scope('reg_norm'):
                    # normalize by the number of positive objects
                    num_positives = tf.reduce_sum(objectness_gt[:, 1])
                    # Assert the condition `num_positives > 0`
                    with tf.control_dependencies(
                            [tf.assert_positive(num_positives)]):
                        localization_loss = localization_loss / num_positives
                        tf.summary.scalar('regression', localization_loss)

            with tf.variable_scope('total_loss'):
                total_loss = objectness_loss + localization_loss

        loss_dict = {
            self.LOSS_RPN_OBJECTNESS: objectness_loss,
            self.LOSS_RPN_REGRESSION: localization_loss,
        }

        return loss_dict, total_loss