jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
        """
        Forward  pass for the Decoder
        
        Arguments:
            x -- Tensor of shape (batch_size, target_seq_len, fully_connected_dim)
            enc_output --  Tensor of shape(batch_size, input_seq_len, fully_connected_dim)
            training -- Boolean, set to true to activate
                        the training mode for dropout layers
            look_ahead_mask -- Boolean mask for the target_input
            padding_mask -- Boolean mask for the second multihead attention layer
        Returns:
            x -- Tensor of shape (batch_size, target_seq_len, fully_connected_dim)
            attention_weights - Dictionary of tensors containing all the attention weights
                                each of shape Tensor of shape (batch_size, num_heads, target_seq_len, input_seq_len)
        """

        seq_len = tf.shape(x)[1]
        attention_weights = {}
        
         # START CODE HERE
        # create word embeddings 
        x = self.embedding(x)  # (batch_size, target_seq_len, fully_connected_dim)
        
        # scale embeddings by multiplying by the square root of their dimension
        x *= tf.math.sqrt(tf.cast(self.embedding_dim, tf.float32))
        
        # calculate positional encodings and add to word embedding
        x += self.pos_encoding[:, :seq_len, :]

        # apply a dropout layer to x
        x = self.dropout(x, training = training)

        # use a for loop to pass x through a stack of decoder layers and update attention_weights (~4 lines total)
        for i in range(self.num_layers):
            # pass x and the encoder output through a stack of decoder layers and save the attention weights
            # of block 1 and 2 (~1 line)
            x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                                 look_ahead_mask, padding_mask)

            #update attention_weights dictionary with the attention weights of block 1 and block 2
            attention_weights['decoder_layer{}_block1_self_att'.format(i+1)] = block1
            attention_weights['decoder_layer{}_block2_decenc_att'.format(i+1)] = block2
        # END CODE HERE
        
        
        # x.shape == (batch_size, target_seq_len, fully_connected_dim)
        return x, attention_weights
//  [python]
def _percentile(self, grad_history: tf.Tensor) -> tf.Tensor:
        """ Compute the clip percentile of the gradient history

        Parameters
        ----------
        grad_history: :class:`tensorflow.Tensor`
            Tge gradient history to calculate the clip percentile for

        Returns
        -------
        :class:`tensorflow.Tensor`
            A rank(:attr:`clip_percentile`) `Tensor`

        Notes
        -----
        Adapted from
        https://github.com/tensorflow/probability/blob/r0.14/tensorflow_probability/python/stats/quantiles.py
        to remove reliance on full tensorflow_probability libraray
        """
        with tf.name_scope("percentile"):
            frac_at_q_or_below = self._clip_percentile / 100.
            sorted_hist = tf.sort(grad_history, axis=-1, direction="ASCENDING")

            num = tf.to_float(tf.shape(grad_history)[-1])

            # get indices
            indices = tf.round((num - 1) * frac_at_q_or_below)
            indices = tf.clip_by_value(tf.to_int32(indices),
                                       0,
                                       tf.shape(grad_history)[-1] - 1)
            gathered_hist = tf.gather(sorted_hist, indices, axis=-1)

            # Propagate NaNs. Apparently tf.is_nan doesn't like other dtypes
            nan_batch_members = tf.reduce_any(tf.math.is_nan(grad_history), axis=None)
            right_rank_matched_shape = tf.pad(tf.shape(nan_batch_members),
                                              paddings=[[0, tf.rank(self._clip_percentile)]],
                                              constant_values=1)
            nan_batch_members = tf.reshape(nan_batch_members, shape=right_rank_matched_shape)

            nan = np.array(np.nan, gathered_hist.dtype.as_numpy_dtype)
            gathered_hist = tf.where(nan_batch_members, nan, gathered_hist)

            return gathered_hist
//  [python]
def _percentile(self, grad_history: tf.Tensor) -> tf.Tensor:
        """ Compute the clip percentile of the gradient history

        Parameters
        ----------
        grad_history: :class:`tensorflow.Tensor`
            Tge gradient history to calculate the clip percentile for

        Returns
        -------
        :class:`tensorflow.Tensor`
            A rank(:attr:`clip_percentile`) `Tensor`

        Notes
        -----
        Adapted from
        https://github.com/tensorflow/probability/blob/r0.14/tensorflow_probability/python/stats/quantiles.py
        to remove reliance on full tensorflow_probability libraray
        """
        with tf.name_scope("percentile"):
            frac_at_q_or_below = self._clip_percentile / 100.
            sorted_hist = tf.sort(grad_history, axis=-1, direction="ASCENDING")

            num = tf.cast(tf.shape(grad_history)[-1], tf.float64)

            # get indices
            indices = tf.round((num - 1) * frac_at_q_or_below)
            indices = tf.clip_by_value(tf.cast(indices, tf.int32),
                                       0,
                                       tf.shape(grad_history)[-1] - 1)
            gathered_hist = tf.gather(sorted_hist, indices, axis=-1)

            # Propagate NaNs. Apparently tf.is_nan doesn't like other dtypes
            nan_batch_members = tf.reduce_any(tf.math.is_nan(grad_history), axis=None)
            right_rank_matched_shape = tf.pad(tf.shape(nan_batch_members),
                                              paddings=[[0, tf.rank(self._clip_percentile)]],
                                              constant_values=1)
            nan_batch_members = tf.reshape(nan_batch_members, shape=right_rank_matched_shape)

            nan = np.array(np.nan, gathered_hist.dtype.as_numpy_dtype)
            gathered_hist = tf.where(nan_batch_members, nan, gathered_hist)

            return gathered_hist
