src/transformers/models/big_bird/modeling_flax_big_bird.py/FlaxBigBirdBlockSparseAttention/transpose_for_scores
class FlaxBigBirdBlockSparseAttention:    def transpose_for_scores(x, n_heads, head_size):
        new_x_shape = x.shape[:-1] + (n_heads, head_size)
        x = x.reshape(*new_x_shape)
        return jnp.transpose(x, axes=(0, 2, 1, 3))