jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_metric_ops(
      self, features_dict: types.TensorTypeMaybeDict,
      predictions_dict: types.TensorTypeMaybeDict,
      labels_dict: types.TensorTypeMaybeDict
  ) -> Dict[str, Tuple[types.TensorType, types.TensorType]]:
    # Note that we have to squeeze predictions, labels, weights so they are all
    # N element vectors (otherwise some of them might be N x 1 tensors, and
    # multiplying a N element vector with a N x 1 tensor uses matrix
    # multiplication rather than element-wise multiplication).
    predictions, labels = self._get_labels_and_predictions(
        predictions_dict, labels_dict)
    predictions = _flatten_to_one_dim(t.cast(predictions, t.float64))
    labels = _flatten_to_one_dim(t.cast(labels, t.float64))
    weights = tf.ones_like(predictions)
    if self._example_weight_key:
      weights = _flatten_to_one_dim(
          t.cast(features_dict[self._example_weight_key], t.float64))
    return {
        self._metric_key(metric_keys.SQUARED_PEARSON_CORRELATION):
            metrics.squared_pearson_correlation(predictions, labels, weights)
    }
//  [python]
def get_metric_ops(
      self, features_dict: types.TensorTypeMaybeDict,
      predictions_dict: types.TensorTypeMaybeDict,
      labels_dict: types.TensorTypeMaybeDict
  ) -> Dict[str, Tuple[types.TensorType, types.TensorType]]:
    # Note that we have to squeeze predictions, labels, weights so they are all
    # N element vectors (otherwise some of them might be N x 1 tensors, and
    # multiplying a N element vector with a N x 1 tensor uses matrix
    # multiplication rather than element-wise multiplication).
    predictions, labels = self._get_labels_and_predictions(
        predictions_dict, labels_dict)
    predictions = _flatten_to_one_dim(tf.cast(predictions, tf.float64))
    labels = _flatten_to_one_dim(tf.cast(labels, tf.float64))
    weights = tf.ones_like(predictions)
    if self._example_weight_key:
      weights = _flatten_to_one_dim(
          tf.cast(features_dict[self._example_weight_key], tf.float64))
    return {
        self._metric_key(metric_keys.SQUARED_PEARSON_CORRELATION):
            metrics.squared_pearson_correlation(predictions, labels, weights)
    }
//  [python]
def prepare_inputs(self):
        """ Prepares Input
        """
        #q1_inputs, self.qmax = clip_sentence(self.q1_inputs, self.q1_len)
        #q2_inputs, self.a1max = clip_sentence(self.q2_inputs, self.q2_len)
        #q3_inputs, self.a2max = clip_sentence(self.q3_inputs, self.q3_len)
        q1_inputs = self.q1_inputs
        q2_inputs = self.q2_inputs
        q3_inputs = self.q3_inputs
        self.qmax = self.args.dmax * self.args.smax
        self.a1max = self.args.dmax * self.args.smax
        self.a2max = self.args.dmax * self.args.smax

        self.q1_mask = tf.cast(q1_inputs, tf.bool)
        self.q2_mask = tf.cast(q2_inputs, tf.bool)
        self.q3_mask = tf.cast(q3_inputs, tf.bool)

        with tf.device('/cpu:0'):
            q1_embed = tf.nn.embedding_lookup(self.embeddings,
                                                    q1_inputs)
            q2_embed = tf.nn.embedding_lookup(self.embeddings,
                                                    q2_inputs)
            q3_embed = tf.nn.embedding_lookup(self.embeddings,
                                                    q3_inputs)

        if(self.args.all_dropout):
            # By default, this is disabled
            q1_embed = tf.nn.dropout(q1_embed, self.emb_dropout)
            q2_embed = tf.nn.dropout(q2_embed, self.emb_dropout)
            q3_embed = tf.nn.dropout(q3_embed, self.emb_dropout)

        # Ignore these. :)
        self.c1_cnn, self.c2_cnn, self.c3_cnn = None, None, None
        self.p1_pos, self.p2_pos, self.p3_pos = None, None, None

        if('TNET' in self.args.rnn_type):
            t_inputs, _ = clip_sentence(self.trans_inputs, self.trans_len)
            self.trans_embed = tf.nn.embedding_lookup(self.embeddings,
                                                        t_inputs)

        if('EXP' in self.args.rnn_type):
            #p_user_inputs1, self.pair_user_q1max = clip_sentence(self.pair_user_inputs1, self.pair_user_len1)
            #p_user_inputs2, self.pair_user_q2max = clip_sentence(self.pair_user_inputs2, self.pair_user_len2)
            p_user_inputs2 = self.pair_user_inputs2
            self.pair_user_q2max = self.args.dmax * self.args.smax            
            pair_user_q2_embed = tf.nn.embedding_lookup(self.embeddings,
                                                        p_user_inputs2)
            self.pair_user_q2_embed = tf.nn.dropout(pair_user_q2_embed, self.emb_dropout)
            self.pair_user_q2_mask = tf.cast(p_user_inputs2, tf.bool)

        self.q1_embed = q1_embed
        self.q2_embed = q2_embed
        self.q3_embed = q3_embed
