jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _encode_tf(self, s: tf.Tensor) -> tf.Tensor:
    """Encode a string tf.Tensor.

    Args:
      s: a tf.Scalar or a 1d tf.Tensor in [batch_size] with dtype tf.string

    Returns:
      a 1d tf.Tensor or
      a tf.RaggedTensor in [batch_size, None] with dtype tf.int32
    """

    def _encode_py_func(text) -> Sequence[int]:
      text = tf.compat.as_text(text.numpy().decode("UTF-8"))
      result = self.encode(text)
      result = tf.expand_dims(result, 0)
      return tf.cast(result, tf.int32)

    def _encode_func():
      result = tf.py_function(func=_encode_py_func, inp=[s], Tout=tf.int32)
      result = tf.squeeze(result, 0)
      result.set_shape([None])
      return result

    def _batch_encode_py_func(strs) -> Sequence[Sequence[int]]:
      texts = []
      for s in strs.numpy():
        texts.append(tf.compat.as_text(s.decode("UTF-8")))
      results = self.model.encode_batch(texts)
      return tf.ragged.constant(results, dtype=tf.int32)

    def _batch_encode_func():
      return tf.py_function(
          func=_batch_encode_py_func,
          inp=[s],
          Tout=tf.RaggedTensorSpec(shape=[s.shape[0], None], dtype=tf.int32),
      )

    # pylint: disable=g-explicit-length-test
    if len(s.shape) > 0:
      return _batch_encode_func()
    return _encode_func()
//  [python]
def filter_greater_than(boxlist, thresh, scope=None):
  """Filter to keep only boxes with score exceeding a given threshold.

  This op keeps the collection of boxes whose corresponding scores are
  greater than the input threshold.

  TODO(jonathanhuang): Change function name to filter_scores_greater_than

  Args:
    boxlist: BoxList holding N boxes.  Must contain a 'scores' field
      representing detection scores.
    thresh: scalar threshold
    scope: name scope.

  Returns:
    a BoxList holding M boxes where M <= N

  Raises:
    ValueError: if boxlist not a BoxList object or if it does not
      have a scores field
  """
  with tf.name_scope(scope or 'FilterGreaterThan'):
    if not isinstance(boxlist, box_list.BoxList):
      raise ValueError('boxlist must be a BoxList')
    if not boxlist.has_field('scores'):
      raise ValueError('input boxlist must have \'scores\' field')
    scores = boxlist.get_field('scores')
    if len(scores.shape.as_list()) > 2:
      raise ValueError('Scores should have rank 1 or 2')
    if len(scores.shape.as_list()) == 2 and scores.shape.as_list()[1] != 1:
      raise ValueError('Scores should have rank 1 or have shape '
                       'consistent with [None, 1]')
    high_score_indices = tf.reshape(tf.where(tf.greater(scores, thresh)), [-1])
    return gather(boxlist, high_score_indices)
//  [python]
def filter_greater_than(boxlist, thresh, scope=None):
  """Filter to keep only boxes with score exceeding a given threshold.

  This op keeps the collection of boxes whose corresponding scores are
  greater than the input threshold.

  TODO(jonathanhuang): Change function name to filter_scores_greater_than

  Args:
    boxlist: BoxList holding N boxes.  Must contain a 'scores' field
      representing detection scores.
    thresh: scalar threshold
    scope: name scope.

  Returns:
    a BoxList holding M boxes where M <= N

  Raises:
    ValueError: if boxlist not a BoxList object or if it does not
      have a scores field
  """
  with tf.name_scope(scope or 'FilterGreaterThan'):
    if not isinstance(boxlist, box_list.BoxList):
      raise ValueError('boxlist must be a BoxList')
    if not boxlist.has_field('scores'):
      raise ValueError('input boxlist must have \'scores\' field')
    scores = boxlist.get_field('scores')
    if len(scores.shape.as_list()) > 2:
      raise ValueError('Scores should have rank 1 or 2')
    if len(scores.shape.as_list()) == 2 and scores.shape.as_list()[1] != 1:
      raise ValueError('Scores should have rank 1 or have shape '
                       'consistent with [None, 1]')
    high_score_indices = tf.cast(
        tf.reshape(tf.where(tf.greater(scores, thresh)), [-1]), tf.int32)
    return gather(boxlist, high_score_indices)
//  [python]
def to_normalized_coordinates(boxlist,
                              height,
                              width,
                              check_range=True,
                              scope=None):
  """Converts absolute box coordinates to normalized coordinates in [0, 1].

  Usually one uses the dynamic shape of the image or conv-layer tensor:
    boxlist = box_list_ops.to_normalized_coordinates(boxlist,
                                                     tf.shape(images)[1],
                                                     tf.shape(images)[2]),

  This function raises an assertion failed error at graph execution time when
  the maximum coordinate is smaller than 1.01 (which means that coordinates are
  already normalized). The value 1.01 is to deal with small rounding errors.

  Args:
    boxlist: BoxList with coordinates in terms of pixel-locations.
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    scope: name scope.

  Returns:
    boxlist with normalized coordinates in [0, 1].
  """
  with tf.name_scope(scope or 'ToNormalizedCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    if check_range:
      max_val = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(
          tf.greater(max_val, 1.01),
          ['max value is lower than 1.01: ', max_val])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, 1 / height, 1 / width)
//  [python]
def to_absolute_coordinates(boxlist,
                            height,
                            width,
                            check_range=True,
                            maximum_normalized_coordinate=1.1,
                            scope=None):
  """Converts normalized box coordinates to absolute pixel coordinates.

  This function raises an assertion failed error when the maximum box coordinate
  value is larger than maximum_normalized_coordinate (in which case coordinates
  are already absolute).

  Args:
    boxlist: BoxList with coordinates in range [0, 1].
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    maximum_normalized_coordinate: Maximum coordinate value to be considered as
      normalized, default to 1.1.
    scope: name scope.

  Returns:
    boxlist with absolute coordinates in terms of the image size.

  """
  with tf.name_scope(scope or 'ToAbsoluteCoordinates'):
    height = tf.to_float(height)
    width = tf.to_float(width)

    # Ensure range of input boxes is correct.
    if check_range:
      box_maximum = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(
          tf.greater_equal(maximum_normalized_coordinate, box_maximum), [
              'maximum box coordinate value is larger '
              'than %f: ' % maximum_normalized_coordinate, box_maximum
          ])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, height, width)