//  [python]
def __call__(self, att_ws_stack, samples):
        output_length = samples["output_length"]
        input_length = samples["input_length"]
        reduction_output_length = (output_length - 1) // self.reduction_factor + 1
        # attn_masks shape: [batch_size, 1, reduction_y_steps, x_steps]
        attn_masks = self._create_attention_masks(input_length, reduction_output_length)
        # length_masks shape: [batch_size, 1, reduction_y_steps, x_steps]
        length_masks = self._create_length_masks(input_length, reduction_output_length)
        att_ws_stack = t.cast(att_ws_stack, dtype=t.float32)
        if len(t.shape(att_ws_stack)) == 3:
            att_ws_stack = t.expand_dims(att_ws_stack, axis=1)
        losses = attn_masks * att_ws_stack
        losses *= length_masks
        loss = t.reduce_sum(losses)
        total_size = t.cast(t.reduce_sum(length_masks), dtype=t.float32)
        return self.guided_attn_weight * loss / total_size
//  [python]
def __call__(self, att_ws_stack, samples):
        output_length = samples["output_length"]
        input_length = samples["input_length"]
        reduction_output_length = (output_length - 1) // self.reduction_factor + 1
        # attn_masks shape: [batch_size, 1, reduction_y_steps, x_steps]
        attn_masks = self._create_attention_masks(input_length, reduction_output_length)
        # length_masks shape: [batch_size, 1, reduction_y_steps, x_steps]
        length_masks = self._create_length_masks(input_length, reduction_output_length)
        att_ws_stack = tf.cast(att_ws_stack, dtype=tf.float32)
        if len(tf.shape(att_ws_stack)) == 3:
            att_ws_stack = tf.expand_dims(att_ws_stack, axis=1)
        losses = attn_masks * att_ws_stack
        losses *= length_masks
        loss = tf.reduce_sum(losses)
        total_size = tf.cast(tf.reduce_sum(length_masks), dtype=tf.float32)
        return self.guided_attn_weight * loss / total_size
//  [python]
def _create_attention_masks(self, input_length, output_length):
        """masks created by attention location

        Args:
            input_length: shape: [batch_size]
            output_length: shape: [batch_size]

        Returns:
            masks: shape: [batch_size, 1, y_steps, x_steps]
        """
        batch_size = tf.shape(input_length)[0]
        input_max_len = tf.reduce_max(input_length)
        output_max_len = tf.reduce_max(output_length)

        # grid_x shape: [output_max_len, input_max_len]
        grid_x, grid_y = tf.meshgrid(tf.range(output_max_len),
                                     tf.range(input_max_len),
                                     indexing='ij')

        grid_x = tf.tile(grid_x[tf.newaxis, :, :], [batch_size, 1, 1])
        grid_y = tf.tile(grid_y[tf.newaxis, :, :], [batch_size, 1, 1])
        input_length = input_length[:, tf.newaxis, tf.newaxis]
        output_length = output_length[:, tf.newaxis, tf.newaxis]
        # masks shape: [batch_size, y_steps, x_steps]
        masks = 1.0 - tf.math.exp(-(grid_y / input_length - grid_x / output_length) ** 2
                                 / (2 * (self.attn_sigma ** 2)))
        masks = tf.expand_dims(masks, axis=1)
        masks = tf.to_float(masks)
        return masks
//  [python]
def _create_attention_masks(self, input_length, output_length):
        """masks created by attention location

        Args:
            input_length: shape: [batch_size]
            output_length: shape: [batch_size]

        Returns:
            masks: shape: [batch_size, 1, y_steps, x_steps]
        """
        batch_size = tf.shape(input_length)[0]
        input_max_len = tf.reduce_max(input_length)
        output_max_len = tf.reduce_max(output_length)

        # grid_x shape: [output_max_len, input_max_len]
        grid_x, grid_y = tf.meshgrid(tf.range(output_max_len),
                                     tf.range(input_max_len),
                                     indexing='ij')

        grid_x = tf.tile(grid_x[tf.newaxis, :, :], [batch_size, 1, 1])
        grid_y = tf.tile(grid_y[tf.newaxis, :, :], [batch_size, 1, 1])
        input_length = input_length[:, tf.newaxis, tf.newaxis]
        output_length = output_length[:, tf.newaxis, tf.newaxis]
        # masks shape: [batch_size, y_steps, x_steps]
        masks = 1.0 - tf.math.exp(-(grid_y / input_length - grid_x / output_length) ** 2
                                 / (2 * (self.attn_sigma ** 2)))
        masks = tf.expand_dims(masks, axis=1)
        masks = tf.cast(masks, dtype=tf.float32)
        return masks
//  [python]
def make_decay_function(step, lr_list):
  """Decay function, stepwise linear from 0.1 to 0.01 of max lr.

  Args:
    step: training step.
    lr_list: [start_decay, end_decay, rate_0, rate_1, ...]

  Returns:
    lr_func: decayed lr function.
    lambda: 10-4 * lr_func()
  """

  start, end = lr_list[0], lr_list[1]
  lr_seq = lr_list[2:]
  seq_num = float(len(lr_seq) - 1)
  seq_len = float(end - start) / seq_num

  def decayed_lr():
    if step < start:
      return lr_seq[0]
    elif step > end:
      return lr_seq[-1]
    else:
      step_trunc = t.cast(step - start, t.float32)
      num = t.cast(step_trunc / seq_len, t.int64)
      lr_start = t.gather(lr_seq, num)
      lr_end = t.gather(lr_seq, num + 1)
      step_start = t.cast(num, t.float32) * seq_len
      alpha = (step_trunc - step_start) / seq_len
      return alpha * lr_end + (1.0 - alpha) * lr_start

  lr_func = t.function(func=decayed_lr)
  return lr_func, lambda: 1e-4 * lr_func()