//  [python]
def get_metric_dictionary(self):
    metrics_dict = {}
    ap_list = []
    for c in self.class_range:
      num_detections_c = tf.shape(self.scores[c])[0]
      _, sorted_detection_indices = tf.math.top_k(
          self.scores[c], k=num_detections_c)
      tp_c = tf.to_float(
          tf.gather(self.tp[c], sorted_detection_indices))
      fp_c = 1.0 - tp_c
      num_gt_c = tf.to_float(self.num_gt[c])
      if num_gt_c > 0 and num_detections_c > 0:
        cumsum_tp_c = tf.math.cumsum(tp_c, axis=0)
        cumsum_fp_c = tf.math.cumsum(fp_c, axis=0)
        precision = cumsum_tp_c / (cumsum_tp_c + cumsum_fp_c)
        recall = cumsum_tp_c / num_gt_c
        recall_mask = tf.math.logical_and(
            tf.greater_equal(recall, 0.0), tf.less_equal(recall, 1.0))
        precision = tf.boolean_mask(precision, recall_mask)
        recall = tf.boolean_mask(recall, recall_mask)
        recall = tf.concat([
            tf.constant([0], dtype=tf.float32),
            recall,
            tf.constant([1], dtype=tf.float32)
        ],
                           axis=0)
        precision = tf.concat([
            tf.constant([0], dtype=tf.float32),
            precision,
            tf.constant([0], dtype=tf.float32)
        ],
                              axis=0)
        num_values = tf.shape(precision)[0]
        for i in tf.range(num_values - 2, -1, -1):
          precision = tf.maximum(
              precision,
              tf.one_hot(i, depth=num_values, dtype=tf.float32) *
              precision[i + 1])
        ap_c = tf.reduce_sum((recall[1:] - recall[:-1]) * precision[:-1])
      else:
        ap_c = tf.constant(0.0, dtype=tf.float32)
      class_name = _get_class_name(class_id=c, label_map=self.label_map)
      metrics_dict[
          self.eval_prefix +
          '_IOU{}_AP/{}'.format(self.iou_threshold, class_name)] = ap_c
      ap_list.append(ap_c)
    mean_ap = _non_nan_mean(ap_list)
    metrics_dict[
        self.eval_prefix +
        '_avg/mean_AP_IOU{}'.format(self.iou_threshold)] = mean_ap
    return metrics_dict
//  [python]
def get_metric_dictionary(self):
    metrics_dict = {}
    ap_list = []
    for c in self.class_range:
      num_detections_c = tf.shape(self.scores[c])[0]
      _, sorted_detection_indices = tf.math.top_k(
          self.scores[c], k=num_detections_c)
      tp_c = tf.cast(
          tf.gather(self.tp[c], sorted_detection_indices), dtype=tf.float32)
      fp_c = 1.0 - tp_c
      num_gt_c = tf.cast(self.num_gt[c], dtype=tf.float32)
      if num_gt_c > 0 and num_detections_c > 0:
        cumsum_tp_c = tf.math.cumsum(tp_c, axis=0)
        cumsum_fp_c = tf.math.cumsum(fp_c, axis=0)
        precision = cumsum_tp_c / (cumsum_tp_c + cumsum_fp_c)
        recall = cumsum_tp_c / num_gt_c
        recall_mask = tf.math.logical_and(
            tf.greater_equal(recall, 0.0), tf.less_equal(recall, 1.0))
        precision = tf.boolean_mask(precision, recall_mask)
        recall = tf.boolean_mask(recall, recall_mask)
        recall = tf.concat([
            tf.constant([0], dtype=tf.float32),
            recall,
            tf.constant([1], dtype=tf.float32)
        ],
                           axis=0)
        precision = tf.concat([
            tf.constant([0], dtype=tf.float32),
            precision,
            tf.constant([0], dtype=tf.float32)
        ],
                              axis=0)
        num_values = tf.shape(precision)[0]
        for i in tf.range(num_values - 2, -1, -1):
          precision = tf.maximum(
              precision,
              tf.one_hot(i, depth=num_values, dtype=tf.float32) *
              precision[i + 1])
        ap_c = tf.reduce_sum((recall[1:] - recall[:-1]) * precision[:-1])
      else:
        ap_c = tf.constant(0.0, dtype=tf.float32)
      class_name = _get_class_name(class_id=c, label_map=self.label_map)
      metrics_dict[
          self.eval_prefix +
          '_IOU{}_AP/{}'.format(self.iou_threshold, class_name)] = ap_c
      ap_list.append(ap_c)
    mean_ap = _non_nan_mean(ap_list)
    metrics_dict[
        self.eval_prefix +
        '_avg/mean_AP_IOU{}'.format(self.iou_threshold)] = mean_ap
    return metrics_dict
//  [python]
def _class_ids(probabilities: tf.Tensor) -> tf.Tensor:
  """Returns class_ids associated with given probabilities tensor.

  Args:
    probabilities: Batch of probablitities (i.e. [[class0, class1, ...], ...]
      with shape [batch, n_classes]).

  Returns:
    Class IDs for N classes (i.e. [[0, 1, ..., n-1], ...] with shape
    [batch, n_classes]).
  """
  n_classes = tf.shape(input=probabilities)[-1]
  # Tensor representing shape of class_ids expanded by batch dims: [1,n_classes]
  expanded_dims = tf.concat(
      [tf.ones_like(tf.shape(input=probabilities))[:-1], [n_classes]], axis=0)
  # Tensor for multiplying tiles by batch size. Shape should be [batch_size,1]
  batch_multiplier = tf.concat([tf.shape(input=probabilities)[:-1], [1]],
                               axis=0)
  # Batches of [0, ..., n_classes]
  return tf.tile(
      input=tf.reshape(tensor=tf.range(n_classes), shape=expanded_dims),
      multiples=batch_multiplier)
//  [python]
def _class_ids(probabilities: tf.Tensor) -> tf.Tensor:
  """Returns class_ids associated with given probabilities tensor.

  Args:
    probabilities: Batch of probablitities (i.e. [[class0, class1, ...], ...]
      with shape [batch, n_classes]).

  Returns:
    Class IDs for N classes (i.e. [[0, 1, ..., n-1], ...] with shape
    [batch, n_classes]).
  """
  n_classes = tf.cast(tf.shape(input=probabilities)[-1], tf.int64)
  # Tensor representing shape of class_ids expanded by batch dims: [1,n_classes]
  expanded_dims = tf.concat(
      [tf.ones_like(tf.shape(input=probabilities))[:-1], [n_classes]], axis=0)
  # Tensor for multiplying tiles by batch size. Shape should be [batch_size,1]
  batch_multiplier = tf.concat([tf.shape(input=probabilities)[:-1], [1]],
                               axis=0)
  # Batches of [0, ..., n_classes]
  return tf.tile(
      input=tf.reshape(tensor=tf.range(n_classes), shape=expanded_dims),
      multiples=batch_multiplier)
