jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def hf_compute_loss(self, labels, logits):
        """CrossEntropyLoss that ignores pad tokens"""
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if self.config.tf_legacy_loss:
            melted_labels = tf.reshape(labels, (-1,))
            active_loss = tf.not_equal(melted_labels, self.config.pad_token_id)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(melted_labels, active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only non-padding labels affect the loss
        loss_mask = tf.to_float(labels != self.config.pad_token_id)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        """CrossEntropyLoss that ignores pad tokens"""
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if self.config.tf_legacy_loss:
            melted_labels = tf.reshape(labels, (-1,))
            active_loss = tf.not_equal(melted_labels, self.config.pad_token_id)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(melted_labels, active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only non-padding labels affect the loss
        loss_mask = tf.cast(labels != self.config.pad_token_id, dtype=unmasked_loss.dtype)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def _relative_position_bucket(relative_position, num_buckets=32, max_distance=128):
        ret = 0
        n = -relative_position

        num_buckets //= 2
        ret += t.cast(t.math.less(n, 0), dtype=relative_position.dtype) * num_buckets
        n = t.math.abs(n)

        # now n is in the range [0, inf)
        max_exact = num_buckets // 2
        is_small = t.math.less(n, max_exact)

        val_if_large = max_exact + t.cast(
            t.math.log(n / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact),
            dtype=relative_position.dtype,
        )

        val_if_large = t.math.minimum(val_if_large, num_buckets - 1)
        ret += t.where(is_small, n, val_if_large)
        return ret
//  [python]
def _relative_position_bucket(relative_position, num_buckets=32, max_distance=128):
        ret = 0
        n = -relative_position

        num_buckets //= 2
        ret += tf.cast(tf.math.less(n, 0), dtype=relative_position.dtype) * num_buckets
        n = tf.math.abs(n)

        # now n is in the range [0, inf)
        max_exact = num_buckets // 2
        is_small = tf.math.less(n, max_exact)

        val_if_large = max_exact + tf.cast(
            tf.math.log(n / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact),
            dtype=relative_position.dtype,
        )

        val_if_large = tf.math.minimum(val_if_large, num_buckets - 1)
        ret += tf.where(is_small, n, val_if_large)
        return ret
//  [python]
def call(self, hidden_states, attention_mask, head_mask, output_attentions, position_bias=None, training=False):
        batch_size = shape_list(hidden_states)[0]

        q = self.q(hidden_states)
        k = self.k(hidden_states)
        v = self.v(hidden_states)

        q = self.transpose_for_scores(q, batch_size)
        k = self.transpose_for_scores(k, batch_size)
        v = self.transpose_for_scores(v, batch_size)

        attention_scores = tf.matmul(q, k, transpose_b=True)
        dk = tf.to_float(shape_list(k)[-1])
        attention_scores = attention_scores / tf.math.sqrt(dk)

        # Apply relative position embedding (precomputed in MPNetEncoder) if provided.
        if position_bias is not None:
            attention_scores += position_bias

        if attention_mask is not None:
            attention_scores = attention_scores + attention_mask

        attention_probs = stable_softmax(attention_scores, axis=-1)

        attention_probs = self.dropout(attention_probs, training=training)

        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        c = tf.matmul(attention_probs, v)
        c = tf.transpose(c, perm=[0, 2, 1, 3])
        c = tf.reshape(c, (batch_size, -1, self.all_head_size))
        o = self.o(c)

        outputs = (o, attention_probs) if output_attentions else (o,)
        return outputs
//  [python]
def call(self, hidden_states, attention_mask, head_mask, output_attentions, position_bias=None, training=False):
        batch_size = shape_list(hidden_states)[0]

        q = self.q(hidden_states)
        k = self.k(hidden_states)
        v = self.v(hidden_states)

        q = self.transpose_for_scores(q, batch_size)
        k = self.transpose_for_scores(k, batch_size)
        v = self.transpose_for_scores(v, batch_size)

        attention_scores = tf.matmul(q, k, transpose_b=True)
        dk = tf.cast(shape_list(k)[-1], attention_scores.dtype)
        attention_scores = attention_scores / tf.math.sqrt(dk)

        # Apply relative position embedding (precomputed in MPNetEncoder) if provided.
        if position_bias is not None:
            attention_scores += position_bias

        if attention_mask is not None:
            attention_scores = attention_scores + attention_mask

        attention_probs = stable_softmax(attention_scores, axis=-1)

        attention_probs = self.dropout(attention_probs, training=training)

        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        c = tf.matmul(attention_probs, v)
        c = tf.transpose(c, perm=[0, 2, 1, 3])
        c = tf.reshape(c, (batch_size, -1, self.all_head_size))
        o = self.o(c)

        outputs = (o, attention_probs) if output_attentions else (o,)
        return outputs
//  [python]
def termination_condition(self, state):
        char_idx = t.cast(t.argmax(state.phi, axis=1), t.int32)
        final_char = char_idx >= self.attention_values_lengths - 1
        past_final_char = char_idx >= self.attention_values_lengths
        output = self.output_function(state)
        es = t.cast(output[:, 2], t.int32)
        is_eos = t.equal(es, t.experimental.numpy.ones_like(es))
        return t.logical_or(t.logical_and(final_char, is_eos), past_final_char)
//  [python]
def termination_condition(self, state):
        char_idx = tf.cast(tf.argmax(state.phi, axis=1), tf.int32)
        final_char = char_idx >= self.attention_values_lengths - 1
        past_final_char = char_idx >= self.attention_values_lengths
        output = self.output_function(state)
        es = tf.cast(output[:, 2], tf.int32)
        is_eos = tf.equal(es, tf.experimental.numpy.ones_like(es))
        return tf.logical_or(tf.logical_and(final_char, is_eos), past_final_char)
//  [python]
def output_function(self, state):
        params = dense_layer(state.h3, self.output_units, scope='gmm', reuse=tfcompat.AUTO_REUSE)
        pis, mus, sigmas, rhos, es = self._parse_parameters(params)
        mu1, mu2 = tf.split(mus, 2, axis=1)
        mus = tf.stack([mu1, mu2], axis=2)
        sigma1, sigma2 = tf.split(sigmas, 2, axis=1)

        covar_matrix = [tf.square(sigma1), rhos * sigma1 * sigma2,
                        rhos * sigma1 * sigma2, tf.square(sigma2)]
        covar_matrix = tf.stack(covar_matrix, axis=2)
        covar_matrix = tf.reshape(covar_matrix, (self.batch_size, self.num_output_mixture_components, 2, 2))

        mvn = tfp.distributions.MultivariateNormalFullCovariance(loc=mus, covariance_matrix=covar_matrix)
        b = tfd.Bernoulli(probs=es)
        c = tfd.Categorical(probs=pis)

        sampled_e = b.sample()
        sampled_coords = mvn.sample()
        sampled_idx = c.sample()

        idx = tf.stack([tf.range(self.batch_size), sampled_idx], axis=1)
        coords = tf.gather_nd(sampled_coords, idx)
        return tf.concat([coords, tf.to_float(sampled_e), axis=1)
//  [python]
def output_function(self, state):
        params = dense_layer(state.h3, self.output_units, scope='gmm', reuse=tfcompat.AUTO_REUSE)
        pis, mus, sigmas, rhos, es = self._parse_parameters(params)
        mu1, mu2 = tf.split(mus, 2, axis=1)
        mus = tf.stack([mu1, mu2], axis=2)
        sigma1, sigma2 = tf.split(sigmas, 2, axis=1)

        covar_matrix = [tf.square(sigma1), rhos * sigma1 * sigma2,
                        rhos * sigma1 * sigma2, tf.square(sigma2)]
        covar_matrix = tf.stack(covar_matrix, axis=2)
        covar_matrix = tf.reshape(covar_matrix, (self.batch_size, self.num_output_mixture_components, 2, 2))

        mvn = tfp.distributions.MultivariateNormalFullCovariance(loc=mus, covariance_matrix=covar_matrix)
        b = tfd.Bernoulli(probs=es)
        c = tfd.Categorical(probs=pis)

        sampled_e = b.sample()
        sampled_coords = mvn.sample()
        sampled_idx = c.sample()

        idx = tf.stack([tf.range(self.batch_size), sampled_idx], axis=1)
        coords = tf.gather_nd(sampled_coords, idx)
        return tf.concat([coords, tf.cast(sampled_e, tf.float32)], axis=1)
//  [python]
def _build_model(self):
        with tf.name_scope('Inputs'):
            self.X_inputs       = tf.placeholder(tf.float32, [None, self.n_Channels, self.n_Samples, 1], name='X_inputs')
            self.y_inputs       = tf.placeholder(tf.int64, None, name='y_inputs')
            self.tf_is_training = tf.placeholder(tf.bool, None, name='is_training')
            self.tf_bn_training = tf.placeholder(tf.bool, None, name='is_training')

        with tf.variable_scope('EEGNet'):
            self.output = EEGNet_Classifier(self.X_inputs, self.n_classes, self.n_Channels, self.n_Samples, 
                        dropoutRate=self.dropout, kernLength=self.fs//2, training=self.tf_is_training,
                        bn_training=self.tf_bn_training)
        
        with tf.name_scope('Accuracy'):
            prediction = tf.argmax(self.output, axis=1)
            correct_prediction = tf.equal(prediction, self.y_inputs)
            self._acc = tf.reduce_mean(tf.to_float(correct_prediction))

        with tf.name_scope('Loss'):
            self._loss = tf.losses.sparse_softmax_cross_entropy(labels=self.y_inputs, logits=self.output)
//  [python]
def _build_model(self):
        with tf.name_scope('Inputs'):
            self.X_inputs       = tf.placeholder(tf.float32, [None, self.n_Channels, self.n_Samples, 1], name='X_inputs')
            self.y_inputs       = tf.placeholder(tf.int64, None, name='y_inputs')
            self.tf_is_training = tf.placeholder(tf.bool, None, name='is_training')
            self.tf_bn_training = tf.placeholder(tf.bool, None, name='is_training')

        with tf.variable_scope('EEGNet'):
            self.output = EEGNet_Classifier(self.X_inputs, self.n_classes, self.n_Channels, self.n_Samples, 
                        dropoutRate=self.dropout, kernLength=self.fs//2, training=self.tf_is_training,
                        bn_training=self.tf_bn_training)
        
        with tf.name_scope('Accuracy'):
            prediction = tf.argmax(self.output, axis=1)
            correct_prediction = tf.equal(prediction, self.y_inputs)
            self._acc = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

        with tf.name_scope('Loss'):
            self._loss = tf.losses.sparse_softmax_cross_entropy(labels=self.y_inputs, logits=self.output)
//  [python]
def __call__(self, inputs, state, scope=None):
        with tfcompat.variable_scope(scope or type(self).__name__, reuse=tfcompat.AUTO_REUSE):
            # lstm 1
            s1_in = tf.concat([state.w, inputs], axis=1)
            cell1 = tfcompat.nn.rnn_cell.LSTMCell(self.lstm_size)
            s1_out, s1_state = cell1(s1_in, state=(state.c1, state.h1))

            # attention
            attention_inputs = tf.concat([state.w, inputs, s1_out], axis=1)
            attention_params = dense_layer(attention_inputs, 3 * self.num_attn_mixture_components, scope='attention')
            alpha, beta, kappa = tf.split(tf.nn.softplus(attention_params), 3, axis=1)
            kappa = state.kappa + kappa / 25.0
            beta = tf.clip_by_value(beta, .01, np.inf)

            kappa_flat, alpha_flat, beta_flat = kappa, alpha, beta
            kappa, alpha, beta = tf.expand_dims(kappa, 2), tf.expand_dims(alpha, 2), tf.expand_dims(beta, 2)

            enum = tf.reshape(tf.range(self.char_len), (1, 1, self.char_len))
            u = tf.tile(enum, (self.batch_size, self.num_attn_mixture_components, 1))
            phi_flat = tf.reduce_sum(alpha * tf.exp(-tf.square(kappa - u) / beta), axis=1)

            phi = tf.expand_dims(phi_flat, 2)
            sequence_mask = tf.cast(tf.sequence_mask(self.attention_values_lengths, maxlen=self.char_len), tf.float32)
            sequence_mask = tf.expand_dims(sequence_mask, 2)
            w = tf.reduce_sum(phi * self.attention_values * sequence_mask, axis=1)

            # lstm 2
            s2_in = tf.concat([inputs, s1_out, w], axis=1)
            cell2 = tfcompat.nn.rnn_cell.LSTMCell(self.lstm_size)
            s2_out, s2_state = cell2(s2_in, state=(state.c2, state.h2))

            # lstm 3
            s3_in = tf.concat([inputs, s2_out, w], axis=1)
            cell3 = tfcompat.nn.rnn_cell.LSTMCell(self.lstm_size)
            s3_out, s3_state = cell3(s3_in, state=(state.c3, state.h3))

            new_state = LSTMAttentionCellState(
                s1_state.h,
                s1_state.c,
                s2_state.h,
                s2_state.c,
                s3_state.h,
                s3_state.c,
                alpha_flat,
                beta_flat,
                kappa_flat,
                w,
                phi_flat,
            )

            return s3_out, new_state
//  [python]
def __call__(self, inputs, state, scope=None):
        with tfcompat.variable_scope(scope or type(self).__name__, reuse=tfcompat.AUTO_REUSE):
            # lstm 1
            s1_in = tf.concat([state.w, inputs], axis=1)
            cell1 = tfcompat.nn.rnn_cell.LSTMCell(self.lstm_size)
            s1_out, s1_state = cell1(s1_in, state=(state.c1, state.h1))

            # attention
            attention_inputs = tf.concat([state.w, inputs, s1_out], axis=1)
            attention_params = dense_layer(attention_inputs, 3 * self.num_attn_mixture_components, scope='attention')
            alpha, beta, kappa = tf.split(tf.nn.softplus(attention_params), 3, axis=1)
            kappa = state.kappa + kappa / 25.0
            beta = tf.clip_by_value(beta, .01, np.inf)

            kappa_flat, alpha_flat, beta_flat = kappa, alpha, beta
            kappa, alpha, beta = tf.expand_dims(kappa, 2), tf.expand_dims(alpha, 2), tf.expand_dims(beta, 2)

            enum = tf.reshape(tf.range(self.char_len), (1, 1, self.char_len))
            u = tf.cast(tf.tile(enum, (self.batch_size, self.num_attn_mixture_components, 1)), tf.float32)
            phi_flat = tf.reduce_sum(alpha * tf.exp(-tf.square(kappa - u) / beta), axis=1)

            phi = tf.expand_dims(phi_flat, 2)
            sequence_mask = tf.cast(tf.sequence_mask(self.attention_values_lengths, maxlen=self.char_len), tf.float32)
            sequence_mask = tf.expand_dims(sequence_mask, 2)
            w = tf.reduce_sum(phi * self.attention_values * sequence_mask, axis=1)

            # lstm 2
            s2_in = tf.concat([inputs, s1_out, w], axis=1)
            cell2 = tfcompat.nn.rnn_cell.LSTMCell(self.lstm_size)
            s2_out, s2_state = cell2(s2_in, state=(state.c2, state.h2))

            # lstm 3
            s3_in = tf.concat([inputs, s2_out, w], axis=1)
            cell3 = tfcompat.nn.rnn_cell.LSTMCell(self.lstm_size)
            s3_out, s3_state = cell3(s3_in, state=(state.c3, state.h3))

            new_state = LSTMAttentionCellState(
                s1_state.h,
                s1_state.c,
                s2_state.h,
                s2_state.c,
                s3_state.h,
                s3_state.c,
                alpha_flat,
                beta_flat,
                kappa_flat,
                w,
                phi_flat,
            )

            return s3_out, new_state
//  [python]
def kl_cutoff_loss(
      self, kl_divergence: types.Tensor, debug_summaries: bool = False
  ) -> types.Tensor:
    # Squared penalization for mean KL divergence above some threshold.
    if self._kl_cutoff_factor <= 0.0:
      return tf.constant(0.0, dtype=tf.float32, name='zero_kl_cutoff_loss')
    kl_cutoff = self._kl_cutoff_factor * self._adaptive_kl_target
    mean_kl = tf.reduce_mean(input_tensor=kl_divergence)
    kl_over_cutoff = tf.maximum(mean_kl - kl_cutoff, 0.0)
    kl_cutoff_loss = self._kl_cutoff_coef * tf.square(kl_over_cutoff)

    if debug_summaries:
      tf.compat.v2.summary.scalar(
          name='kl_cutoff_count',
          data=tf.reduce_sum(
              input_tensor=tf.to_int64(kl_divergence > kl_cutoff)
          ),
          step=self.train_step_counter,
      )
      tf.compat.v2.summary.scalar(
          name='kl_cutoff_loss',
          data=kl_cutoff_loss,
          step=self.train_step_counter,
      )

    return tf.identity(kl_cutoff_loss, name='kl_cutoff_loss')
//  [python]
def kl_cutoff_loss(
      self, kl_divergence: types.Tensor, debug_summaries: bool = False
  ) -> types.Tensor:
    # Squared penalization for mean KL divergence above some threshold.
    if self._kl_cutoff_factor <= 0.0:
      return tf.constant(0.0, dtype=tf.float32, name='zero_kl_cutoff_loss')
    kl_cutoff = self._kl_cutoff_factor * self._adaptive_kl_target
    mean_kl = tf.reduce_mean(input_tensor=kl_divergence)
    kl_over_cutoff = tf.maximum(mean_kl - kl_cutoff, 0.0)
    kl_cutoff_loss = self._kl_cutoff_coef * tf.square(kl_over_cutoff)

    if debug_summaries:
      tf.compat.v2.summary.scalar(
          name='kl_cutoff_count',
          data=tf.reduce_sum(
              input_tensor=tf.cast(kl_divergence > kl_cutoff, dtype=tf.int64)
          ),
          step=self.train_step_counter,
      )
      tf.compat.v2.summary.scalar(
          name='kl_cutoff_loss',
          data=kl_cutoff_loss,
          step=self.train_step_counter,
      )

    return tf.identity(kl_cutoff_loss, name='kl_cutoff_loss')
//  [python]
def compute_loss(labels, predictions,loss_object,batch_size):

    per_example_loss = loss_object(labels, predictions)

    per_example_loss /= tf.to_float(tf.reduce_prod(tf.shape(labels)[1:]))

    return tf.nn.compute_average_loss(per_example_loss, global_batch_size = batch_size)
//  [python]
def compute_loss(labels, predictions,loss_object,batch_size):

    per_example_loss = loss_object(labels, predictions)

    per_example_loss /= tf.cast(tf.reduce_prod(tf.shape(labels)[1:]), tf.float32)

    return tf.nn.compute_average_loss(per_example_loss, global_batch_size = batch_size)
//  [python]
logits = []
costs = []
i = 0
accuracy = 0
# for output in outputs:
for i in range(0, max_word_length):
	output = outputs[-i - 1]
	uniform = tf.random_uniform([num_hidden, classes], minval=-1. / width, maxval=1. / width)
	weights = tf.Variable(uniform, name="weights_%d" % i)
	uniform_bias = tf.random_uniform([classes], minval=-1. / width, maxval=1. / width)
	bias = tf.Variable(uniform_bias, name="bias_dense_%d" % i)
	y_ = outputY = tf.matmul(output, weights, name="dense_%d" % i) + bias

	cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_, labels=y[:, i, :]), name="cost")  # prediction, target
	costs.append(cost)
	logits.append(y_)

	correct_pred = tf.equal(tf.argmax(outputY, 1), tf.argmax(y[:, i], 1))
	accuraci = tf.reduce_mean(tf.to_float(correct_pred))
	accuracy += accuraci
//  [python]
logits = []
costs = []
i = 0
accuracy = 0
# for output in outputs:
for i in range(0, max_word_length):
	output = outputs[-i - 1]
	uniform = tf.random_uniform([num_hidden, classes], minval=-1. / width, maxval=1. / width)
	weights = tf.Variable(uniform, name="weights_%d" % i)
	uniform_bias = tf.random_uniform([classes], minval=-1. / width, maxval=1. / width)
	bias = tf.Variable(uniform_bias, name="bias_dense_%d" % i)
	y_ = outputY = tf.matmul(output, weights, name="dense_%d" % i) + bias

	cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_, labels=y[:, i, :]), name="cost")  # prediction, target
	costs.append(cost)
	logits.append(y_)

	correct_pred = tf.equal(tf.argmax(outputY, 1), tf.argmax(y[:, i], 1))
	accuraci = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
	accuracy += accuraci
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c, top_span_cdd_pool_flag=None):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k) # [k]

    num_cdd_in_pool = tf.reduce_sum(tf.to_int32(top_span_cdd_pool_flag))
    num_cdd_in_dial = k - num_cdd_in_pool
    top_span_range_cdd = tf.concat([tf.zeros(num_cdd_in_pool, tf.int32), tf.range(1, num_cdd_in_dial + 1)], 0)
    antecedent_offsets = tf.expand_dims(top_span_range_cdd, 1) - tf.expand_dims(top_span_range_cdd, 0) # [k, k]

    antecedents_mask = antecedent_offsets >= 1 # [k, k]
    fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0) # [k, k]
    fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask)) # [k, k]
    fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb) # [k, k]

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, top_span_mention_scores, c, top_span_cdd_pool_flag=None):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k) # [k]

    num_cdd_in_pool = tf.reduce_sum(tf.cast(top_span_cdd_pool_flag, tf.int32))
    num_cdd_in_dial = k - num_cdd_in_pool
    top_span_range_cdd = tf.concat([tf.zeros(num_cdd_in_pool, tf.int32), tf.range(1, num_cdd_in_dial + 1)], 0)
    antecedent_offsets = tf.expand_dims(top_span_range_cdd, 1) - tf.expand_dims(top_span_range_cdd, 0) # [k, k]

    antecedents_mask = antecedent_offsets >= 1 # [k, k]
    fast_antecedent_scores = tf.expand_dims(top_span_mention_scores, 1) + tf.expand_dims(top_span_mention_scores, 0) # [k, k]
    fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask)) # [k, k]
    fast_antecedent_scores += self.get_fast_antecedent_scores(top_span_emb) # [k, k]

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def find_top_k(user, item_factors, k):
    all_items = tf.matmul(a = tf.expand_dims(input = user, axis = 0), b = tf.transpose(a = item_factors))
    topk = tf.nn.top_k(input = all_items, k = k)
    return tf.to_int64(x = topk.indices)
