def METHOD_NAME(self, s):
    ss = ops.modal_dot(s, s, transpose_a=True)
    i_s = tf.eye(self.k, dtype=ss.dtype)
    ortho_loss = tf.norm(
        ss / tf.norm(ss, axis=(-1, -2), keepdims=True) - i_s / tf.norm(i_s),
        axis=(-1, -2),
    )
    return ortho_loss