src/transformers/models/esm/modeling_tf_esm.py/TFEsmContactPredictionHead/call
class TFEsmContactPredictionHead:    def call(self, tokens, attentions):
        # remove eos token attentions
        eos_mask = tf.cast(tokens != self.eos_idx, attentions.dtype)
        eos_mask = tf.expand_dims(eos_mask, 1) * tf.expand_dims(eos_mask, 2)
        attentions = attentions * eos_mask[:, None, None, :, :]
        attentions = attentions[..., :-1, :-1]
        # remove cls token attentions
        attentions = attentions[..., 1:, 1:]
        batch_size, layers, heads, seqlen, _ = shape_list(attentions)
        attentions = tf.reshape(attentions, (batch_size, layers * heads, seqlen, seqlen))

        # features: batch x channels x tokens x tokens (symmetric)
        attentions = average_product_correct(symmetrize(attentions))
        attentions = tf.transpose(attentions, perm=(0, 2, 3, 1))
        return tf.squeeze(self.regression(attentions), 3)