//  [python]
def make_decay_function(step, lr_list):
  """Decay function, stepwise linear from 0.1 to 0.01 of max lr.

  Args:
    step: training step.
    lr_list: [start_decay, end_decay, rate_0, rate_1, ...]

  Returns:
    lr_func: decayed lr function.
    lambda: 10-4 * lr_func()
  """

  start, end = lr_list[0], lr_list[1]
  lr_seq = lr_list[2:]
  seq_num = float(len(lr_seq) - 1)
  seq_len = float(end - start) / seq_num

  def decayed_lr():
    if step < start:
      return lr_seq[0]
    elif step > end:
      return lr_seq[-1]
    else:
      step_trunc = tf.cast(step - start, tf.float32)
      num = tf.cast(step_trunc / seq_len, tf.int64)
      lr_start = tf.gather(lr_seq, num)
      lr_end = tf.gather(lr_seq, num + 1)
      step_start = tf.cast(num, tf.float32) * seq_len
      alpha = (step_trunc - step_start) / seq_len
      return alpha * lr_end + (1.0 - alpha) * lr_start

  lr_func = tf.function(func=decayed_lr)
  return lr_func, lambda: 1e-4 * lr_func()
//  [python]
def _gather_batched(self, labels, match_indices, mask, mask_val):
    """Gather based on batched labels."""
    batch_size = labels.shape[0]
    if batch_size == 1:
      if mask is not None:
        result = self._gather_unbatched(
            tf.squeeze(labels, axis=0), tf.squeeze(match_indices, axis=0),
            tf.squeeze(mask, axis=0), mask_val)
      else:
        result = self._gather_unbatched(
            tf.squeeze(labels, axis=0), tf.squeeze(match_indices, axis=0),
            None, mask_val)
      return tf.expand_dims(result, axis=0)
    else:
      indices_shape = tf.shape(match_indices)
      indices_dtype = match_indices.dtype
      batch_indices = (tf.expand_dims(
          tf.range(indices_shape[0], dtype=indices_dtype), axis=-1) *
                       tf.ones([1, indices_shape[-1]], dtype=indices_dtype))
      gather_nd_indices = tf.stack(
          [batch_indices, match_indices], axis=-1)
      targets = tf.gather_nd(labels, gather_nd_indices)
      if mask is None:
        return targets
      else:
        masked_targets = tf.to_float(mask_val) * tf.ones_like(
            mask, dtype=labels.dtype)
        return tf.where(mask, masked_targets, targets)
//  [python]
def _gather_batched(self, labels, match_indices, mask, mask_val):
    """Gather based on batched labels."""
    batch_size = labels.shape[0]
    if batch_size == 1:
      if mask is not None:
        result = self._gather_unbatched(
            tf.squeeze(labels, axis=0), tf.squeeze(match_indices, axis=0),
            tf.squeeze(mask, axis=0), mask_val)
      else:
        result = self._gather_unbatched(
            tf.squeeze(labels, axis=0), tf.squeeze(match_indices, axis=0),
            None, mask_val)
      return tf.expand_dims(result, axis=0)
    else:
      indices_shape = tf.shape(match_indices)
      indices_dtype = match_indices.dtype
      batch_indices = (tf.expand_dims(
          tf.range(indices_shape[0], dtype=indices_dtype), axis=-1) *
                       tf.ones([1, indices_shape[-1]], dtype=indices_dtype))
      gather_nd_indices = tf.stack(
          [batch_indices, match_indices], axis=-1)
      targets = tf.gather_nd(labels, gather_nd_indices)
      if mask is None:
        return targets
      else:
        masked_targets = tf.cast(mask_val, labels.dtype) * tf.ones_like(
            mask, dtype=labels.dtype)
        return tf.where(mask, masked_targets, targets)
//  [python]
def __call__(self, outputs, samples):
        """
        Its corresponding log value is calculated to make it Gaussian.
        Args:
            outputs: it contains four elements:
                before_outs: outputs before postnet, shape: [batch, y_steps, feat_dim]
                teacher_outs: teacher outputs, shape: [batch, y_steps, feat_dim]
                after_outs: outputs after postnet, shape: [batch, y_steps, feat_dim]
                duration_sequences: duration predictions from teacher model, shape: [batch, x_steps]
                pred_duration_sequences: duration predictions from trained predictor
                    shape: [batch, x_steps]
            samples: samples from dataset

        """
        final_loss = {}
        before_outs, teacher_outs, after_outs, duration_sequences, pred_duration_sequences = outputs
        # perform masking for padded values
        output_length = samples["output_length"]
        output = samples["output"] # [batch, y_steps, feat_dim]
        input_length = samples['input_length']
        x_steps = tf.reduce_max(input_length)
        y_steps = tf.shape(output)[1]
        feat_dim = tf.shape(output)[2]
        if self.use_mask:
            mask = tf.sequence_mask(output_length, y_steps, dtype=tf.float32)
            mask = tf.tile(tf.expand_dims(mask, axis=-1), [1, 1, feat_dim])
        else:
            mask = tf.ones_like(output)
        total_size = tf.cast(tf.reduce_sum(mask), dtype=tf.float32)

        if self.teacher_guide:
            l1_loss = tf.abs(after_outs - teacher_outs) + tf.abs(before_outs - teacher_outs)
        else:
            l1_loss = tf.abs(after_outs - output) + tf.abs(before_outs - output)
        l1_loss *= mask
        final_loss['l1_loss'] = tf.reduce_sum(l1_loss) / total_size

        teacher_duration = tf.math.log(tf.cast(duration_sequences, dtype=tf.float32) + self.eps)
        if self.use_mask:
            mask = tf.sequence_mask(input_length, x_steps, dtype=tf.float32)
        else:
            mask = tf.ones_like(teacher_duration)
        total_size = tf.cast(tf.reduce_sum(mask), dtype=tf.float32)
        duration_loss = tf.square(teacher_duration - pred_duration_sequences)
        duration_loss *= mask
        final_loss['duration_loss'] = tf.reduce_sum(duration_loss) / total_size

        return final_loss
