[Bug Report] Should we convert tensor dtype before eigenvalue decomposition

[notoookay]

If you are submitting a bug report, please fill in the following details and use the tag [bug].

Describe the bug If we load the model with bfloat16, calculating OV eigenvalue could be buggy as torch.linalg.eig doesn't support bfloat16.

Code example

from transformer_lens import HookedTransformer

model = HookedTransformer.from_pretrained('google/gemma-2-2b-it', dtype='bfloat16')
ov_circuit = model.OV
eigs = ov_circuit.eigenvalues  # this would cause error

and we get

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[1], line 5
      3 model = HookedTransformer.from_pretrained('google/gemma-2-2b-it', dtype='bfloat16', device='cpu')
      4 ov_circuit = model.OV
----> 5 eigs = ov_circuit.eigenvalues  # this would cause error

File ~/github/exp_hallucination/.venv/lib/python3.11/site-packages/transformer_lens/FactoredMatrix.py:194, in FactoredMatrix.eigenvalues(self)
    191 @property
    192 def eigenvalues(self) -> Float[torch.Tensor, "*leading_dims mdim"]:
    193     """Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv, so Av is an eigenvector of AB with eigenvalue k."""
--> 194     return torch.linalg.eig(self.BA).eigenvalues

RuntimeError: "linalg_eig_out_cpu" not implemented for 'BFloat16'

Maybe we could convert self.BA to higher precision: torch.linalg.eig(self.BA.to(torch.float32)).eigenvalues

System Info Describe the characteristic of your environment:

• Describe how transformer_lens was installed (pip, docker, source, ...): uv
• What OS are you using? (Linux, MacOS, Windows): Ubuntu 20.04
• Python version (We support 3.7--3.10 currently): 3.11

Additional context No.

Checklist

• [x] I have checked that there is no similar issue in the repo (required)