//  [python]
def confusion_matrix_metric_ops(
      self,
      features_dict: types.TensorTypeMaybeDict,
      predictions_dict: types.TensorTypeMaybeDict,
      labels_dict: types.TensorTypeMaybeDict,
  ) -> Tuple[Dict[str, List[types.TensorType]], Dict[str,
                                                     List[types.TensorType]]]:
    """Metric ops for computing confusion matrix at the given thresholds.

    This is factored out because it's common to AucPlots and
    ConfusionMatrixAtThresholds.

    Args:
      features_dict: Features dict.
      predictions_dict: Predictions dict.
      labels_dict: Labels dict.

    Returns:
      (value_ops, update_ops) for the confusion matrix.
    """
    # Note that we have to squeeze predictions, labels, weights so they are all
    # N element vectors (otherwise some of them might be N x 1 tensors, and
    # multiplying a N element vector with a N x 1 tensor uses matrix
    # multiplication rather than element-wise multiplication).
    predictions, labels = self._get_labels_and_predictions(
        predictions_dict, labels_dict)
    prediction_tensor = _flatten_to_one_dim(t.cast(predictions, t.float64))
    label_tensor = _flatten_to_one_dim(t.cast(labels, t.float64))
    squeezed_weights = t.ones_like(prediction_tensor)
    if self._example_weight_key:
      squeezed_weights = _flatten_to_one_dim(
          t.cast(features_dict[self._example_weight_key], t.float64))
    prediction_tensor, label_tensor, squeezed_weights = (
        _create_predictions_labels_weights_for_fractional_labels(
            prediction_tensor, label_tensor, squeezed_weights))

    # TODO(b/72239826): Expose _confusion_matrix_at_thresholds for OSS?
    values, update_ops = metrics_impl._confusion_matrix_at_thresholds(  # pylint: disable=protected-access
        label_tensor, prediction_tensor, self._thresholds, squeezed_weights)

    values['precision'] = math.divide_no_nan(values['tp'],
                                             (values['tp'] + values['fp']))
    values['recall'] = math.divide_no_nan(values['tp'],
                                          (values['tp'] + values['fn']))
    return (values, update_ops)
//  [python]
def confusion_matrix_metric_ops(
      self,
      features_dict: types.TensorTypeMaybeDict,
      predictions_dict: types.TensorTypeMaybeDict,
      labels_dict: types.TensorTypeMaybeDict,
  ) -> Tuple[Dict[str, List[types.TensorType]], Dict[str,
                                                     List[types.TensorType]]]:
    """Metric ops for computing confusion matrix at the given thresholds.

    This is factored out because it's common to AucPlots and
    ConfusionMatrixAtThresholds.

    Args:
      features_dict: Features dict.
      predictions_dict: Predictions dict.
      labels_dict: Labels dict.

    Returns:
      (value_ops, update_ops) for the confusion matrix.
    """
    # Note that we have to squeeze predictions, labels, weights so they are all
    # N element vectors (otherwise some of them might be N x 1 tensors, and
    # multiplying a N element vector with a N x 1 tensor uses matrix
    # multiplication rather than element-wise multiplication).
    predictions, labels = self._get_labels_and_predictions(
        predictions_dict, labels_dict)
    prediction_tensor = _flatten_to_one_dim(tf.cast(predictions, tf.float64))
    label_tensor = _flatten_to_one_dim(tf.cast(labels, tf.float64))
    squeezed_weights = tf.ones_like(prediction_tensor)
    if self._example_weight_key:
      squeezed_weights = _flatten_to_one_dim(
          tf.cast(features_dict[self._example_weight_key], tf.float64))
    prediction_tensor, label_tensor, squeezed_weights = (
        _create_predictions_labels_weights_for_fractional_labels(
            prediction_tensor, label_tensor, squeezed_weights))

    # TODO(b/72239826): Expose _confusion_matrix_at_thresholds for OSS?
    values, update_ops = metrics_impl._confusion_matrix_at_thresholds(  # pylint: disable=protected-access
        label_tensor, prediction_tensor, self._thresholds, squeezed_weights)

    values['precision'] = math.divide_no_nan(values['tp'],
                                             (values['tp'] + values['fp']))
    values['recall'] = math.divide_no_nan(values['tp'],
                                          (values['tp'] + values['fn']))
    return (values, update_ops)
//  [python]
def get_metric_ops(  # pytype: disable=signature-mismatch  # overriding-return-type-checks
      self, features_dict: types.TensorTypeMaybeDict,
      predictions_dict: types.TensorTypeMaybeDict,
      labels_dict: types.TensorTypeMaybeDict) -> Dict[str, Tuple[
          types.TensorOrOperationType, types.TensorOrOperationType]]:
    predictions, labels = self._get_labels_and_predictions(
        predictions_dict, labels_dict)
    prediction_tensor = _flatten_to_one_dim(t.cast(predictions, t.float64))
    label_tensor = _flatten_to_one_dim(t.cast(labels, t.float64))
    squeezed_weights = tf.ones_like(prediction_tensor)
    if self._example_weight_key:
      squeezed_weights = _flatten_to_one_dim(
          t.cast(features_dict[self._example_weight_key], t.float64))
    metric_fn = self._metric_fn(label_tensor, prediction_tensor,
                                squeezed_weights)

    return {self._metric_key(self._metric_name): metric_fn}