//  [python]
def find_top_k(user, item_factors, k):
    all_items = tf.matmul(a = tf.expand_dims(input = user, axis = 0), b = tf.transpose(a = item_factors))
    topk = tf.nn.top_k(input = all_items, k = k)
    return tf.cast(x = topk.indices, dtype = tf.int64)
//  [python]
def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int):
    if pad_token_id is None:
        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
    pad_token_id = tf.cast(pad_token_id, input_ids.dtype)

    if decoder_start_token_id is None:
        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
    decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype)

    start_tokens = tf.fill((shape_list(input_ids)[0], 1), decoder_start_token_id)
    shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)
    # replace possible -100 values in labels by `pad_token_id`
    shifted_input_ids = tf.where(
        shifted_input_ids == -100, tf.fill(shape_list(shifted_input_ids), pad_token_id), shifted_input_ids
    )

    # "Verify that `labels` has only positive values and -100"
    assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype))

    # Make sure the assertion op is called by wrapping the result in an identity no-op
    with tf.control_dependencies([assert_gte0]):
        shifted_input_ids = tf.identity(shifted_input_ids)

    return shifted_input_ids
//  [python]
def load_img_and_reshape(path_to_img):
    max_dim = 512
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_image(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)

    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    tf.print(f"raw_img_shape {shape}")
    long_dim = max(shape)
    scale = max_dim / long_dim

    new_shape = tf.to_int32(shape * scale)
    tf.print(f"new_shape {new_shape}")
    img = tf.image.resize(img, new_shape)
    img = img[tf.newaxis, :]
    return img
