Clarify some core attention code via variable renaming and function extraction.
Clarify some core attention code via variable renaming and function extraction.

Also fix a docstring inaccuracy.

PiperOrigin-RevId: 358521061

google/trax,google/trax