//  [python]
def to_absolute_coordinates(boxlist,
                            height,
                            width,
                            check_range=True,
                            maximum_normalized_coordinate=1.1,
                            scope=None):
  """Converts normalized box coordinates to absolute pixel coordinates.

  This function raises an assertion failed error when the maximum box coordinate
  value is larger than maximum_normalized_coordinate (in which case coordinates
  are already absolute).

  Args:
    boxlist: BoxList with coordinates in range [0, 1].
    height: Maximum value for height of absolute box coordinates.
    width: Maximum value for width of absolute box coordinates.
    check_range: If True, checks if the coordinates are normalized or not.
    maximum_normalized_coordinate: Maximum coordinate value to be considered as
      normalized, default to 1.1.
    scope: name scope.

  Returns:
    boxlist with absolute coordinates in terms of the image size.

  """
  with tf.name_scope(scope or 'ToAbsoluteCoordinates'):
    height = tf.cast(height, tf.float32)
    width = tf.cast(width, tf.float32)

    # Ensure range of input boxes is correct.
    if check_range:
      box_maximum = tf.reduce_max(boxlist.get())
      max_assert = tf.Assert(
          tf.greater_equal(maximum_normalized_coordinate, box_maximum), [
              'maximum box coordinate value is larger '
              'than %f: ' % maximum_normalized_coordinate, box_maximum
          ])
      with tf.control_dependencies([max_assert]):
        width = tf.identity(width)

    return scale(boxlist, height, width)
//  [python]
def boolean_mask(boxlist,
                 indicator,
                 fields=None,
                 scope=None,
                 use_static_shapes=False,
                 indicator_sum=None):
  """Select boxes from BoxList according to indicator and return new BoxList.

  `boolean_mask` returns the subset of boxes that are marked as "True" by the
  indicator tensor. By default, `boolean_mask` returns boxes corresponding to
  the input index list, as well as all additional fields stored in the boxlist
  (indexing into the first dimension).  However one can optionally only draw
  from a subset of fields.

  Args:
    boxlist: BoxList holding N boxes
    indicator: a rank-1 boolean tensor
    fields: (optional) list of fields to also gather from.  If None (default),
      all fields are gathered from.  Pass an empty fields list to only gather
      the box coordinates.
    scope: name scope.
    use_static_shapes: Whether to use an implementation with static shape
      gurantees.
    indicator_sum: An integer containing the sum of `indicator` vector. Only
      required if `use_static_shape` is True.

  Returns:
    subboxlist: a BoxList corresponding to the subset of the input BoxList
      specified by indicator
  Raises:
    ValueError: if `indicator` is not a rank-1 boolean tensor.
  """
  with tf.name_scope(scope or 'BooleanMask'):
    if indicator.shape.ndims != 1:
      raise ValueError('indicator should have rank 1')
    if indicator.dtype != tf.bool:
      raise ValueError('indicator should be a boolean tensor')
    if use_static_shapes:
      if not (indicator_sum and isinstance(indicator_sum, int)):
        raise ValueError('`indicator_sum` must be a of type int')
      selected_positions = tf.to_float(indicator)
      indexed_positions = tf.to_int32(
          tf.multiply(tf.cumsum(selected_positions), selected_positions))
      one_hot_selector = tf.one_hot(
          indexed_positions - 1, indicator_sum, dtype=tf.float32)
      sampled_indices = tf.to_int32(
          tf.tensordot(
              tf.to_float(tf.range(tf.shape(indicator)[0])),
              one_hot_selector,
              axes=[0, 0]))
      return gather(boxlist, sampled_indices, use_static_shapes=True)
    else:
      subboxlist = box_list.BoxList(tf.boolean_mask(boxlist.get(), indicator))
      if fields is None:
        fields = boxlist.get_extra_fields()
      for field in fields:
        if not boxlist.has_field(field):
          raise ValueError('boxlist must contain all specified fields')
        subfieldlist = tf.boolean_mask(boxlist.get_field(field), indicator)
        subboxlist.add_field(field, subfieldlist)
      return subboxlist
//  [python]
def boolean_mask(boxlist,
                 indicator,
                 fields=None,
                 scope=None,
                 use_static_shapes=False,
                 indicator_sum=None):
  """Select boxes from BoxList according to indicator and return new BoxList.

  `boolean_mask` returns the subset of boxes that are marked as "True" by the
  indicator tensor. By default, `boolean_mask` returns boxes corresponding to
  the input index list, as well as all additional fields stored in the boxlist
  (indexing into the first dimension).  However one can optionally only draw
  from a subset of fields.

  Args:
    boxlist: BoxList holding N boxes
    indicator: a rank-1 boolean tensor
    fields: (optional) list of fields to also gather from.  If None (default),
      all fields are gathered from.  Pass an empty fields list to only gather
      the box coordinates.
    scope: name scope.
    use_static_shapes: Whether to use an implementation with static shape
      gurantees.
    indicator_sum: An integer containing the sum of `indicator` vector. Only
      required if `use_static_shape` is True.

  Returns:
    subboxlist: a BoxList corresponding to the subset of the input BoxList
      specified by indicator
  Raises:
    ValueError: if `indicator` is not a rank-1 boolean tensor.
  """
  with tf.name_scope(scope or 'BooleanMask'):
    if indicator.shape.ndims != 1:
      raise ValueError('indicator should have rank 1')
    if indicator.dtype != tf.bool:
      raise ValueError('indicator should be a boolean tensor')
    if use_static_shapes:
      if not (indicator_sum and isinstance(indicator_sum, int)):
        raise ValueError('`indicator_sum` must be a of type int')
      selected_positions = tf.cast(indicator, dtype=tf.float32)
      indexed_positions = tf.cast(
          tf.multiply(tf.cumsum(selected_positions), selected_positions),
          dtype=tf.int32)
      one_hot_selector = tf.one_hot(
          indexed_positions - 1, indicator_sum, dtype=tf.float32)
      sampled_indices = tf.cast(
          tf.tensordot(
              tf.cast(tf.range(tf.shape(indicator)[0]), dtype=tf.float32),
              one_hot_selector,
              axes=[0, 0]),
          dtype=tf.int32)
      return gather(boxlist, sampled_indices, use_static_shapes=True)
    else:
      subboxlist = box_list.BoxList(tf.boolean_mask(boxlist.get(), indicator))
      if fields is None:
        fields = boxlist.get_extra_fields()
      for field in fields:
        if not boxlist.has_field(field):
          raise ValueError('boxlist must contain all specified fields')
        subfieldlist = tf.boolean_mask(boxlist.get_field(field), indicator)
        subboxlist.add_field(field, subfieldlist)
      return subboxlist