//  [python]
def __call__(self, outputs, samples):
        final_loss = {}
        before_outs, teacher_outs, after_outs, duration_sequences, pred_duration_sequences, pred_f0, pred_energy = outputs
        # perform masking for padded values
        output_length = samples["output_length"]
        output = samples["output"] # [batch, y_steps, feat_dim]
        input_length = samples['input_length']
        f0_gst = samples["f0"]
        energy_gst = samples["energy"]
        x_steps = tf.reduce_max(input_length)
        y_steps = tf.shape(output)[1]
        feat_dim = tf.shape(output)[2]
        if self.use_mask:
            mask = tf.sequence_mask(output_length, y_steps, dtype=tf.float32)
            mask = tf.tile(tf.expand_dims(mask, axis=-1), [1, 1, feat_dim])
        else:
            mask = tf.ones_like(output)
        total_size = tf.cast(tf.reduce_sum(mask), dtype=tf.float32)

        l1_loss = tf.abs(after_outs - output) + tf.abs(before_outs - output)
        l1_loss *= mask
        final_loss['l1_loss'] = tf.reduce_sum(l1_loss) / total_size

        teacher_duration = tf.math.log(tf.cast(duration_sequences, dtype=tf.float32) + self.eps)
        if self.use_mask:
            mask = tf.sequence_mask(input_length, x_steps, dtype=tf.float32)
        else:
            mask = tf.ones_like(teacher_duration)
        total_size = tf.cast(tf.reduce_sum(mask), dtype=tf.float32)
        duration_loss = tf.square(teacher_duration - pred_duration_sequences)
        duration_loss *= mask
        final_loss['duration_loss'] = tf.reduce_sum(duration_loss) / total_size

        f0_loss = tf.square(f0_gst - pred_f0)
        f0_loss *= mask
        final_loss['f0_loss'] = tf.reduce_sum(f0_loss) / total_size

        energy_loss = tf.square(energy_gst - pred_energy)
        energy_loss *= mask
        final_loss['energy_loss'] = tf.reduce_sum(energy_loss) / total_size
        return final_loss
//  [python]
def __call__(self, outputs, samples, logit_length=None):
        """
        Args:
            outputs: contain elements below:
                att_ws_stack: shape: [batch, y_steps, x_steps]
        """
        final_loss = {}
        before_outs, after_outs, logits_stack, att_ws_stack = outputs
        # perform masking for padded values
        output_length = samples["output_length"]
        output = samples["output"]
        y_steps = tf.shape(output)[1]
        batch = tf.shape(output)[0]
        feat_dim = tf.shape(output)[2]
        if self.mask_decoder:
            mask = tf.sequence_mask(
                output_length, y_steps, dtype=tf.float32
            )
            mask = tf.tile(tf.expand_dims(mask, axis=-1), [1, 1, feat_dim])
        else:
            mask = tf.ones_like(output)
        total_size = tf.to_float(tf.reduce_sum(mask))
        if self.l1_loss_weight > 0:
            l1_loss = tf.abs(after_outs - output) + tf.abs(before_outs - output)
            l1_loss *= mask
            final_loss['l1_loss'] = tf.reduce_sum(l1_loss) / total_size * \
                                    self.l1_loss_weight
        mse_loss = tf.square(after_outs - output) + tf.square(before_outs - output)
        mse_loss *= mask

        indexes = tf.tile(tf.range(y_steps)[tf.newaxis, :], [batch, 1]) # [batch, y_steps]
        end_index = tf.tile((output_length - 1)[:, tf.newaxis], [1, y_steps])
        zeroes = tf.zeros_like(indexes, dtype=tf.float32)
        ones = tf.ones_like(indexes, dtype=tf.float32)
        labels = tf.where(end_index <= indexes, ones, zeroes) # [batch, y_steps]
        # bce_loss is used for stop token prediction
        bce_loss = tf.nn.weighted_cross_entropy_with_logits(labels=labels,
                                                            logits=logits_stack[:, :, 0],
                                                            pos_weight=self.pos_weight)
        bce_loss = bce_loss[:, :, tf.newaxis]
        bce_loss *= mask
        final_loss['mse_loss'] = tf.reduce_sum(mse_loss) / total_size
        final_loss['bce_loss'] = tf.reduce_sum(bce_loss) / total_size
        if self.guided_attn_loss_function is not None and \
                self.guided_attn_loss_function.guided_attn_weight > 0:
            final_loss['guided_attn_loss'] = self.guided_attn_loss_function(att_ws_stack, samples)
        if self.regularization_weight > 0:
            computed_vars = [var for var in self.model.trainable_variables
                             if 'bias' not in var.name and \
                                'projection' not in var.name and \
                                'e_function' not in var.name and \
                                'embedding' not in var.name and \
                                'rnn' not in var.name and \
                                'zone_out' not in var.name]
            final_loss['regularization_loss'] = tf.add_n([tf.nn.l2_loss(v) for v in computed_vars]) * \
                          self.regularization_weight

        return final_loss