//  [python]
def get_metric_ops(  # pytype: disable=signature-mismatch  # overriding-return-type-checks
      self, features_dict: types.TensorTypeMaybeDict,
      predictions_dict: types.TensorTypeMaybeDict,
      labels_dict: types.TensorTypeMaybeDict) -> Dict[str, Tuple[
          types.TensorOrOperationType, types.TensorOrOperationType]]:
    predictions, labels = self._get_labels_and_predictions(
        predictions_dict, labels_dict)
    prediction_tensor = _flatten_to_one_dim(tf.cast(predictions, tf.float64))
    label_tensor = _flatten_to_one_dim(tf.cast(labels, tf.float64))
    squeezed_weights = tf.ones_like(prediction_tensor)
    if self._example_weight_key:
      squeezed_weights = _flatten_to_one_dim(
          tf.cast(features_dict[self._example_weight_key], tf.float64))
    metric_fn = self._metric_fn(label_tensor, prediction_tensor,
                                squeezed_weights)

    return {self._metric_key(self._metric_name): metric_fn}
//  [python]
def frame_swap(
    frames: tf.Tensor, frame_mask: Optional[tf.Tensor] = None
) -> tf.Tensor:
  """Self-weighted average pooling over all frames of a video.

  It does the following operation independently for each feature:
    x_pooled = (sum_i x_i * |x_i|) / (sum_i |x_i|).
  Basically the weight for the feature in each frame is determined by the
  magnitude of the feature itself.

  Paper: https://research.google/pubs/pub48351/

  Args:
    frames: A tensor with shape [batch_size, max_frames, feature_size].
    frame_mask:  A tensor with shape [batch_size, max_frames, 1].

  Returns:
    A tensor with shape [batch_size, feature_size].
  """
  weights = tf.abs(frames)
  if frame_mask is not None:
    weights *= t.cast(frame_mask, weights.dtype)
  # We set axis to 1 to reduce the dimension corresponding to max_frames.
  return weighted_average_pooling(frames, weights, axis=1)
//  [python]
def frame_swap(
    frames: tf.Tensor, frame_mask: Optional[tf.Tensor] = None
) -> tf.Tensor:
  """Self-weighted average pooling over all frames of a video.

  It does the following operation independently for each feature:
    x_pooled = (sum_i x_i * |x_i|) / (sum_i |x_i|).
  Basically the weight for the feature in each frame is determined by the
  magnitude of the feature itself.

  Paper: https://research.google/pubs/pub48351/

  Args:
    frames: A tensor with shape [batch_size, max_frames, feature_size].
    frame_mask:  A tensor with shape [batch_size, max_frames, 1].

  Returns:
    A tensor with shape [batch_size, feature_size].
  """
  weights = tf.abs(frames)
  if frame_mask is not None:
    weights *= tf.cast(frame_mask, weights.dtype)
  # We set axis to 1 to reduce the dimension corresponding to max_frames.
  return weighted_average_pooling(frames, weights, axis=1)
//  [python]
def test_table_roundtrip(self):
    export_path = os.path.join(tempfile.mkdtemp(), 'export')

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        input_string = tf.compat.v1.placeholder(tf.string)
        # Map string through a table, in this case based on a constant tensor.
        table_keys = ['cat', 'dog', 'giraffe']
        initializer = tf.lookup.KeyValueTensorInitializer(
            keys=table_keys,
            values=tf.range(len(table_keys)),
            key_dtype=tf.string,
            value_dtype=tf.int64)
        table = tf.lookup.StaticHashTable(initializer, default_value=-1)

        output = table.lookup(input_string)
        inputs = {'input': input_string}
        outputs = {'output': output}
        saved_transform_io.write_saved_transform_from_session(
            session, inputs, outputs, export_path)

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        # Using a computed input gives confidence that the graphs are fused.
        input_string = tf.constant('dog')
        inputs = {'input': input_string}
        _, outputs = (
            saved_transform_io.partially_apply_saved_transform_internal(
                export_path, inputs))
        session.run(tf.compat.v1.tables_initializer())
        result = session.run(outputs['output'])
        self.assertEqual(1, result)
//  [python]
def test_table_roundtrip(self):
    export_path = os.path.join(tempfile.mkdtemp(), 'export')

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        input_string = tf.compat.v1.placeholder(tf.string)
        # Map string through a table, in this case based on a constant tensor.
        table_keys = ['cat', 'dog', 'giraffe']
        initializer = tf.lookup.KeyValueTensorInitializer(
            keys=table_keys,
            values=tf.cast(tf.range(len(table_keys)), tf.int64),
            key_dtype=tf.string,
            value_dtype=tf.int64)
        table = tf.lookup.StaticHashTable(initializer, default_value=-1)

        output = table.lookup(input_string)
        inputs = {'input': input_string}
        outputs = {'output': output}
        saved_transform_io.write_saved_transform_from_session(
            session, inputs, outputs, export_path)

    with tf.compat.v1.Graph().as_default():
      with tf.compat.v1.Session().as_default() as session:
        # Using a computed input gives confidence that the graphs are fused.
        input_string = tf.constant('dog')
        inputs = {'input': input_string}
        _, outputs = (
            saved_transform_io.partially_apply_saved_transform_internal(
                export_path, inputs))
        session.run(tf.compat.v1.tables_initializer())
        result = session.run(outputs['output'])
        self.assertEqual(1, result)
//  [python]
def distorted_inputs(batch_size, Delta2, epsilon2, W_conv1Noise):
  """Construct distorted input for CIFAR training using the Reader ops.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.

  Raises:
    ValueError: If no data_dir
  """
  if not FLAGS.data_dir:
    raise ValueError('Please supply a data_dir')
  data_dir = os.path.join(FLAGS.data_dir, 'cifar-10-batches-bin')
  beta = redistributeNoise(os.getcwd() + '/LRP_0_25_v12.txt')
  images, labels = cifar10_input.distorted_inputs(beta = beta, data_dir=data_dir,
                            batch_size=batch_size, Delta2 = Delta2, epsilon2 = epsilon2, W_conv1Noise = W_conv1Noise)
    
  if FLAGS.use_fp16:
    images = tf.to_float(images)
    labels = tf.to_float(labels)
  return images, labels