//  [python]
def visualize_tensors(name, imgs, scale_func=lambda x: (x + 1.) * 128., max_outputs=1):
    """Generate tensor for TensorBoard (casting, clipping)

    Args:
        name: name for visualization operation
        *imgs: multiple tensors as list
        scale_func: scale input tensors to fit range [0, 255]

    Example:
        visualize_tensors('viz1', [img1])
        visualize_tensors('viz2', [img1, img2, img3], max_outputs=max(30, BATCH))
    """
    xy = scale_func(tf.concat(imgs, axis=2))
    xy = tf.cast(tf.clip_by_value(xy, 0, 255), tf.uint8, name='viz')
    tf.summary.image(name, xy, max_outputs=30)
//  [python]
def _remove_low_score_boxes(outputs,
                            score_thresh,
                            min_num_boxes=100,
                            epsilon=0.0001):
  """remove the boxes that have an score lower than threshold.

  Args:
    outputs: Output dictionary.
    score_thresh: A float corresponding to score threshold.
    min_num_boxes: Minimum number of boxes.
    epsilon: A very small value.
  """
  max_scores = tf.reduce_max(
      outputs[standard_fields.DetectionResultFields.objects_score], axis=1)
  max_score = tf.reduce_max(max_scores)
  score_mask = tf.greater(max_scores,
                          tf.minimum(score_thresh, max_score - epsilon))
  high_score_indices = tf.cast(
      tf.reshape(tf.where(score_mask), [-1]), dtype=tf.int32)
  top_k_indices = _get_top_k_indices(outputs=outputs, k=min_num_boxes)
  high_score_indices = tf.cond(
      tf.shape(high_score_indices)[0] > min_num_boxes,
      lambda: high_score_indices, lambda: top_k_indices)

  for tensor_name in standard_fields.get_output_object_fields():
    if tensor_name in outputs and outputs[tensor_name] is not None:
      outputs[tensor_name] = tf.gather(outputs[tensor_name], high_score_indices)
//  [python]
def _multi_category_focal_loss(cls, onehot_labels, logits, classes_weights, gamma=2.0):
        """

        :param onehot_labels:
        :param logits:
        :param classes_weights:
        :param gamma:
        :return:
        """
        epsilon = 1.e-7
        alpha = tf.mul(onehot_labels, classes_weights)
        alpha = tf.cast(alpha, tf.float32)
        gamma = float(gamma)
        y_true = tf.cast(onehot_labels, tf.float32)
        y_pred = tf.nn.softmax(logits, dim=-1)
        y_pred = tf.clip_by_value(y_pred, epsilon, 1. - epsilon)
        y_t = tf.mul(y_true, y_pred) + tf.mul(1-y_true, 1-y_pred)
        ce = -tf.log(y_t)
        weight = tf.pow(tf.sub(1., y_t), gamma)
        fl = tf.mul(tf.mul(weight, ce), alpha)
        loss = tf.reduce_mean(fl)
        
        return loss
//  [python]
def _multi_category_focal_loss(cls, onehot_labels, logits, classes_weights, gamma=2.0):
        """

        :param onehot_labels:
        :param logits:
        :param classes_weights:
        :param gamma:
        :return:
        """
        epsilon = 1.e-7
        alpha = tf.multiply(onehot_labels, classes_weights)
        alpha = tf.cast(alpha, tf.float32)
        gamma = float(gamma)
        y_true = tf.cast(onehot_labels, tf.float32)
        y_pred = tf.nn.softmax(logits, dim=-1)
        y_pred = tf.clip_by_value(y_pred, epsilon, 1. - epsilon)
        y_t = tf.multiply(y_true, y_pred) + tf.multiply(1-y_true, 1-y_pred)
        ce = -tf.log(y_t)
        weight = tf.pow(tf.subtract(1., y_t), gamma)
        fl = tf.multiply(tf.multiply(weight, ce), alpha)
        loss = tf.reduce_mean(fl)
        
        return loss
//  [python]
def cov_loss(self, matrix):
        n_rows = tf.to_float(tf.shape(matrix)[0])
        matrix = matrix - tf.reduce_mean(matrix, axis=0)
        cov = tf.matmul(matrix, matrix, transpose_a=True) / n_rows
        cov = tf.matrix_set_diag(cov, tf.zeros(self.config.embed_size, dtype=tf.float32))

        f2_norm = tf.sqrt(tf.reduce_sum(tf.square(cov)))
        return f2_norm
//  [python]
def cov_loss(self, matrix):
        n_rows = tf.cast(tf.shape(matrix)[0], tf.float32)
        matrix = matrix - tf.reduce_mean(matrix, axis=0)
        cov = tf.matmul(matrix, matrix, transpose_a=True) / n_rows
        cov = tf.matrix_set_diag(cov, tf.zeros(self.config.embed_size, dtype=tf.float32))

        f2_norm = tf.sqrt(tf.reduce_sum(tf.square(cov)))
        return f2_norm