//  [python]
def __call__(self, outputs, samples, logit_length=None):
        """
        Args:
            outputs: contain elements below:
                att_ws_stack: shape: [batch, y_steps, x_steps]
        """
        final_loss = {}
        before_outs, after_outs, logits_stack, att_ws_stack = outputs
        # perform masking for padded values
        output_length = samples["output_length"]
        output = samples["output"]
        y_steps = tf.shape(output)[1]
        batch = tf.shape(output)[0]
        feat_dim = tf.shape(output)[2]
        if self.mask_decoder:
            mask = tf.sequence_mask(
                output_length, y_steps, dtype=tf.float32
            )
            mask = tf.tile(tf.expand_dims(mask, axis=-1), [1, 1, feat_dim])
        else:
            mask = tf.ones_like(output)
        total_size = tf.cast(tf.reduce_sum(mask), dtype=tf.float32)
        if self.l1_loss_weight > 0:
            l1_loss = tf.abs(after_outs - output) + tf.abs(before_outs - output)
            l1_loss *= mask
            final_loss['l1_loss'] = tf.reduce_sum(l1_loss) / total_size * \
                                    self.l1_loss_weight
        mse_loss = tf.square(after_outs - output) + tf.square(before_outs - output)
        mse_loss *= mask

        indexes = tf.tile(tf.range(y_steps)[tf.newaxis, :], [batch, 1]) # [batch, y_steps]
        end_index = tf.tile((output_length - 1)[:, tf.newaxis], [1, y_steps])
        zeroes = tf.zeros_like(indexes, dtype=tf.float32)
        ones = tf.ones_like(indexes, dtype=tf.float32)
        labels = tf.where(end_index <= indexes, ones, zeroes) # [batch, y_steps]
        # bce_loss is used for stop token prediction
        bce_loss = tf.nn.weighted_cross_entropy_with_logits(labels=labels,
                                                            logits=logits_stack[:, :, 0],
                                                            pos_weight=self.pos_weight)
        bce_loss = bce_loss[:, :, tf.newaxis]
        bce_loss *= mask
        final_loss['mse_loss'] = tf.reduce_sum(mse_loss) / total_size
        final_loss['bce_loss'] = tf.reduce_sum(bce_loss) / total_size
        if self.guided_attn_loss_function is not None and \
                self.guided_attn_loss_function.guided_attn_weight > 0:
            final_loss['guided_attn_loss'] = self.guided_attn_loss_function(att_ws_stack, samples)
        if self.regularization_weight > 0:
            computed_vars = [var for var in self.model.trainable_variables
                             if 'bias' not in var.name and \
                                'projection' not in var.name and \
                                'e_function' not in var.name and \
                                'embedding' not in var.name and \
                                'rnn' not in var.name and \
                                'zone_out' not in var.name]
            final_loss['regularization_loss'] = tf.add_n([tf.nn.l2_loss(v) for v in computed_vars]) * \
                          self.regularization_weight

        return final_loss
//  [python]
def frechet_inception_distance(avg_pools1, avg_pools2):
  """FID for evaluating a generative model."""
  m = tf.reduce_mean(avg_pools1, axis=0)
  m_w = tf.reduce_mean(avg_pools2, axis=0)

  # Calculate the unbiased covariance matrix of first activations.
  num_examples_real = tf.to_float(tf.shape(input=avg_pools1)[0])
  sigma = (
      num_examples_real / (num_examples_real - 1) *
      tfp.stats.covariance(avg_pools1))

  # Calculate the unbiased covariance matrix of second activations.
  num_examples_generated = tf.to_float(tf.shape(input=avg_pools2)[0])
  sigma_w = (
      num_examples_generated / (num_examples_generated - 1) *
      tfp.stats.covariance(avg_pools2))

  # Find the Tr(sqrt(sigma sigma_w)) component of FID
  sqrt_trace_component = _trace_sqrt_product(sigma, sigma_w)

  # First the covariance component.
  # Here, note that trace(A + B) = trace(A) + trace(B)
  trace = tf.linalg.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component

  # Next the distance between means.
  mean = tf.reduce_sum(tf.math.squared_difference(m, m_w))
  fid = trace + mean
  return fid
//  [python]
def frechet_inception_distance(avg_pools1, avg_pools2):
  """FID for evaluating a generative model."""
  m = tf.reduce_mean(avg_pools1, axis=0)
  m_w = tf.reduce_mean(avg_pools2, axis=0)

  # Calculate the unbiased covariance matrix of first activations.
  num_examples_real = tf.cast(tf.shape(input=avg_pools1)[0], tf.float32)
  sigma = (
      num_examples_real / (num_examples_real - 1) *
      tfp.stats.covariance(avg_pools1))

  # Calculate the unbiased covariance matrix of second activations.
  num_examples_generated = tf.cast(tf.shape(input=avg_pools2)[0], tf.float32)
  sigma_w = (
      num_examples_generated / (num_examples_generated - 1) *
      tfp.stats.covariance(avg_pools2))

  # Find the Tr(sqrt(sigma sigma_w)) component of FID
  sqrt_trace_component = _trace_sqrt_product(sigma, sigma_w)

  # First the covariance component.
  # Here, note that trace(A + B) = trace(A) + trace(B)
  trace = tf.linalg.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component

  # Next the distance between means.
  mean = tf.reduce_sum(tf.math.squared_difference(m, m_w))
  fid = trace + mean
  return fid
//  [python]
def _reduce_panoptic_metrics(self, logs: Dict[str, Any]):
    """Updates the per class and mean panoptic metrics in the logs."""
    result = self.panoptic_quality_metric.result()
    valid_thing_classes = result['valid_thing_classes']
    valid_stuff_classes = result['valid_stuff_classes']
    valid_classes = valid_stuff_classes | valid_thing_classes
    num_categories = tf.math.count_nonzero(valid_classes, dtype=tf.float32)
    num_thing_categories = tf.math.count_nonzero(
        valid_thing_classes, dtype=tf.float32
    )
    num_stuff_categories = tf.math.count_nonzero(
        valid_stuff_classes, dtype=tf.float32
    )
    valid_thing_classes = tf.to_float(valid_thing_classes)
    valid_stuff_classes = tf.to_float(valid_stuff_classes)

    logs['panoptic_quality/All_num_categories'] = num_categories
    logs['panoptic_quality/Things_num_categories'] = num_thing_categories
    logs['panoptic_quality/Stuff_num_categories'] = num_stuff_categories
    for metric in ['pq', 'sq', 'rq']:
      metric_per_class = result[f'{metric}_per_class']
      logs[f'panoptic_quality/All_{metric}'] = tf.math.divide_no_nan(
          tf.reduce_sum(metric_per_class), num_categories
      )
      logs[f'panoptic_quality/Things_{metric}'] = tf.math.divide_no_nan(
          tf.reduce_sum(metric_per_class * valid_thing_classes),
          num_thing_categories,
      )
      logs[f'panoptic_quality/Stuff_{metric}'] = tf.math.divide_no_nan(
          tf.reduce_sum(metric_per_class * valid_stuff_classes),
          num_stuff_categories,
      )
      if self.task_config.panoptic_quality_evaluator.report_per_class_metrics:
        for i, is_valid in enumerate(valid_classes.numpy()):
          if is_valid:
            logs[f'panoptic_quality/{metric}/class_{i}'] = metric_per_class[i]
