src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py/TFWav2Vec2Attention/_shape
class TFWav2Vec2Attention:    def _shape(self, tensor: tf.Tensor, seq_len: int, bsz: int):
        return tf.transpose(tf.reshape(tensor, (bsz, seq_len, self.num_heads, self.head_dim)), (0, 2, 1, 3))