//  [python]
def get_box_in_a_single_layer(self,datas,threshold):
        bboxes,clses,scores,indexs = self.box2box_transform.apply_deltas(datas)
        ids = datas['id_embedding']
        B,H,W,C = btf.combined_static_and_dynamic_shape(ids)
        ids = tf.reshape(ids,[B,H*W,C])
        ids = wmlt.batch_gather(ids,indexs)
        mask = tf.to_int32(tf.greater_equal(scores,threshold))
        length = tf.reduce_sum(mask,axis=-1)
        return bboxes,clses,scores,ids,length
//  [python]
def get_box_in_a_single_layer(self,datas,threshold):
        bboxes,clses,scores,indexs = self.box2box_transform.apply_deltas(datas)
        ids = datas['id_embedding']
        B,H,W,C = btf.combined_static_and_dynamic_shape(ids)
        ids = tf.reshape(ids,[B,H*W,C])
        ids = wmlt.batch_gather(ids,indexs)
        mask = tf.cast(tf.greater_equal(scores,threshold),tf.int32)
        length = tf.reduce_sum(mask,axis=-1)
        return bboxes,clses,scores,ids,length
//  [python]
def knn_graph_from_points_unbatched(points, k, distance_upper_bound, mask=None):
  """Returns the distances and indices of the neighbors of each point.

  Note that each point will have at least k neighbors unless the number of
  points is less than k. In that case, the python function that is wrapped in
  py_function will raise a value error.

  Args:
    points: A tf.float32 tensor of size [N, D] where D is the point dimensions.
    k: Number of neighbors for each point.
    distance_upper_bound: Only build the graph using points that are closer than
      this distance.
    mask: If not None, A tf.bool tensor of size [N]. If None, it is ignored.
      If not None, knn will be applied to only points where the mask is True.
      The points where the mask is False will have themselves as their
      neighbors.

  Returns:
    distances: A tf.float32 tensor of size [N, k].
    indices: A tf.int32 tensor of [N, k].
  """

  def fn(np_points, np_mask):
    return np_knn_graph_from_points_unbatched(
        points=np_points,
        k=k,
        distance_upper_bound=distance_upper_bound,
        mask=np_mask)

  num_points = tf.shape(points)[0]
  if mask is None:
    mask = tf.cast(tf.ones([num_points], dtype=tf.int32), dtype=tf.bool)
  else:
    mask = tf.reshape(mask, [num_points])
  distances, indices = tf.compat.v1.py_func(fn, [points, mask],
                                            [tf.float32, tf.int32])
  distances = tf.reshape(distances, [num_points, k])
  indices = tf.reshape(indices, [num_points, k])
  return distances, indices
//  [python]
def __call__(self, global_step):
    warmup_lr = self._params.warmup_learning_rate
    warmup_steps = self._params.warmup_steps
    init_lr = self._params.init_learning_rate
    lr_levels = self._params.learning_rate_levels
    lr_steps = self._params.learning_rate_steps
    linear_warmup = (
        warmup_lr + t.cast(global_step, dtype=t.float32) / warmup_steps *
        (init_lr - warmup_lr))
    learning_rate = tf.where(global_step < warmup_steps, linear_warmup, init_lr)

    for next_learning_rate, start_step in zip(lr_levels, lr_steps):
      learning_rate = tf.where(global_step >= start_step, next_learning_rate,
                               learning_rate)
    return learning_rate
//  [python]
def __call__(self, global_step):
    warmup_lr = self._params.warmup_learning_rate
    warmup_steps = self._params.warmup_steps
    init_lr = self._params.init_learning_rate
    lr_levels = self._params.learning_rate_levels
    lr_steps = self._params.learning_rate_steps
    linear_warmup = (
        warmup_lr + tf.cast(global_step, dtype=tf.float32) / warmup_steps *
        (init_lr - warmup_lr))
    learning_rate = tf.where(global_step < warmup_steps, linear_warmup, init_lr)

    for next_learning_rate, start_step in zip(lr_levels, lr_steps):
      learning_rate = tf.where(global_step >= start_step, next_learning_rate,
                               learning_rate)
    return learning_rate
//  [python]
def __call__(self, global_step):
    global_step = t.cast(global_step, dtype=t.float32)
    warmup_lr = self._params.warmup_learning_rate
    warmup_steps = self._params.warmup_steps
    init_lr = self._params.init_learning_rate
    total_steps = self._total_steps
    linear_warmup = (
        warmup_lr + global_step / warmup_steps * (init_lr - warmup_lr))
    cosine_learning_rate = (
        init_lr * (t.cos(np.pi * (global_step - warmup_steps) /
                          (total_steps - warmup_steps)) + 1.0) / 2.0)
    learning_rate = t.where(global_step < warmup_steps, linear_warmup,
                             cosine_learning_rate)
    return learning_rate
//  [python]
def __call__(self, global_step):
    global_step = tf.cast(global_step, dtype=tf.float32)
    warmup_lr = self._params.warmup_learning_rate
    warmup_steps = self._params.warmup_steps
    init_lr = self._params.init_learning_rate
    total_steps = self._total_steps
    linear_warmup = (
        warmup_lr + global_step / warmup_steps * (init_lr - warmup_lr))
    cosine_learning_rate = (
        init_lr * (tf.cos(np.pi * (global_step - warmup_steps) /
                          (total_steps - warmup_steps)) + 1.0) / 2.0)
    learning_rate = tf.where(global_step < warmup_steps, linear_warmup,
                             cosine_learning_rate)
    return learning_rate
//  [python]
def mask_mse_duration(self, y_true, y_pred):
        mask = tf.to_int32(tf.not_equal(y_true, 0))
        y_true = tf.to_float(y_true)
        y_pred = tf.to_float(y_pred)
        loss = tf.square(y_pred - y_true)

        loss = mask * loss
        loss = tf.reduce_sum(loss) / tf.reduce_sum(mask)

        return loss