//  [python]
def distorted_inputs(batch_size, Delta2, epsilon2, W_conv1Noise):
  """Construct distorted input for CIFAR training using the Reader ops.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.

  Raises:
    ValueError: If no data_dir
  """
  if not FLAGS.data_dir:
    raise ValueError('Please supply a data_dir')
  data_dir = os.path.join(FLAGS.data_dir, 'cifar-10-batches-bin')
  beta = redistributeNoise(os.getcwd() + '/LRP_0_25_v12.txt')
  images, labels = cifar10_input.distorted_inputs(beta = beta, data_dir=data_dir,
                            batch_size=batch_size, Delta2 = Delta2, epsilon2 = epsilon2, W_conv1Noise = W_conv1Noise)
    
  if FLAGS.use_fp16:
    images = tf.cast(images, tf.float16)
    labels = tf.cast(labels, tf.float16)
  return images, labels
//  [python]
def loss(logits, labels):
  """You can also add L2Loss to all the trainable variables.

  Add summary for "Loss" and "Loss/avg".
  Args:
    logits: Logits from inference().
    labels: Labels from distorted_inputs or inputs(). 1-D tensor
            of shape [batch_size]

  Returns:
    Loss tensor of type float.
  """
  # Calculate the average cross entropy loss across the batch.
  labels = tf.to_int64(labels)
  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=labels, logits=logits, name='cross_entropy_per_example')
  cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')
  tf.add_to_collection('losses', cross_entropy_mean)

  # The total loss is defined as the cross entropy loss plus all of the weight
  # decay terms (L2 loss).
  return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def loss(logits, labels):
  """You can also add L2Loss to all the trainable variables.

  Add summary for "Loss" and "Loss/avg".
  Args:
    logits: Logits from inference().
    labels: Labels from distorted_inputs or inputs(). 1-D tensor
            of shape [batch_size]

  Returns:
    Loss tensor of type float.
  """
  # Calculate the average cross entropy loss across the batch.
  labels = tf.cast(labels, tf.int64)
  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=labels, logits=logits, name='cross_entropy_per_example')
  cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')
  tf.add_to_collection('losses', cross_entropy_mean)

  # The total loss is defined as the cross entropy loss plus all of the weight
  # decay terms (L2 loss).
  return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def inputs(eval_data):
  """Construct input for CIFAR evaluation using the Reader ops.

  Args:
    eval_data: bool, indicating if one should use the train or eval data set.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.

  Raises:
    ValueError: If no data_dir
  """
  if not FLAGS.data_dir:
    raise ValueError('Please supply a data_dir')
  data_dir = os.path.join(FLAGS.data_dir, 'cifar-10-batches-bin')
  images, labels = cifar10_input.inputs(eval_data=eval_data,
                                        data_dir=data_dir,
                                        batch_size=2500)
  if FLAGS.use_fp16:
    images = tf.to_float(images)
    labels = tf.to_float(labels)
  return images, labels
//  [python]
def inputs(eval_data):
  """Construct input for CIFAR evaluation using the Reader ops.

  Args:
    eval_data: bool, indicating if one should use the train or eval data set.

  Returns:
    images: Images. 4D tensor of [batch_size, IMAGE_SIZE, IMAGE_SIZE, 3] size.
    labels: Labels. 1D tensor of [batch_size] size.

  Raises:
    ValueError: If no data_dir
  """
  if not FLAGS.data_dir:
    raise ValueError('Please supply a data_dir')
  data_dir = os.path.join(FLAGS.data_dir, 'cifar-10-batches-bin')
  images, labels = cifar10_input.inputs(eval_data=eval_data,
                                        data_dir=data_dir,
                                        batch_size=2500)
  if FLAGS.use_fp16:
    images = tf.cast(images, tf.float16)
    labels = tf.cast(labels, tf.float16)
  return images, labels
//  [python]
def lossDPSGD(logits, labels):
    """Add L2Loss to all the trainable variables.
        
    Add summary for "Loss" and "Loss/avg".
    Args:
    logits: Logits from inference().
    labels: Labels from distorted_inputs or inputs(). 1-D tensor
    of shape [batch_size]
    
    Returns:
    Loss tensor of type float.
    """
    # Calculate the average cross entropy loss across the batch.
    labels = tf.argmax(tf.cast(labels, tf.int64), 1)
    print(labels)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits, name='cross_entropy_per_example')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')
    tf.add_to_collection('losses', cross_entropy_mean)
        
    # The total loss is defined as the cross entropy loss plus all of the weight
    # decay terms (L2 loss).
    return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def _convert_label(self, image, label):
        r""" Maps chars in label to integers  according to table.txt

        Args:
            image ( :class:`tensorflow.tensor`) : Image tensor
            label  (str)      : label

        Returns:
            :class:`tensorflow.Tensor` : Image tensor. 
            :class:`tensorflow.Tensor` : Label sparse tensor.
        """
        chars = tf.strings.unicode_split(label, input_encoding="UTF-8")
        mapped_label = tf.ragged.map_flat_values(self.table.lookup, chars)
        sparse_label = mapped_label.to_sparse()
        label = tf.to_int32(sparse_label)
        return image, label
//  [python]
def _convert_label(self, image, label):
        r""" Maps chars in label to integers  according to table.txt

        Args:
            image ( :class:`tensorflow.tensor`) : Image tensor
            label  (str)      : label

        Returns:
            :class:`tensorflow.Tensor` : Image tensor. 
            :class:`tensorflow.Tensor` : Label sparse tensor.
        """
        chars = tf.strings.unicode_split(label, input_encoding="UTF-8")
        mapped_label = tf.ragged.map_flat_values(self.table.lookup, chars)
        sparse_label = mapped_label.to_sparse()
        label = tf.cast(sparse_label, tf.int32)
        return image, label
