Fix FactoredMatrix bug

[callummcdougall]

Description

The function _convert_to_slice was added recently, which causes indexing with elements to work (previously this didn't work). But when this change was made, it caused indexing with sequences to work.

Type of change

The solution is to make _convert_to_slice only alter integer elements of the index.

To further explain, _convert_to_slice will take an integer like 1 and convert it into the slice slice(1, 2). This is good because now the indexed tensors A and B have the correct number of dimensions (i.e. they don't lose a dimension). But this isn't how we should treat things like t.tensor([0, 1, 2]) when they're used as an index, because they're already in the correct form for indexing (i.e. indexing with them won't delete a dimension).

I've now added tests to tests/unit/factored_matrix/test_get_item.py which would have failed on both the version of TransformerLens before the recent change to FactoredMatrix, and the version after this change, but which now all pass.

Screenshots

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

[jbloomAus]

This was causing issues in ARENA tutorials so expedited review. I'm planning on reviewing other PR's / issues this weekend.