//  [python]
def mask_mse_duration(self, y_true, y_pred):
        mask = tf.cast(tf.not_equal(y_true, 0), 1)
        y_true = tf.cast(y_true, tf.float32)
        y_pred = tf.cast(y_pred, tf.float32)
        loss = tf.square(y_pred - y_true)

        loss = mask * loss
        loss = tf.reduce_sum(loss) / tf.reduce_sum(mask)

        return loss
//  [python]
def mask_mse_mel(self, y_true, y_pred):
        T1=tf.shape(y_true)[1]
        T2=tf.shape(y_pred)[1]
        T=tf.reduce_min([T1,T2])
        y_true=y_true[:,:T,:]
        y_pred=y_pred[:,:T,:]
        y_true = tf.to_float(y_true)
        y_pred = tf.to_float(y_pred)
        mask = tf.cast(tf.not_equal(y_true, -(self.config['max_abs_value']+0.1)), 1)
        loss = tf.square(y_pred - y_true)
        loss = mask * loss
        loss = tf.reduce_sum(loss) / tf.reduce_sum(mask)
        return loss
//  [python]
def mask_mse_mel(self, y_true, y_pred):
        T1=tf.shape(y_true)[1]
        T2=tf.shape(y_pred)[1]
        T=tf.reduce_min([T1,T2])
        y_true=y_true[:,:T,:]
        y_pred=y_pred[:,:T,:]
        y_true = tf.cast(y_true, tf.float32)
        y_pred = tf.cast(y_pred, tf.float32)
        mask = tf.cast(tf.not_equal(y_true, -(self.config['max_abs_value']+0.1)), 1)
        loss = tf.square(y_pred - y_true)
        loss = mask * loss
        loss = tf.reduce_sum(loss) / tf.reduce_sum(mask)
        return loss
//  [python]
def local_model_fn(model, features, labels, mode):
    logits, local_logits, optimizer = model_def(features)

    if mode == tf.estimator.ModeKeys.PREDICT:
        return model.make_spec(mode=mode, predictions=logits)

    y = labels['y']
    if mode == tf.estimator.ModeKeys.EVAL:
        classes = tf.argmax(logits, axis=1)
        acc_pair = tf.metrics.accuracy(y, classes)
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y,
                                                              logits=logits)
        loss = tf.math.reduce_mean(loss)
        return model.make_spec(
            mode=mode, loss=loss, eval_metric_ops={'accuracy': acc_pair})

    global_step = tf.train.get_or_create_global_step()
    # mode == tf.estimator.ModeKeys.TRAIN
    logits_grad = model.send('logits', logits, require_grad=True)
    remote_train_op = model.minimize(
        optimizer, logits,
        grad_loss=logits_grad,
        global_step=global_step)

    # mode == tf.estimator.ModeKeys.TRAIN
    local_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=y, logits=local_logits)
    local_loss = tf.math.reduce_mean(local_loss)

    local_train_op = optimizer.minimize(
        local_loss, global_step=global_step)

    correct = tf.nn.in_top_k(predictions=local_logits, targets=y, k=1)
    acc = tf.reduce_mean(input_tensor=tf.to_float(correct))
    logging_hook = tf.estimator.LoggingTensorHook(
        {"local loss": local_loss, "local acc": acc}, every_n_iter=10)
    return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=local_loss,
            train_op=local_train_op,
            training_hooks=[logging_hook])
//  [python]
def local_model_fn(model, features, labels, mode):
    logits, local_logits, optimizer = model_def(features)

    if mode == tf.estimator.ModeKeys.PREDICT:
        return model.make_spec(mode=mode, predictions=logits)

    y = labels['y']
    if mode == tf.estimator.ModeKeys.EVAL:
        classes = tf.argmax(logits, axis=1)
        acc_pair = tf.metrics.accuracy(y, classes)
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y,
                                                              logits=logits)
        loss = tf.math.reduce_mean(loss)
        return model.make_spec(
            mode=mode, loss=loss, eval_metric_ops={'accuracy': acc_pair})

    global_step = tf.train.get_or_create_global_step()
    # mode == tf.estimator.ModeKeys.TRAIN
    logits_grad = model.send('logits', logits, require_grad=True)
    remote_train_op = model.minimize(
        optimizer, logits,
        grad_loss=logits_grad,
        global_step=global_step)

    # mode == tf.estimator.ModeKeys.TRAIN
    local_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=y, logits=local_logits)
    local_loss = tf.math.reduce_mean(local_loss)

    local_train_op = optimizer.minimize(
        local_loss, global_step=global_step)

    correct = tf.nn.in_top_k(predictions=local_logits, targets=y, k=1)
    acc = tf.reduce_mean(input_tensor=tf.cast(correct, tf.float32))
    logging_hook = tf.estimator.LoggingTensorHook(
        {"local loss": local_loss, "local acc": acc}, every_n_iter=10)
    return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=local_loss,
            train_op=local_train_op,
            training_hooks=[logging_hook])
//  [python]
def zero_out_clipped_grads(grad, x, clip_min, clip_max):
    """
    Helper function to erase entries in the gradient where the update would be
    clipped.
    :param grad: The gradient
    :param x: The current input
    :param clip_min: Minimum input component value
    :param clip_max: Maximum input component value
    """
    signed_grad = tf.sign(grad)

    # Find input components that lie at the boundary of the input range, and
    # where the gradient points in the wrong direction.
    clip_low = tf.logical_and(
        tf.less_equal(x, t.cast(clip_min, x.dtype)), tf.less(signed_grad, 0)
    )
    clip_high = tf.logical_and(
        tf.greater_equal(x, t.cast(clip_max, x.dtype)), tf.greater(signed_grad, 0)
    )
    clip = tf.logical_or(clip_low, clip_high)
    grad = tf.where(clip, mul(grad, 0), grad)

    return grad
