scale by key_size in relational memory _multihead_attention
scale by key_size in relational memory _multihead_attention

GIT_ORIGIN_REV_ID=b92d54997e7df890458b672540c4f9832da9e8aa
PiperOrigin-RevId: 222243055

deepmind/sonnet,deepmind/sonnet