src/transformers/models/groupvit/modeling_tf_groupvit.py/TFGroupViTAssignAttention/call
class TFGroupViTAssignAttention:    def call(self, query: tf.Tensor, key: tf.Tensor, training: bool = False):
        value = key
        # [batch_size, query_length, channels]
        query = self.q_proj(query)

        # [batch_size, key_length, channels]
        key = self.k_proj(key)

        # [batch_size, key_length, channels]
        value = self.v_proj(value)

        # [batch_size, query_length, key_length]
        raw_attn = tf.matmul(query, key, transpose_b=True) * self.scale

        attn = self.get_attn(raw_attn, training=training)
        soft_attn = self.get_attn(raw_attn, training=training, gumbel=False, hard=False)

        attn = attn / (tf.math.reduce_sum(attn, axis=-1, keepdims=True) + self.assign_eps)

        out = tf.matmul(attn, value)

        out = self.proj(out)

        return out, soft_attn