//  [python]
def zero_out_clipped_grads(grad, x, clip_min, clip_max):
    """
    Helper function to erase entries in the gradient where the update would be
    clipped.
    :param grad: The gradient
    :param x: The current input
    :param clip_min: Minimum input component value
    :param clip_max: Maximum input component value
    """
    signed_grad = tf.sign(grad)

    # Find input components that lie at the boundary of the input range, and
    # where the gradient points in the wrong direction.
    clip_low = tf.logical_and(
        tf.less_equal(x, tf.cast(clip_min, x.dtype)), tf.less(signed_grad, 0)
    )
    clip_high = tf.logical_and(
        tf.greater_equal(x, tf.cast(clip_max, x.dtype)), tf.greater(signed_grad, 0)
    )
    clip = tf.logical_or(clip_low, clip_high)
    grad = tf.where(clip, mul(grad, 0), grad)

    return grad
//  [python]
def clip_by_value(t, clip_value_min, clip_value_max, name=None):
    """
    A wrapper for clip_by_value that casts the clipping range if needed.
    """

    def cast_clip(clip):
        """
        Cast clipping range argument if needed.
        """
        if t.dtype in (tf.float32, tf.float64):
            if hasattr(clip, "dtype"):
                # Convert to tf dtype in case this is a numpy dtype
                clip_dtype = tf.as_dtype(clip.dtype)
                if clip_dtype != t.dtype:
                    return t.cast(clip, t.dtype)
        return clip

    clip_value_min = cast_clip(clip_value_min)
    clip_value_max = cast_clip(clip_value_max)

    return tf.clip_by_value(t, clip_value_min, clip_value_max, name)
//  [python]
def clip_by_value(t, clip_value_min, clip_value_max, name=None):
    """
    A wrapper for clip_by_value that casts the clipping range if needed.
    """

    def cast_clip(clip):
        """
        Cast clipping range argument if needed.
        """
        if t.dtype in (tf.float32, tf.float64):
            if hasattr(clip, "dtype"):
                # Convert to tf dtype in case this is a numpy dtype
                clip_dtype = tf.as_dtype(clip.dtype)
                if clip_dtype != t.dtype:
                    return tf.cast(clip, t.dtype)
        return clip

    clip_value_min = cast_clip(clip_value_min)
    clip_value_max = cast_clip(clip_value_max)

    return tf.clip_by_value(t, clip_value_min, clip_value_max, name)
//  [python]
def op_with_scalar_cast(a, b, f):
    """
    Builds the graph to compute f(a, b).
    If only one of the two arguments is a scalar and the operation would
    cause a type error without casting, casts the scalar to match the
    tensor.
    :param a: a tf-compatible array or scalar
    :param b: a tf-compatible array or scalar
    """

    try:
        return f(a, b)
    except (TypeError, ValueError):
        pass

    def is_scalar(x):
        """Return True if `x` is a scalar"""
        if hasattr(x, "get_shape"):
            shape = x.get_shape()
            return shape.ndims == 0
        if hasattr(x, "ndim"):
            return x.ndim == 0
        assert isinstance(x, (int, float))
        return True

    a_scalar = is_scalar(a)
    b_scalar = is_scalar(b)

    if a_scalar and b_scalar:
        raise TypeError("Trying to apply " + str(f) + " with mixed types")

    if a_scalar and not b_scalar:
        a = t.cast(a, b.dtype)

    if b_scalar and not a_scalar:
        b = t.cast(b, a.dtype)

    return f(a, b)
//  [python]
def op_with_scalar_cast(a, b, f):
    """
    Builds the graph to compute f(a, b).
    If only one of the two arguments is a scalar and the operation would
    cause a type error without casting, casts the scalar to match the
    tensor.
    :param a: a tf-compatible array or scalar
    :param b: a tf-compatible array or scalar
    """

    try:
        return f(a, b)
    except (TypeError, ValueError):
        pass

    def is_scalar(x):
        """Return True if `x` is a scalar"""
        if hasattr(x, "get_shape"):
            shape = x.get_shape()
            return shape.ndims == 0
        if hasattr(x, "ndim"):
            return x.ndim == 0
        assert isinstance(x, (int, float))
        return True

    a_scalar = is_scalar(a)
    b_scalar = is_scalar(b)

    if a_scalar and b_scalar:
        raise TypeError("Trying to apply " + str(f) + " with mixed types")

    if a_scalar and not b_scalar:
        a = tf.cast(a, b.dtype)

    if b_scalar and not a_scalar:
        b = tf.cast(b, a.dtype)

    return f(a, b)
//  [python]
def main():
    global_batch_size = 1024
    slot_num = 10
    nnz_per_slot = 5

    from tensorflow.python.keras.engine import base_layer_utils

    base_layer_utils.enable_v2_dtype_behavior()

    policy = tf.keras.mixed_precision.experimental.Policy("mixed_float16")
    tf.keras.mixed_precision.experimental.set_policy(policy)

    dataset = utility.get_dataset(
        global_batch_size // hvd.size(), read_batchsize=global_batch_size // hvd.size()
    )

    sok_init_op = sok.Init(global_batch_size=global_batch_size)

    model = utility.SOKDenseDemo(
        max_vocabulary_size_per_gpu=1024,
        embedding_vec_size=8,
        slot_num=slot_num,
        nnz_per_slot=nnz_per_slot,
        num_dense_layers=0,
    )

    optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)
    optimizer = sok.tf.keras.mixed_precision.LossScaleOptimizer(optimizer, 1024)

    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction="none")

    def _replica_loss(labels, logits):
        loss = loss_fn(labels, logits)
        dtype = loss.dtype
        loss = tf.cast(loss, tf.float32)
        loss = tf.nn.compute_average_loss(loss, global_batch_size=global_batch_size)
        return tf.cast(loss, dtype)

    def train_step(inputs, labels):
        logit = model(inputs, training=True)
        loss = _replica_loss(labels, logit)
        scaled_loss = optimizer.get_scaled_loss(loss)
        scaled_gradients = tf.gradients(scaled_loss, model.trainable_variables)
        emb_vars, other_vars = sok.split_embedding_variable_from_others(model.trainable_variables)
        scaled_emb_grads, scaled_other_grads = (
            scaled_gradients[: len(emb_vars)],
            scaled_gradients[len(emb_vars) :],
        )
        emb_grads = optimizer.get_unscaled_gradients(scaled_emb_grads)
        other_grads = optimizer.get_unscaled_gradients(scaled_other_grads)
        other_grads = [hvd.allreduce(grad) for grad in other_grads]
        with sok.OptimizerScope(emb_vars):
            emb_train_op = optimizer.apply_gradients(zip(emb_grads, emb_vars))
        other_train_op = optimizer.apply_gradients(zip(other_grads, other_vars))
        total_loss = hvd.allreduce(loss)
        with tf.control_dependencies([emb_train_op, other_train_op]):
            return tf.identity(total_loss)

    train_iterator = dataset.make_initializable_iterator()
    iterator_init = train_iterator.initializer
    inputs, labels = train_iterator.get_next()

    loss = train_step(inputs, labels)

    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

    with tf.Session() as sess:
        sess.run(sok_init_op)
        sess.run([init_op, iterator_init])

        for step in range(10):
            loss_v = sess.run(loss)
            if hvd.local_rank() == 0:
                print("[INFO]: step {}, loss {}".format(step, loss_v))
