Why does ALBERT use einsum in PyTorch implementation while in TF one it does not?
# ❓ Questions & Help
I wanted to learn internals of ALBERT model from your implementation (which is BTW really clean in comparison to the original one - good job!), but I've stumbled upon weird looking part in the `AlbertAttention`: https://github.com/huggingface/transformers/blob/6af3306a1da0322f58861b1fbb62ce5223d97b8a/src/transformers/modeling_albert.py#L258

Why does PyTorch version use `einsum`-based notation while calculating hidden state (with manual usage of `dense` layer's weights), while the TensorFlow version just reshapes the `context_layer` and does standard "forward" on dense layer?

https://github.com/huggingface/transformers/blob/6af3306a1da0322f58861b1fbb62ce5223d97b8a/src/transformers/modeling_tf_albert.py#L296

I would really like to know the explanation of this implementation - @LysandreJik cloud you shed some light here?

The two implementations are equivalent, but the  Pytorch version  is cumbersome.
I think  the code

```python
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

        # Should find a better way to do this
        w = (
            self.dense.weight.t()
            .view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
            .to(context_layer.dtype)
        )
        b = self.dense.bias.to(context_layer.dtype)

        projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
```
should be rewritten by

```python
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_shape = context_layer.size()[:-2] + (-1,)
        context_layer = context_layer.view(*new_shape)

        projected_context_layer = self.dense(context_layer)
```<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>Ping @LysandreJik <|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>Bump<|||||>Ping @LysandreJik again<|||||>For no particular reason, but it might not have been the best choice according to [this thread on performance](https://github.com/huggingface/transformers/issues/6771).<|||||>OK, thanks.