//  [python]
def _reduce_panoptic_metrics(self, logs: Dict[str, Any]):
    """Updates the per class and mean panoptic metrics in the logs."""
    result = self.panoptic_quality_metric.result()
    valid_thing_classes = result['valid_thing_classes']
    valid_stuff_classes = result['valid_stuff_classes']
    valid_classes = valid_stuff_classes | valid_thing_classes
    num_categories = tf.math.count_nonzero(valid_classes, dtype=tf.float32)
    num_thing_categories = tf.math.count_nonzero(
        valid_thing_classes, dtype=tf.float32
    )
    num_stuff_categories = tf.math.count_nonzero(
        valid_stuff_classes, dtype=tf.float32
    )
    valid_thing_classes = tf.cast(valid_thing_classes, dtype=tf.float32)
    valid_stuff_classes = tf.cast(valid_stuff_classes, dtype=tf.float32)

    logs['panoptic_quality/All_num_categories'] = num_categories
    logs['panoptic_quality/Things_num_categories'] = num_thing_categories
    logs['panoptic_quality/Stuff_num_categories'] = num_stuff_categories
    for metric in ['pq', 'sq', 'rq']:
      metric_per_class = result[f'{metric}_per_class']
      logs[f'panoptic_quality/All_{metric}'] = tf.math.divide_no_nan(
          tf.reduce_sum(metric_per_class), num_categories
      )
      logs[f'panoptic_quality/Things_{metric}'] = tf.math.divide_no_nan(
          tf.reduce_sum(metric_per_class * valid_thing_classes),
          num_thing_categories,
      )
      logs[f'panoptic_quality/Stuff_{metric}'] = tf.math.divide_no_nan(
          tf.reduce_sum(metric_per_class * valid_stuff_classes),
          num_stuff_categories,
      )
      if self.task_config.panoptic_quality_evaluator.report_per_class_metrics:
        for i, is_valid in enumerate(valid_classes.numpy()):
          if is_valid:
            logs[f'panoptic_quality/{metric}/class_{i}'] = metric_per_class[i]
//  [python]
def call(self, scores, mask=None):

    if mask is not None:
      for _ in range(len(scores.shape) - len(mask.shape)):
        mask = tf.expand_dims(mask, axis=self._mask_expansion_axes)

      # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
      # masked positions, this operation will create a tensor which is 0.0 for
      # positions we want to attend and -1.e9 for masked positions.
      adder = (1.0 - tf.to_float(mask)) * _large_compatible_negative(
          scores.dtype)
      # Since we are adding it to the raw scores before the softmax, this is
      # effectively the same as removing these entirely.
      scores += adder

    if len(self._normalization_axes) == 1:
      return tf.nn.softmax(scores, axis=self._normalization_axes[0])
    else:
      return tf.math.exp(scores - tf.math.reduce_logsumexp(
          scores, axis=self._normalization_axes, keepdims=True))
//  [python]
def call(self, scores, mask=None):

    if mask is not None:
      for _ in range(len(scores.shape) - len(mask.shape)):
        mask = tf.expand_dims(mask, axis=self._mask_expansion_axes)

      # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
      # masked positions, this operation will create a tensor which is 0.0 for
      # positions we want to attend and -1.e9 for masked positions.
      adder = (1.0 - tf.cast(mask, scores.dtype)) * _large_compatible_negative(
          scores.dtype)
      # Since we are adding it to the raw scores before the softmax, this is
      # effectively the same as removing these entirely.
      scores += adder

    if len(self._normalization_axes) == 1:
      return tf.nn.softmax(scores, axis=self._normalization_axes[0])
    else:
      return tf.math.exp(scores - tf.math.reduce_logsumexp(
          scores, axis=self._normalization_axes, keepdims=True))
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync

    with tf.GradientTape() as tape:
      model_kwargs = {
          'image_info': labels['image_info'],
          'anchor_boxes': labels['anchor_boxes'],
          'gt_boxes': labels['gt_boxes'],
          'gt_classes': labels['gt_classes'],
          'training': True,
      }
      if self.task_config.model.include_mask:
        model_kwargs['gt_masks'] = labels['gt_masks']
        if self.task_config.model.outer_boxes_scale > 1.0:
          model_kwargs['gt_outer_boxes'] = labels['gt_outer_boxes']
      outputs = model(images, **model_kwargs)
      outputs = tf.nest.map_structure(
          lambda x: tf.to_float(x), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses)
      scaled_loss = losses['total_loss'] / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])

    if (self.task_config.segmentation_evaluation.report_train_mean_iou and
        self.segmentation_train_mean_iou is not None):
      segmentation_labels = {
          'masks': labels['gt_segmentation_mask'],
          'valid_masks': labels['gt_segmentation_valid_mask'],
          'image_info': labels['image_info']
      }
      self.process_metrics(
          metrics=[self.segmentation_train_mean_iou],
          labels=segmentation_labels,
          model_outputs=outputs['segmentation_outputs'])
      logs.update({
          self.segmentation_train_mean_iou.name:
              self.segmentation_train_mean_iou.result()
      })

    return logs
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync

    with tf.GradientTape() as tape:
      model_kwargs = {
          'image_info': labels['image_info'],
          'anchor_boxes': labels['anchor_boxes'],
          'gt_boxes': labels['gt_boxes'],
          'gt_classes': labels['gt_classes'],
          'training': True,
      }
      if self.task_config.model.include_mask:
        model_kwargs['gt_masks'] = labels['gt_masks']
        if self.task_config.model.outer_boxes_scale > 1.0:
          model_kwargs['gt_outer_boxes'] = labels['gt_outer_boxes']
      outputs = model(images, **model_kwargs)
      outputs = tf.nest.map_structure(
          lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses)
      scaled_loss = losses['total_loss'] / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])

    if (self.task_config.segmentation_evaluation.report_train_mean_iou and
        self.segmentation_train_mean_iou is not None):
      segmentation_labels = {
          'masks': labels['gt_segmentation_mask'],
          'valid_masks': labels['gt_segmentation_valid_mask'],
          'image_info': labels['image_info']
      }
      self.process_metrics(
          metrics=[self.segmentation_train_mean_iou],
          labels=segmentation_labels,
          model_outputs=outputs['segmentation_outputs'])
      logs.update({
          self.segmentation_train_mean_iou.name:
              self.segmentation_train_mean_iou.result()
      })

    return logs