//  [python]
def frame_pooling(frames, method="average", num_frames=None):
  """Pools over the frames of a video.

  Args:
    frames: tensor of shape [batch_size, num_frames, feature_size].
    method: string indicating pooling method, one of: "average", "max",
      "attention", or "none".
    num_frames: optional tensor of shape [batch_size] indicating valid number of
      frames for each video.

  Returns:
    tensor of shape [batch_size, feature_size] for average, max, or
    attention pooling, and shape [batch_size*num_frames, feature_size]
    for none pooling.
  Raises:
    ValueError: if method is other than "average", "max", "attention", or
    "none".
  """
  frame_mask = None
  if num_frames is not None:
    max_frames = frames.shape.as_list()[1]
    # Generate binary mask from number of frames.
    frame_mask = tf.sequence_mask(num_frames, max_frames, frames.dtype)
    frame_mask = tf.expand_dims(frame_mask, axis=2)

  if method == "average":
    if num_frames is None:
      reduced = tf.reduce_mean(frames, 1)
    else:
      num_frames = tf.reshape(tf.to_float(num_frames), [-1, 1])
      reduced = tf.reduce_sum(frames * frame_mask, 1) / num_frames
  elif method == "max":
    if num_frames is not None:
      frame_mask = tf.cast(frame_mask, tf.bool)
      frames = tf.where(
          frame_mask,
          frames,
          tf.ones_like(frames, dtype=frames.dtype)
          * _large_compatible_negative(frames.dtype),
      )
    # Magic to avoid loss NaN when bfloat16 is enabled.
    # See yaqs/5377152819545505792 and b/214396297 for more discussion.
    reduced = tf.reduce_max(frames, 1) + tf.reduce_mean(frames, 1) * 0
  elif method == "swap":
    # Note we assume the frames are in the shape of
    # [batch_size, num_frames, feature_size]. Otherwise this function might
    # fail.
    reduced = frame_swap(frames, frame_mask)
  elif method == "none":
    feature_size = frames.shape.as_list()[2]
    reduced = tf.reshape(frames, [-1, feature_size])
  else:
    raise ValueError("Unrecognized pooling method: %s" % method)

  return reduced
//  [python]
def frame_pooling(frames, method="average", num_frames=None):
  """Pools over the frames of a video.

  Args:
    frames: tensor of shape [batch_size, num_frames, feature_size].
    method: string indicating pooling method, one of: "average", "max",
      "attention", or "none".
    num_frames: optional tensor of shape [batch_size] indicating valid number of
      frames for each video.

  Returns:
    tensor of shape [batch_size, feature_size] for average, max, or
    attention pooling, and shape [batch_size*num_frames, feature_size]
    for none pooling.
  Raises:
    ValueError: if method is other than "average", "max", "attention", or
    "none".
  """
  frame_mask = None
  if num_frames is not None:
    max_frames = frames.shape.as_list()[1]
    # Generate binary mask from number of frames.
    frame_mask = tf.sequence_mask(num_frames, max_frames, frames.dtype)
    frame_mask = tf.expand_dims(frame_mask, axis=2)

  if method == "average":
    if num_frames is None:
      reduced = tf.reduce_mean(frames, 1)
    else:
      num_frames = tf.reshape(tf.cast(num_frames, frames.dtype), [-1, 1])
      reduced = tf.reduce_sum(frames * frame_mask, 1) / num_frames
  elif method == "max":
    if num_frames is not None:
      frame_mask = tf.cast(frame_mask, tf.bool)
      frames = tf.where(
          frame_mask,
          frames,
          tf.ones_like(frames, dtype=frames.dtype)
          * _large_compatible_negative(frames.dtype),
      )
    # Magic to avoid loss NaN when bfloat16 is enabled.
    # See yaqs/5377152819545505792 and b/214396297 for more discussion.
    reduced = tf.reduce_max(frames, 1) + tf.reduce_mean(frames, 1) * 0
  elif method == "swap":
    # Note we assume the frames are in the shape of
    # [batch_size, num_frames, feature_size]. Otherwise this function might
    # fail.
    reduced = frame_swap(frames, frame_mask)
  elif method == "none":
    feature_size = frames.shape.as_list()[2]
    reduced = tf.reshape(frames, [-1, feature_size])
  else:
    raise ValueError("Unrecognized pooling method: %s" % method)

  return reduced