//  [python]
def clip_eta(eta, ord, eps):
    """
    Helper function to clip the perturbation to epsilon norm ball.
    :param eta: A tensor with the current perturbation.
    :param ord: Order of the norm (mimics Numpy).
                Possible values: np.inf, 1 or 2.
    :param eps: Epsilon, bound of the perturbation.
    """

    # Clipping perturbation eta to self.ord norm ball
    if ord not in [np.inf, 1, 2]:
        raise ValueError("ord must be np.inf, 1, or 2.")
    reduc_ind = list(xrange(1, len(eta.get_shape())))
    avoid_zero_div = 1e-12
    if ord == np.inf:
        eta = clip_by_value(eta, -eps, eps)
    elif ord == 1:
        # Implements a projection algorithm onto the l1-ball from
        # (Duchi et al. 2008) that runs in time O(d*log(d)) where d is the
        # input dimension.
        # Paper link (Duchi et al. 2008): https://dl.acm.org/citation.cfm?id=1390191

        eps = t.cast(eps, eta.dtype)

        dim = tf.reduce_prod(tf.shape(eta)[1:])
        eta_flat = tf.reshape(eta, (-1, dim))
        abs_eta = tf.abs(eta_flat)

        if "sort" in dir(tf):
            mu = -tf.sort(-abs_eta, axis=-1)
        else:
            # `tf.sort` is only available in TF 1.13 onwards
            mu = tf.nn.top_k(abs_eta, k=dim, sorted=True)[0]
        cumsums = tf.cumsum(mu, axis=-1)
        js = t.cast(tf.divide(1, tf.range(1, dim + 1)), eta.dtype)
        t = t.cast(tf.greater(mu - js * (cumsums - eps), 0), eta.dtype)

        rho = tf.argmax(t * cumsums, axis=-1)
        rho_val = tf.reduce_max(t * cumsums, axis=-1)
        theta = tf.divide(rho_val - eps, t.cast(1 + rho, eta.dtype))

        eta_sgn = tf.sign(eta_flat)
        eta_proj = eta_sgn * tf.maximum(abs_eta - theta[:, tf.newaxis], 0)
        eta_proj = tf.reshape(eta_proj, tf.shape(eta))

        norm = tf.reduce_sum(tf.abs(eta), reduc_ind)
        eta = tf.where(tf.greater(norm, eps), eta_proj, eta)

    elif ord == 2:
        # avoid_zero_div must go inside sqrt to avoid a divide by zero
        # in the gradient through this operation
        norm = tf.sqrt(
            tf.maximum(
                avoid_zero_div, reduce_sum(tf.square(eta), reduc_ind, keepdims=True)
            )
        )
        # We must *clip* to within the norm ball, not *normalize* onto the
        # surface of the ball
        factor = tf.minimum(1.0, div(eps, norm))
        eta = eta * factor
    return eta
//  [python]
def clip_eta(eta, ord, eps):
    """
    Helper function to clip the perturbation to epsilon norm ball.
    :param eta: A tensor with the current perturbation.
    :param ord: Order of the norm (mimics Numpy).
                Possible values: np.inf, 1 or 2.
    :param eps: Epsilon, bound of the perturbation.
    """

    # Clipping perturbation eta to self.ord norm ball
    if ord not in [np.inf, 1, 2]:
        raise ValueError("ord must be np.inf, 1, or 2.")
    reduc_ind = list(xrange(1, len(eta.get_shape())))
    avoid_zero_div = 1e-12
    if ord == np.inf:
        eta = clip_by_value(eta, -eps, eps)
    elif ord == 1:
        # Implements a projection algorithm onto the l1-ball from
        # (Duchi et al. 2008) that runs in time O(d*log(d)) where d is the
        # input dimension.
        # Paper link (Duchi et al. 2008): https://dl.acm.org/citation.cfm?id=1390191

        eps = tf.cast(eps, eta.dtype)

        dim = tf.reduce_prod(tf.shape(eta)[1:])
        eta_flat = tf.reshape(eta, (-1, dim))
        abs_eta = tf.abs(eta_flat)

        if "sort" in dir(tf):
            mu = -tf.sort(-abs_eta, axis=-1)
        else:
            # `tf.sort` is only available in TF 1.13 onwards
            mu = tf.nn.top_k(abs_eta, k=dim, sorted=True)[0]
        cumsums = tf.cumsum(mu, axis=-1)
        js = tf.cast(tf.divide(1, tf.range(1, dim + 1)), eta.dtype)
        t = tf.cast(tf.greater(mu - js * (cumsums - eps), 0), eta.dtype)

        rho = tf.argmax(t * cumsums, axis=-1)
        rho_val = tf.reduce_max(t * cumsums, axis=-1)
        theta = tf.divide(rho_val - eps, tf.cast(1 + rho, eta.dtype))

        eta_sgn = tf.sign(eta_flat)
        eta_proj = eta_sgn * tf.maximum(abs_eta - theta[:, tf.newaxis], 0)
        eta_proj = tf.reshape(eta_proj, tf.shape(eta))

        norm = tf.reduce_sum(tf.abs(eta), reduc_ind)
        eta = tf.where(tf.greater(norm, eps), eta_proj, eta)

    elif ord == 2:
        # avoid_zero_div must go inside sqrt to avoid a divide by zero
        # in the gradient through this operation
        norm = tf.sqrt(
            tf.maximum(
                avoid_zero_div, reduce_sum(tf.square(eta), reduc_ind, keepdims=True)
            )
        )
        # We must *clip* to within the norm ball, not *normalize* onto the
        # surface of the ball
        factor = tf.minimum(1.0, div(eps, norm))
        eta = eta * factor
    return eta
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_mask': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'segment_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'label_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
    }
    if self._include_sentence_id:
      name_to_features['sentence_id'] = tf.io.FixedLenFeature([], tf.int64)
      name_to_features['sub_sentence_id'] = tf.io.FixedLenFeature([], tf.int64)

    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in example:
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example
//  [python]
def _decode(self, record: tf.Tensor):
    """Decodes a serialized tf.Example."""
    name_to_features = {
        'input_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'input_mask': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'segment_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
        'label_ids': tf.io.FixedLenFeature([self._seq_length], tf.int64),
    }
    if self._include_sentence_id:
      name_to_features['sentence_id'] = tf.io.FixedLenFeature([], tf.int64)
      name_to_features['sub_sentence_id'] = tf.io.FixedLenFeature([], tf.int64)

    example = tf.io.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in example:
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example
//  [python]
def _direct_assign(self, value):
        self._value.assign(tf.cast(value, self._value.dtype))