//  [python]
def load_img_and_reshape(path_to_img):
    max_dim = 512
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_image(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)

    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    tf.print(f"raw_img_shape {shape}")
    long_dim = max(shape)
    scale = max_dim / long_dim

    new_shape = tf.cast(shape * scale, tf.int32)
    tf.print(f"new_shape {new_shape}")
    img = tf.image.resize(img, new_shape)
    img = img[tf.newaxis, :]
    return img
//  [python]
def reset(self):
    obs = ag.utils.wrap_py_func(self.env.reset, tf.float64, ())
    obs = tf.reshape(obs, (1, 4))
    obs = tf.cast(obs, tf.float32)
    return obs
//  [python]
def _build_trig_vector(length, key_dim):
  """Builds the trig vector."""
  tf_dtype = tf_keras.mixed_precision.global_policy().compute_dtype
  position_ids = t.cast(t.range(length), dtype=tf_dtype)
  position_ids = t.expand_dims(position_ids, axis=0)
  steps = key_dim // 2
  # 2 (i - 1) / key_dim = (i - 1) / steps: (-1 achieved with zero-indexing)
  wavenumber_exponent = -t.cast(t.range(steps), dtype=tf_dtype) / steps
  wavenumbers = t.pow(
      t.constant(10000.0, dtype=tf_dtype), wavenumber_exponent
  )
  vec = t.einsum('bl,d->bld', position_ids, wavenumbers)
  sin_vec = t.repeat(t.sin(vec), repeats=2, axis=-1)
  cos_vec = t.repeat(t.cos(vec), repeats=2, axis=-1)
  sin_vec, cos_vec = t.expand_dims(sin_vec, 2), t.expand_dims(cos_vec, 2)
  return sin_vec, cos_vec