//  [python]
def TaylorExpAdLM(logits, labels, perturbW):
    """You can also add L2Loss to all the trainable variables.
        Add summary for "Loss" and "Loss/avg".
        Args:
        logits: Logits from inference().
        labels: Labels from distorted_inputs or inputs(). 1-D tensor
        of shape [batch_size]
        
        Returns:
        Loss tensor of type float.
        """
    # Calculate the average cross entropy loss across the batch.
    labels = tf.cast(labels, tf.float32)
    
    ### Taylor for benign x
    zeros2 = array_ops.zeros_like(logits, dtype=logits.dtype)
    cond2 = (logits >= zeros2)
    relu_logits_benign = array_ops.where(cond2, logits, zeros2)
    neg_abs_logits_benign = array_ops.where(cond2, -logits, logits)
    Taylor_benign = math_ops.add(relu_logits_benign - logits * labels, math.log(2.0) + 0.5*neg_abs_logits_benign + 1.0/8.0*neg_abs_logits_benign**2)
    
    zeros1 = array_ops.zeros_like(perturbW, dtype=perturbW.dtype)
    cond1 = (perturbW >= zeros1)
    perturbW = array_ops.where(cond1, perturbW, -perturbW)
    
    cross_entropy_mean = tf.reduce_mean(Taylor_benign, name='cross_entropy') + tf.reduce_mean(perturbW, name = 'perturbW');
    
    tf.add_to_collection('losses', cross_entropy_mean)
    
    return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def _decode_and_resize(self, image_filename, label_filename):
        r""" Reads images. Reads and one hot encodes labels.

        Args:
            image_filename (str) : Path of image file.
            label_filename (str) : Path of label file.

        Returns:
            :class:`tensorflow.Tensor`  : Image tensor.
            :class:`tensorflow.Tensor`  : Label tensor.
        """
        image = tf.io.read_file(image_filename)
        image = tf.io.decode_png(image, channels=3)
        image = tf.image.convert_image_dtype(image, tf.float32)
        resize = (tf.shape(image[:, :, 0])//32)*32
        image = tf.image.resize(image, resize)
        label = tf.io.read_file(label_filename)
        label = tf.io.decode_png(label, channels=3)
        label = tf.image.resize(label, resize)
        label = tf.to_int32((label[:, :, 0] > 100))
        label = tf.one_hot(label, self.n_classes)
        return image, label
//  [python]
def _decode_and_resize(self, image_filename, label_filename):
        r""" Reads images. Reads and one hot encodes labels.

        Args:
            image_filename (str) : Path of image file.
            label_filename (str) : Path of label file.

        Returns:
            :class:`tensorflow.Tensor`  : Image tensor.
            :class:`tensorflow.Tensor`  : Label tensor.
        """
        image = tf.io.read_file(image_filename)
        image = tf.io.decode_png(image, channels=3)
        image = tf.image.convert_image_dtype(image, tf.float32)
        resize = (tf.shape(image[:, :, 0])//32)*32
        image = tf.image.resize(image, resize)
        label = tf.io.read_file(label_filename)
        label = tf.io.decode_png(label, channels=3)
        label = tf.image.resize(label, resize)
        label = tf.cast((label[:, :, 0] > 100), tf.int32)
        label = tf.one_hot(label, self.n_classes)
        return image, label
//  [python]
def TaylorExp(logits, labels, adv_logits, b_labels, L, alpha, perturbW):
    """You can also add L2Loss to all the trainable variables.
        Add summary for "Loss" and "Loss/avg".
        Args:
        logits: Logits from inference().
        labels: Labels from distorted_inputs or inputs(). 1-D tensor
        of shape [batch_size]
    
        Returns:
        Loss tensor of type float.
        """
    # Calculate the average cross entropy loss across the batch.
    labels = tf.cast(labels, tf.float32)
    b_labels = tf.cast(b_labels, tf.float32)
    
    # Differentially private sparse cross entropy error based on Taylor Expansion
    zeros = array_ops.zeros_like(adv_logits, dtype=adv_logits.dtype)
    cond = (adv_logits >= zeros)
    relu_logits = array_ops.where(cond, adv_logits, zeros)
    neg_abs_logits = array_ops.where(cond, -adv_logits, adv_logits)
    Taylor_adv = math_ops.add(relu_logits - adv_logits * b_labels, math.log(2.0) + 0.5*neg_abs_logits + 1.0/8.0*neg_abs_logits**2)
    
    ### Taylor for benign x
    zeros2 = array_ops.zeros_like(logits, dtype=logits.dtype)
    cond2 = (logits >= zeros2)
    relu_logits_benign = array_ops.where(cond2, logits, zeros2)
    neg_abs_logits_benign = array_ops.where(cond2, -logits, logits)
    Taylor_benign = math_ops.add(relu_logits_benign - logits * labels, math.log(2.0) + 0.5*neg_abs_logits_benign + 1.0/8.0*neg_abs_logits_benign**2)
    
    zeros1 = array_ops.zeros_like(perturbW, dtype=perturbW.dtype)
    cond1 = (perturbW >= zeros1)
    perturbW = array_ops.where(cond1, perturbW, -perturbW)
    
    ### Adversarial training loss
    adv_loss = (1/(1 + alpha))*(Taylor_benign + alpha * Taylor_adv)

    cross_entropy_mean = tf.reduce_mean(adv_loss, name='cross_entropy') + tf.reduce_mean(perturbW, name = 'perturbW');
    
    tf.add_to_collection('losses', cross_entropy_mean)
    
    return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def _remove_low_score_instances(scores,
                                instance_embeddings,
                                score_threshold,
                                min_num_instances=100):
  """Removes the instances that have an score lower than threshold.

  Args:
    scores: A tf.float32 tensor of size [num_voxels, num_classes].
    instance_embeddings: A tf.float32 tensor of
      size [num_voxels, embedding_dims].
    score_threshold: A float corresponding to score threshold.
    min_num_instances: Minimum number of instances.

  Returns:
    High scores and their corresponding instance embeddings.
  """
  max_score = tf.reduce_max(scores, axis=1)
  score_mask = tf.greater_equal(max_score, score_threshold)
  high_score_indices = tf.to_int32(
      tf.reshape(tf.where(score_mask), [-1]))
  top_k_indices = _get_top_k_indices(scores=scores, k=min_num_instances)
  high_score_indices = tf.cond(
      tf.shape(high_score_indices)[0] > min_num_instances,
      lambda: high_score_indices, lambda: top_k_indices)
  scores = tf.gather(scores, high_score_indices)
  instance_embeddings = tf.gather(instance_embeddings, high_score_indices)
  return scores, instance_embeddings
//  [python]
def _remove_low_score_instances(scores,
                                instance_embeddings,
                                score_threshold,
                                min_num_instances=100):
  """Removes the instances that have an score lower than threshold.

  Args:
    scores: A tf.float32 tensor of size [num_voxels, num_classes].
    instance_embeddings: A tf.float32 tensor of
      size [num_voxels, embedding_dims].
    score_threshold: A float corresponding to score threshold.
    min_num_instances: Minimum number of instances.

  Returns:
    High scores and their corresponding instance embeddings.
  """
  max_score = tf.reduce_max(scores, axis=1)
  score_mask = tf.greater_equal(max_score, score_threshold)
  high_score_indices = tf.cast(
      tf.reshape(tf.where(score_mask), [-1]), dtype=tf.int32)
  top_k_indices = _get_top_k_indices(scores=scores, k=min_num_instances)
  high_score_indices = tf.cond(
      tf.shape(high_score_indices)[0] > min_num_instances,
      lambda: high_score_indices, lambda: top_k_indices)
  scores = tf.gather(scores, high_score_indices)
  instance_embeddings = tf.gather(instance_embeddings, high_score_indices)
  return scores, instance_embeddings
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)
    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    pos_encoding = angle_rads[np.newaxis, ...]
    return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)
    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    pos_encoding = angle_rads[np.newaxis, ...]
    return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def __init__(self, d_model, warmup_steps=4000):
        super(CustomSchedule, self).__init__()

        self.d_model = d_model
        self.d_model = tf.cast(self.d_model, tf.float32)

        self.warmup_steps = warmup_steps