//  [python]
def _get_NN_prediction(self, state):
        assert state.shape.rank == 5  # Batch, H, W, Channel, History
        state = tf.transpose(state, [0, 1, 2, 4, 3])  # swap channel & history, to be compatible with old models
        image = tf.reshape(state, [-1] + list(STATE_SHAPE[:2]) + [STATE_SHAPE[2] * FRAME_HISTORY])
        image = tf.cast(image, tf.float32)

        image = image / 255.0
        with argscope(Conv2D, activation=tf.nn.relu):
            l = Conv2D('conv0', image, 32, 5)
            l = MaxPooling('pool0', l, 2)
            l = Conv2D('conv1', l, 32, 5)
            l = MaxPooling('pool1', l, 2)
            l = Conv2D('conv2', l, 64, 4)
            l = MaxPooling('pool2', l, 2)
            l = Conv2D('conv3', l, 64, 3)

        l = FullyConnected('fc0', l, 512)
        l = PReLU('prelu', l)
        logits = FullyConnected('fc-pi', l, NUM_ACTIONS)    # unnormalized policy
        value = FullyConnected('fc-v', l, 1)
        return logits, value
//  [python]
def scanFunc(self, state, b_input):
		ifeatures, outBox, cropBox, angle = b_input
		cropFeatures = tf.image.crop_to_bounding_box(ifeatures, outBox[1], outBox[0], outBox[3], outBox[2])
		rotateCropedFeatures = tf.contrib.image.rotate(cropFeatures, angle)

		textImgFeatures = tf.image.crop_to_bounding_box(rotateCropedFeatures, cropBox[1], cropBox[0], cropBox[3], cropBox[2])

		# resize keep ratio 	
		w = tf.to_int32(tf.ceil(tf.multiply(tf.divide(self.fix_RoiHeight, cropBox[3]), tf.to_float(cropBox[2]))))
		resize_textImgFeatures = tf.image.resize_images(textImgFeatures, (self.fix_RoiHeight, w), 1)
		w = tf.minimum(w, self.max_RoiWidth)
		pad_or_crop_textImgFeatures = tf.image.crop_to_bounding_box(resize_textImgFeatures, 0, 0, self.fix_RoiHeight, w)
		# pad
		pad_or_crop_textImgFeatures = tf.image.pad_to_bounding_box(pad_or_crop_textImgFeatures, 0, 0, self.fix_RoiHeight, self.max_RoiWidth)
		
		return [pad_or_crop_textImgFeatures, w]
//  [python]
def scanFunc(self, state, b_input):
		ifeatures, outBox, cropBox, angle = b_input
		cropFeatures = tf.image.crop_to_bounding_box(ifeatures, outBox[1], outBox[0], outBox[3], outBox[2])
		rotateCropedFeatures = tf.contrib.image.rotate(cropFeatures, angle)

		textImgFeatures = tf.image.crop_to_bounding_box(rotateCropedFeatures, cropBox[1], cropBox[0], cropBox[3], cropBox[2])

		# resize keep ratio 	
		w = tf.cast(tf.ceil(tf.multiply(tf.divide(self.fix_RoiHeight, cropBox[3]), tf.cast(cropBox[2], tf.float64))), tf.int32)
		resize_textImgFeatures = tf.image.resize_images(textImgFeatures, (self.fix_RoiHeight, w), 1)
		w = tf.minimum(w, self.max_RoiWidth)
		pad_or_crop_textImgFeatures = tf.image.crop_to_bounding_box(resize_textImgFeatures, 0, 0, self.fix_RoiHeight, w)
		# pad
		pad_or_crop_textImgFeatures = tf.image.pad_to_bounding_box(pad_or_crop_textImgFeatures, 0, 0, self.fix_RoiHeight, self.max_RoiWidth)
		
		return [pad_or_crop_textImgFeatures, w]
//  [python]
def _dot_op(self, features):
        batch_size = tf.shape(features)[0]
        matrixdot = tf.matmul(features, features, transpose_b=True)
        feature_dim = matrixdot.shape[-1]

        ones_mat = tf.ones_like(matrixdot)
        lower_tri_mat = ones_mat - tf.linalg.band_part(ones_mat, 0, -1)
        lower_tri_mask = tf.to_int32(lower_tri_mat) > 0
        result = tf.boolean_mask(matrixdot, lower_tri_mask)
        output_dim = feature_dim * (feature_dim - 1) // 2

        return tf.reshape(result, (batch_size, output_dim))
//  [python]
def _dot_op(self, features):
        batch_size = tf.shape(features)[0]
        matrixdot = tf.matmul(features, features, transpose_b=True)
        feature_dim = matrixdot.shape[-1]

        ones_mat = tf.ones_like(matrixdot)
        lower_tri_mat = ones_mat - tf.linalg.band_part(ones_mat, 0, -1)
        lower_tri_mask = tf.cast(lower_tri_mat, tf.bool)
        result = tf.boolean_mask(matrixdot, lower_tri_mask)
        output_dim = feature_dim * (feature_dim - 1) // 2

        return tf.reshape(result, (batch_size, output_dim))
