src/transformers/models/gptj/modeling_tf_gptj.py/TFGPTJAttention/_split_heads
class TFGPTJAttention:    def _split_heads(self, hidden_states: tf.Tensor, rotary: bool) -> tf.Tensor:
        """
        Splits hidden dim into attn_head_size and num_attention_heads
        """
        new_shape = shape_list(hidden_states)[:-1] + [self.num_attention_heads, self.head_dim]
        hidden_states = tf.reshape(hidden_states, new_shape)
        if rotary:
            return hidden_states
        if len(shape_list(hidden_states)) == 4:
            return tf.transpose(hidden_states, (0, 2, 1, 3))  # (batch, head, seq_length, head_features)
        if len(shape_list(hidden_states)) == 5:
            return tf.transpose(hidden_states, (0, 1, 3, 2, 4))  # (batch, blocks, head, block_length, head_features)
        raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(shape_list(hidden_states))}")