//  [python]
def accuracy_function(real, pred):
    accuracies = tf.equal(real, tf.argmax(pred, axis=2))

    mask = tf.math.logical_not(tf.math.equal(real, 0))
    accuracies = tf.math.logical_and(mask, accuracies)

    accuracies = tf.to_float(accuracies)
    mask = tf.to_float(mask)
    return tf.reduce_sum(accuracies)/tf.reduce_sum(mask)
//  [python]
def accuracy_function(real, pred):
    accuracies = tf.equal(real, tf.argmax(pred, axis=2))

    mask = tf.math.logical_not(tf.math.equal(real, 0))
    accuracies = tf.math.logical_and(mask, accuracies)

    accuracies = tf.cast(accuracies, dtype=tf.float32)
    mask = tf.cast(mask, dtype=tf.float32)
    return tf.reduce_sum(accuracies)/tf.reduce_sum(mask)
//  [python]
def call(self, x, training, mask):
        outputs = {}

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        outputs['encoder_layer_embedding'] = x.numpy()

        x *= tf.sqrt(tf.to_float(self.d_model))
        outputs['encoder_layer_multiply_sqrt'] = x.numpy()

        x += self.pos_encoding[:, :seq_len, :]
        x = self.dropout(x, training=training)
        outputs['encoder_layer_pos_encoding'] = x.numpy()

        for i in range(self.num_layers):
            x, enc_out = self.enc_layers[i](x, training, mask, i)
            outputs.update(enc_out)

        return x, outputs
//  [python]
def call(self, x, training, mask):
        outputs = {}

        seq_len = tf.shape(x)[1]

        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        outputs['encoder_layer_embedding'] = x.numpy()

        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        outputs['encoder_layer_multiply_sqrt'] = x.numpy()

        x += self.pos_encoding[:, :seq_len, :]
        x = self.dropout(x, training=training)
        outputs['encoder_layer_pos_encoding'] = x.numpy()

        for i in range(self.num_layers):
            x, enc_out = self.enc_layers[i](x, training, mask, i)
            outputs.update(enc_out)

        return x, outputs
//  [python]
def call(self, x, enc_output, training,
             look_ahead_mask, padding_mask):
        outputs = {}

        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
        outputs['decoder_layer_embedding'] = x.numpy()

        x *= tf.sqrt(tf.to_float(self.d_model))
        outputs['decoder_layer_mult_sqrt'] = x.numpy()

        x += self.pos_encoding[:, :seq_len, :]
        x = self.dropout(x, training=training)
        outputs['decoder_layer_pos_encoding'] = x.numpy()

        for i in range(self.num_layers):
            x, block1, block2, dec_outs = self.dec_layers[i](x, enc_output, training,
                                                   look_ahead_mask, padding_mask, i)

            attention_weights[f'decoder_layer{i+1}_block1'] = block1
            attention_weights[f'decoder_layer{i+1}_block2'] = block2
            outputs.update(dec_outs)

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights, outputs
//  [python]
def call(self, x, enc_output, training,
             look_ahead_mask, padding_mask):
        outputs = {}

        seq_len = tf.shape(x)[1]
        attention_weights = {}

        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
        outputs['decoder_layer_embedding'] = x.numpy()

        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
        outputs['decoder_layer_mult_sqrt'] = x.numpy()

        x += self.pos_encoding[:, :seq_len, :]
        x = self.dropout(x, training=training)
        outputs['decoder_layer_pos_encoding'] = x.numpy()

        for i in range(self.num_layers):
            x, block1, block2, dec_outs = self.dec_layers[i](x, enc_output, training,
                                                   look_ahead_mask, padding_mask, i)

            attention_weights[f'decoder_layer{i+1}_block1'] = block1
            attention_weights[f'decoder_layer{i+1}_block2'] = block2
            outputs.update(dec_outs)

        # x.shape == (batch_size, target_seq_len, d_model)
        return x, attention_weights, outputs
//  [python]
def scaled_dot_product_attention(q, k, v, mask, outputs, name):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead)
    but it must be broadcastable for addition.

    Args:
      q: query shape == (..., seq_len_q, depth)
      k: key shape == (..., seq_len_k, depth)
      v: value shape == (..., seq_len_v, depth_v)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)
    outputs[f'{name}_sat_matmul_qk'] = matmul_qk.numpy()
    
    # scale matmul_qk
    dk = tf.to_float(tf.shape(k)[-1])
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
    outputs[f'{name}_sat_logits'] = scaled_attention_logits.numpy()
    
    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)

    outputs[f'{name}_sat_mask'] = scaled_attention_logits.numpy()
    # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)
    outputs[f'{name}_sat_softmax'] = attention_weights.numpy()
    
    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)
    outputs[f'{name}_sat_matmul_atten'] = output.numpy()
    
    return output, attention_weights