//  [python]
def mask_padding_embedding_lookup(embeddings, embedding_dim,
                                  input, padding_id):
    """ mask padding tf.nn.embedding_lookup.
    padding_id must be zero.

    ref(@ay27): https://github.com/tensorflow/tensorflow/issues/2373
    """

    assert padding_id == 0

    mask_padding_zero_op = tf.scatter_update(
        embeddings, padding_id, tf.zeros([embedding_dim], dtype=tf.float32),
        name="mask_padding_zero_op")
    with tf.control_dependencies([mask_padding_zero_op]):
        output = tf.nn.embedding_lookup(
            embeddings, tf.to_int32(input, name="lookup_idx_cast"),
            name="embedding_lookup")
    return output
//  [python]
def mask_padding_embedding_lookup(embeddings, embedding_dim,
                                  input, padding_id):
    """ mask padding tf.nn.embedding_lookup.
    padding_id must be zero.

    ref(@ay27): https://github.com/tensorflow/tensorflow/issues/2373
    """

    assert padding_id == 0

    mask_padding_zero_op = tf.scatter_update(
        embeddings, padding_id, tf.zeros([embedding_dim], dtype=tf.float32),
        name="mask_padding_zero_op")
    with tf.control_dependencies([mask_padding_zero_op]):
        output = tf.nn.embedding_lookup(
            embeddings, tf.cast(input, tf.int32, name="lookup_idx_cast"),
            name="embedding_lookup")
    return output
//  [python]
def build_graph(self, state, action, futurereward, action_prob):
        logits, value = self._get_NN_prediction(state)
        value = tf.squeeze(value, [1], name='pred_value')  # (B,)
        policy = tf.nn.softmax(logits, name='policy')
        if not self.training:
            return
        log_probs = tf.log(policy + 1e-6)

        log_pi_a_given_s = tf.reduce_sum(
            log_probs * tf.one_hot(action, NUM_ACTIONS), 1)
        advantage = tf.subtract(tf.stop_gradient(value), futurereward, name='advantage')

        pi_a_given_s = tf.reduce_sum(policy * tf.one_hot(action, NUM_ACTIONS), 1)  # (B,)
        importance = tf.stop_gradient(tf.clip_by_value(pi_a_given_s / (action_prob + 1e-8), 0, 10))

        policy_loss = tf.reduce_sum(log_pi_a_given_s * advantage * importance, name='policy_loss')
        xentropy_loss = tf.reduce_sum(policy * log_probs, name='xentropy_loss')
        value_loss = tf.nn.l2_loss(value - futurereward, name='value_loss')

        pred_reward = tf.reduce_mean(value, name='predict_reward')
        advantage = tf.sqrt(tf.reduce_mean(tf.square(advantage)), name='rms_advantage')
        entropy_beta = tf.get_variable('entropy_beta', shape=[],
                                       initializer=tf.constant_initializer(0.01), trainable=False)
        cost = tf.add_n([policy_loss, xentropy_loss * entropy_beta, value_loss])
        cost = tf.truediv(cost, tf.to_float(tf.shape(futurereward)[0]), name='cost')
        summary.add_moving_summary(policy_loss, xentropy_loss,
                                   value_loss, pred_reward, advantage,
                                   cost, tf.reduce_mean(importance, name='importance'))
        return cost
//  [python]
def build_graph(self, state, action, futurereward, action_prob):
        logits, value = self._get_NN_prediction(state)
        value = tf.squeeze(value, [1], name='pred_value')  # (B,)
        policy = tf.nn.softmax(logits, name='policy')
        if not self.training:
            return
        log_probs = tf.log(policy + 1e-6)

        log_pi_a_given_s = tf.reduce_sum(
            log_probs * tf.one_hot(action, NUM_ACTIONS), 1)
        advantage = tf.subtract(tf.stop_gradient(value), futurereward, name='advantage')

        pi_a_given_s = tf.reduce_sum(policy * tf.one_hot(action, NUM_ACTIONS), 1)  # (B,)
        importance = tf.stop_gradient(tf.clip_by_value(pi_a_given_s / (action_prob + 1e-8), 0, 10))

        policy_loss = tf.reduce_sum(log_pi_a_given_s * advantage * importance, name='policy_loss')
        xentropy_loss = tf.reduce_sum(policy * log_probs, name='xentropy_loss')
        value_loss = tf.nn.l2_loss(value - futurereward, name='value_loss')

        pred_reward = tf.reduce_mean(value, name='predict_reward')
        advantage = tf.sqrt(tf.reduce_mean(tf.square(advantage)), name='rms_advantage')
        entropy_beta = tf.get_variable('entropy_beta', shape=[],
                                       initializer=tf.constant_initializer(0.01), trainable=False)
        cost = tf.add_n([policy_loss, xentropy_loss * entropy_beta, value_loss])
        cost = tf.truediv(cost, tf.cast(tf.shape(futurereward)[0], tf.float32), name='cost')
        summary.add_moving_summary(policy_loss, xentropy_loss,
                                   value_loss, pred_reward, advantage,
                                   cost, tf.reduce_mean(importance, name='importance'))
        return cost
//  [python]
def average_pooling(embeds, nonzeros):
    with tf.name_scope('average_pooling'):
        embeds_sum = tf.reduce_sum(embeds, 1)
        embeds_mean = embeds_sum / tf.to_float(nonzeros)
    return embeds_mean
//  [python]
def average_pooling(embeds, nonzeros):
    with tf.name_scope('average_pooling'):
        embeds_sum = tf.reduce_sum(embeds, 1)
        embeds_mean = embeds_sum / tf.cast(nonzeros, tf.float32)
    return embeds_mean
//  [python]
def sgd_lr_fasttext_decay(features, params):
    """sgd lr fasttext decay."""

    lr = params['opts'].lr
    total_tokens = params['total_tokens']

    processed_tokens = get_processed_tokens()
    current_tokens = tf.reduce_sum(features[model_keys.TOKENS_COL])
    processed_tokens = tf.assign_add(processed_tokens, current_tokens)
    tf.summary.scalar("processed_tokens", processed_tokens)
    new_lr = lr * (1.0 - (tf.to_float(processed_tokens)
                          / tf.to_float(total_tokens)))
    new_lr = tf.maximum(new_lr, 1e-8)
    return new_lr