def symmetrize(x):
    "Make layer symmetric in final two dimensions, used for contact prediction."
    return x + tf.linalg.matrix_transpose(x)  # Transposes last two dimensions only
Make layer symmetric in final two dimensions, used for contact prediction.
dpkd/transformers/src/transformers/models/esm/modeling_tf_esm.py