Complete type checking for OLMo support (builds on #816)

[taziksh]

Description

This PR builds on #816 and resolves all remaining type checking issues, allowing the OLMo implementation to pass mypy validation.

Changes:

• Added type assertions for OLMo2/OLMoE decoder layers to resolve union-attr errors
• Bumped minimum Python version to 3.10 for jaxtyping mypy plugin compatibility (Python 3.8 reached EOL Oct 2024)
• Fixed jaxtyping dimension annotation in create_alibi_multipliers (head_idx → n_heads)

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• New and existing unit tests pass locally with my changes

[taziksh]

@jonasrohw this is my PR with the type fixes

[taziksh]

Not sure who to tag to have this PR reviewed! @bryce13950 any pointers would be appreciated

[taziksh]

Bumping this again for visibility. OLMo 2’s open-source training data makes it a great fit for TransformerLens and could enable some interesting experiments. Would appreciate a review when possible.