//  [python]
def _build_trig_vector(length, key_dim):
  """Builds the trig vector."""
  tf_dtype = tf_keras.mixed_precision.global_policy().compute_dtype
  position_ids = tf.cast(tf.range(length), dtype=tf_dtype)
  position_ids = tf.expand_dims(position_ids, axis=0)
  steps = key_dim // 2
  # 2 (i - 1) / key_dim = (i - 1) / steps: (-1 achieved with zero-indexing)
  wavenumber_exponent = -tf.cast(tf.range(steps), dtype=tf_dtype) / steps
  wavenumbers = tf.pow(
      tf.constant(10000.0, dtype=tf_dtype), wavenumber_exponent
  )
  vec = tf.einsum('bl,d->bld', position_ids, wavenumbers)
  sin_vec = tf.repeat(tf.sin(vec), repeats=2, axis=-1)
  cos_vec = tf.repeat(tf.cos(vec), repeats=2, axis=-1)
  sin_vec, cos_vec = tf.expand_dims(sin_vec, 2), tf.expand_dims(cos_vec, 2)
  return sin_vec, cos_vec
//  [python]
def _get_cross_entropy_and_save_actions(self, inputs):
    logits, actions = self(inputs)  # pylint:disable=not-callable
    self._current_actions = actions
    labels = 1.0 - tf.to_float(actions)
    return tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)
//  [python]
def _get_cross_entropy_and_save_actions(self, inputs):
    logits, actions = self(inputs)  # pylint:disable=not-callable
    self._current_actions = actions
    labels = 1.0 - tf.cast(actions, tf.float32)
    return tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            tf.keras.regularizers.l2(WEIGHT_DECAY)(w) / tf.to_float(tf.size(w))
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """
        # Optimizer object
        optimizer = keras.optimizers.SGD(
            lr=learning_rate, momentum=momentum,
            clipnorm=GRADIENT_CLIP_NORM)
        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            tf.keras.regularizers.l2(WEIGHT_DECAY)(w) / tf.cast(tf.size(w), tf.float32)
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def unique_counts(
    x: Union[tf.Tensor, tf.Variable],
    /,
) -> Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]]:
    Results = namedtuple("Results", ["values", "counts"])
    v, _, c = tf.unique_with_counts(tf.sort(tf.reshape(x, [-1])))
    v = tf.cast(v, dtype=x.dtype)
    c = tf.cast(c, dtype=tf.int64)
    return Results(v, c)
//  [python]
def roi_pooling(self, feature_maps, rois, img_shape, scope):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping_' + scope):
            img_h, img_w = tf.to_float(img_shape[1]), tf.to_float(img_shape[2])
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features
//  [python]
def roi_pooling(self, feature_maps, rois, img_shape, scope):
        '''
        Here use roi warping as roi_pooling

        :param featuremaps_dict: feature map to crop
        :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
        :return:
        '''

        with tf.variable_scope('ROI_Warping_' + scope):
            img_h, img_w = tf.cast(img_shape[1], tf.float32), tf.cast(img_shape[2], tf.float32)
            N = tf.shape(rois)[0]
            x1, y1, x2, y2 = tf.unstack(rois, axis=1)

            normalized_x1 = x1 / img_w
            normalized_x2 = x2 / img_w
            normalized_y1 = y1 / img_h
            normalized_y2 = y2 / img_h

            normalized_rois = tf.transpose(
                tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')

            normalized_rois = tf.stop_gradient(normalized_rois)

            cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
                                                            box_ind=tf.zeros(shape=[N, ],
                                                                             dtype=tf.int32),
                                                            crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
                                                            name='CROP_AND_RESIZE'
                                                            )
            roi_features = slim.max_pool2d(cropped_roi_features,
                                           [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
                                           stride=cfgs.ROI_POOL_KERNEL_SIZE)

        return roi_features
//  [python]
def unique_inverse(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: Optional[int] = None,
) -> Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]]:
    Results = namedtuple("Results", ["values", "inverse_indices"])
    if axis is None:
        x = tf.reshape(x, shape=(-1,))
        axis = 0

    flat_tensor = tf.reshape(x, -1)
    values = tf.unique(tf.sort(flat_tensor))[0]
    values = t.cast(values, dtype=x.dtype)

    values_list = values.numpy().tolist()
    inverse_indices = [values_list.index(val) for val in flat_tensor.numpy().tolist()]
    inverse_indices = tf.reshape(tf.convert_to_tensor(inverse_indices), x.shape)
    inverse_indices = t.cast(inverse_indices, dtype=tf.int64)
    return Results(values, inverse_indices)
//  [python]
def unique_inverse(
    x: Union[tf.Tensor, tf.Variable],
    /,
    *,
    axis: Optional[int] = None,
) -> Tuple[Union[tf.Tensor, tf.Variable], Union[tf.Tensor, tf.Variable]]:
    Results = namedtuple("Results", ["values", "inverse_indices"])
    if axis is None:
        x = tf.reshape(x, shape=(-1,))
        axis = 0

    flat_tensor = tf.reshape(x, -1)
    values = tf.unique(tf.sort(flat_tensor))[0]
    values = tf.cast(values, dtype=x.dtype)

    values_list = values.numpy().tolist()
    inverse_indices = [values_list.index(val) for val in flat_tensor.numpy().tolist()]
    inverse_indices = tf.reshape(tf.convert_to_tensor(inverse_indices), x.shape)
    inverse_indices = tf.cast(inverse_indices, dtype=tf.int64)
    return Results(values, inverse_indices)
//  [python]
def get_accuracy(logits, actuals):
    # Calulate if each output is correct
    batch_acc = tf.equal(tf.argmax(logits, 1), tf.cast(actuals, tf.int64))
    # Convert logical to float
    batch_acc = tf.to_float(batch_acc)
    return batch_acc
//  [python]
def get_accuracy(logits, actuals):
    # Calulate if each output is correct
    batch_acc = tf.equal(tf.argmax(logits, 1), tf.cast(actuals, tf.int64))
    # Convert logical to float
    batch_acc = tf.cast(batch_acc, tf.float32)
    return batch_acc
//  [python]
def train(self, cart_pole_env, optimizer, discount_rate, num_games,
            max_steps_per_game):
    var_list = tf.trainable_variables()
    grad_list = [
        tf.TensorArray(tf.float32, 0, dynamic_size=True) for _ in var_list
    ]

    step_counts = []
    discounted_rewards = []
    ag.set_element_type(discounted_rewards, tf.float32)
    ag.set_element_type(step_counts, tf.int32)

    # Note: we use a shared object, cart_pole_env here. Because calls to the
    # object's method are made through py_func, TensorFlow cannot detect its
    # data dependencies. Hence we must manually synchronize access to it
    # and ensure the control dependencies are set in such a way that
    # calls to reset(), take_one_step, etc. are made in the correct order.
    sync_counter = tf.constant(0)

    for _ in tf.range(num_games):
      with tf.control_dependencies([sync_counter]):
        obs = cart_pole_env.reset()
        with tf.control_dependencies([obs]):
          sync_counter += 1

        game_rewards = []
        ag.set_element_type(game_rewards, tf.float32)

        for step in tf.range(max_steps_per_game):
          logits, actions = self(obs)  # pylint:disable=not-callable
          logits = tf.reshape(logits, ())
          actions = tf.reshape(actions, ())

          labels = 1.0 - tf.to_float(actions)
          loss = tf.nn.sigmoid_cross_entropy_with_logits(
              labels=labels, logits=logits)
          grads = tf.gradients(loss, var_list)

          for i in range(len(grads)):
            grad_list[i].append(grads[i])

          with tf.control_dependencies([sync_counter]):
            obs, reward, done = cart_pole_env.step(actions)
            with tf.control_dependencies([obs]):
              sync_counter += 1
            obs = tf.reshape(obs, (1, 4))

          game_rewards.append(reward)
          if reward < 0.1 or done:
            step_counts.append(step + 1)
            break

        discounted_rewards = graph_append_discounted_rewards(
            discounted_rewards, game_rewards, discount_rate)

    discounted_rewards = ag.stack(discounted_rewards)
    discounted_rewards.set_shape((None,))
    mean, variance = tf.nn.moments(discounted_rewards, [0])
    normalized_rewards = (discounted_rewards - mean) / tf.sqrt(variance)

    for i in range(len(grad_list)):
      g = ag.stack(grad_list[i])

      # This block just adjusts the shapes to match for multiplication.
      r = normalized_rewards
      if r.shape.ndims < g.shape.ndims:
        r = tf.expand_dims(r, -1)
      if r.shape.ndims < g.shape.ndims:
        r = tf.expand_dims(r, -1)

      grad_list[i] = tf.reduce_mean(g * r, axis=0)

    optimizer.apply_gradients(
        zip(grad_list, var_list), global_step=tf.train.get_global_step())

    return ag.stack(step_counts)
//  [python]
def train(self, cart_pole_env, optimizer, discount_rate, num_games,
            max_steps_per_game):
    var_list = tf.trainable_variables()
    grad_list = [
        tf.TensorArray(tf.float32, 0, dynamic_size=True) for _ in var_list
    ]

    step_counts = []
    discounted_rewards = []
    ag.set_element_type(discounted_rewards, tf.float32)
    ag.set_element_type(step_counts, tf.int32)

    # Note: we use a shared object, cart_pole_env here. Because calls to the
    # object's method are made through py_func, TensorFlow cannot detect its
    # data dependencies. Hence we must manually synchronize access to it
    # and ensure the control dependencies are set in such a way that
    # calls to reset(), take_one_step, etc. are made in the correct order.
    sync_counter = tf.constant(0)

    for _ in tf.range(num_games):
      with tf.control_dependencies([sync_counter]):
        obs = cart_pole_env.reset()
        with tf.control_dependencies([obs]):
          sync_counter += 1

        game_rewards = []
        ag.set_element_type(game_rewards, tf.float32)

        for step in tf.range(max_steps_per_game):
          logits, actions = self(obs)  # pylint:disable=not-callable
          logits = tf.reshape(logits, ())
          actions = tf.reshape(actions, ())

          labels = 1.0 - tf.cast(actions, tf.float32)
          loss = tf.nn.sigmoid_cross_entropy_with_logits(
              labels=labels, logits=logits)
          grads = tf.gradients(loss, var_list)

          for i in range(len(grads)):
            grad_list[i].append(grads[i])

          with tf.control_dependencies([sync_counter]):
            obs, reward, done = cart_pole_env.step(actions)
            with tf.control_dependencies([obs]):
              sync_counter += 1
            obs = tf.reshape(obs, (1, 4))

          game_rewards.append(reward)
          if reward < 0.1 or done:
            step_counts.append(step + 1)
            break

        discounted_rewards = graph_append_discounted_rewards(
            discounted_rewards, game_rewards, discount_rate)

    discounted_rewards = ag.stack(discounted_rewards)
    discounted_rewards.set_shape((None,))
    mean, variance = tf.nn.moments(discounted_rewards, [0])
    normalized_rewards = (discounted_rewards - mean) / tf.sqrt(variance)

    for i in range(len(grad_list)):
      g = ag.stack(grad_list[i])

      # This block just adjusts the shapes to match for multiplication.
      r = normalized_rewards
      if r.shape.ndims < g.shape.ndims:
        r = tf.expand_dims(r, -1)
      if r.shape.ndims < g.shape.ndims:
        r = tf.expand_dims(r, -1)

      grad_list[i] = tf.reduce_mean(g * r, axis=0)

    optimizer.apply_gradients(
        zip(grad_list, var_list), global_step=tf.train.get_global_step())

    return ag.stack(step_counts)
//  [python]
def scaled_dot_product_attention(q, k, v, mask, is_training=True, rel_key_emb=None, rel_val_emb=None):
        # attn_dim: num_joints for spatial and seq_len for temporal
        '''
        The scaled dot product attention mechanism introduced in the Transformer
        :param q: the query vectors matrix (..., attn_dim, d_model/num_heads)
        :param k: the key vector matrix (..., attn_dim, d_model/num_heads)
        :param v: the value vector matrix (..., attn_dim, d_model/num_heads)
        :param mask: a mask for attention
        :return: the updated encoding and the attention weights matrix
        '''

        matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., num_heads, attn_dim, attn_dim)

        batch_size = tf.shape(q)[0]
        heads = tf.shape(q)[1]
        length = tf.shape(q)[2]
        
        if rel_key_emb is not None:
            q_t = tf.transpose(q, [2, 0, 1, 3])
            q_t_r = tf.reshape(q_t, [length, heads*batch_size, -1])
            q_tz_matmul = tf.matmul(q_t_r, rel_key_emb, transpose_b=True)
            q_tz_matmul_r = tf.reshape(q_tz_matmul, [length, batch_size, heads, -1])
            q_tz_matmul_r_t = tf.transpose(q_tz_matmul_r, [1, 2, 0, 3])
            matmul_qk += q_tz_matmul_r_t

        # scale matmul_qk
        dk = tf.to_float(tf.shape(k)[-1])
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        # add the mask to the scaled tensor.
        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        # 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)  # (..., num_heads, attn_dim, attn_dim)

        # attention_weights = tf.layers.dropout(attention_weights, training=is_training, rate=0.2)
        output = tf.matmul(attention_weights, v)  # (..., num_heads, attn_dim, depth)
        
        if rel_val_emb is not None:
            w_t = tf.transpose(attention_weights, [2, 0, 1, 3])
            w_t_r = tf.reshape(w_t, [length, heads*batch_size, -1])
            w_tz_matmul = tf.matmul(w_t_r, rel_val_emb, transpose_b=False)
            w_tz_matmul_r = tf.reshape(w_tz_matmul, [length, batch_size, heads, -1])
            w_tz_matmul_r_t = tf.transpose(w_tz_matmul_r, [1, 2, 0, 3])
            output += w_tz_matmul_r_t

        return output, attention_weights