Throw error when using attn_in with grouped query attention

[degenfabian]

Description

When using attn_in with models that use GroupedQueryAttention, TransformerLens crashes because use_attn_in does not account for the different number of query and key/value heads when using GQA. For models with GQA use_split_qkv_input should be used instead, because it implements hooks for query, key and value heads and therefore can account for the different number of heads for each of them. This PR implements a more meaningful error message that informs the user to use split_qkv_input when working with models with GQA instead of use_attn_in.

This PR is not linked to a specific issue.

After adding a test case, it failed because of a beartype error that stated that rotary_base needs to be an integer instead of a float. I adjusted this accordingly in the configuration of google/gemma-2b

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility