Question about Rotary Embedding Sequence in Model Code vs. Diagrams

[littlepsilon]

Hello,

I have been studying the Gemma-Pytorch implementation and noticed a potential discrepancy in the code related to the application of rotary embeddings in the attention mechanism.

Specifically, the rotary embeddings are applied before calculating the attention scores in the code between lines 213 and 331, which seems to contradict several diagrams that depict these embeddings being applied after the attention calculation.
code: https://github.com/google/gemma_pytorch/blob/main/gemma/model.py#L213
diagram: https://developers.googleblog.com/en/gemma-explained-new-in-gemma-2/

Could you please clarify whether the embeddings should indeed be applied before the attention scores are calculated, or if the diagrams might have represented this differently?

Thank you!