//  [python]
def call(self, inputs):
    # Just eat the circular import.
    from model import ModelBuilder
    # Expect (batch, 219 = 109 + 109 + 1) int64.
    images = inputs["inputs"]
    # Encode (batch, 109) int64 to (batch, 64, 256) float32 for the "before"/"previous" position.
    before = self.encoder_hider.encoder(images[:, 0:ModelBuilder.input_planes_count], training=False)
    # Encode (batch, 109) int64 to (batch, 64, 256) float32 for the "after"/"current" position.
    after = self.encoder_hider.encoder(images[:, ModelBuilder.input_planes_count:2*ModelBuilder.input_planes_count], training=False)
    # Expand (batch, 1) int64 to (batch, 1, 256) of all ones or all zeros.
    side_to_move = tf.cast(images[:, -1:], tf.bool)
    side_to_move = tf.cast(side_to_move, tf.float32)
    side_to_move = tf.tile(side_to_move[:, :, tf.newaxis], [1, 1, ModelBuilder.filter_count])
    # Concatenate encoder outputs to (batch, 129, 256) float32.
    encoder_outputs = tf.concat([before, after, side_to_move], axis=1)
    # Decode.
    inputs = {**inputs, "inputs": encoder_outputs }
    return self.decoder(inputs)
//  [python]
def _embedding_linear(self, embedding_matrix, x):
    """Uses embeddings as linear transformation weights."""
    batch_size = tf.shape(x)[0]
    length = tf.shape(x)[1]
    hidden_size = tf.shape(x)[2]
    vocab_size = tf.shape(embedding_matrix)[0]

    x = tf.reshape(x, [-1, hidden_size])
    logits = tf.matmul(
        tf.to_float(x),
        tf.to_float(embedding_matrix),
        transpose_b=True)

    return tf.reshape(logits, [batch_size, length, vocab_size])
//  [python]
def _embedding_linear(self, embedding_matrix, x):
    """Uses embeddings as linear transformation weights."""
    batch_size = tf.shape(x)[0]
    length = tf.shape(x)[1]
    hidden_size = tf.shape(x)[2]
    vocab_size = tf.shape(embedding_matrix)[0]

    x = tf.reshape(x, [-1, hidden_size])
    logits = tf.matmul(
        tf.cast(x, dtype=self._dtype),
        tf.cast(embedding_matrix, self._dtype),
        transpose_b=True)

    return tf.reshape(logits, [batch_size, length, vocab_size])
//  [python]
def evaluate(sentence):
    sentence = list(sentence[0])
    sentence.insert(0, '<SOS>')
    sentence.append('<EOS>')
    sentence = token_Japanese.texts_to_sequences(sentence)
    # For some reason tokenizing JP returns [[1], [2], [3]] when we want [1, 2, 3]
    sentence = [sentence[i][0] for i in range(len(sentence))]
    sentence = pad([sentence], length = len(Japanese[0]))
    sentence = tf.convert_to_tensor(np.array(sentence),dtype=tf.int64)

    decoder_input = tf.convert_to_tensor(np.array(token_English.texts_to_sequences(['SOS'])),dtype=tf.int64)

    outputs = None
    for i in range(len(Japanese[0])):
        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
            sentence, decoder_input)
        #tf.print(enc_padding_mask.shape, combined_mask.shape, dec_padding_mask.shape)
        # predictions.shape == (batch_size, seq_len, vocab_size)
        predictions, attention_weights, out = transformer([sentence,
                                                     decoder_input,
                                                     enc_padding_mask,
                                                     combined_mask,
                                                     dec_padding_mask],
                                                     False)

        # select the last word from the seq_len dimension
        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

        predicted_id = tf.to_int64(tf.argmax(predictions, axis=-1))
        
        #tf.print(predicted_id,decoder_input)
        # return the result if the predicted_id is equal to the end token
        if predicted_id == token_English.texts_to_sequences(['EOS']):
            out['input_sentence'] = sentence.numpy()
            return tf.squeeze(decoder_input, axis=0), attention_weights , sentence, out

        # concatentate the predicted_id to the output which is given to the decoder
        # as its input.
        decoder_input = tf.concat([decoder_input, predicted_id], axis=1)
        outputs = out

    outputs['input_sentence'] = sentence.numpy()
    return tf.squeeze(decoder_input, axis=0), attention_weights, sentence, outputs