Fix type checking (mypy)

[alan-cooney]

Currently mypy type checks are not passing - with a large number of easy to resolve errors.

poetry run mypy transformer_lens

Resolving these would help make the codebase more robust, and the corresponding checks in /.github/workflows/checks.yml can then be enabled (they're commented out).

https://github.com/neelnanda-io/TransformerLens/blob/72b9be00c12b5be37baaa16b57e5df5e2aff3aad/.github/workflows/checks.yml#L58

[derpyplops]

I tried a bunch to fix it with stub packages but mainly didn't really get anywhere

[neelnanda-io]

Hmm, when I run mypy, most of what I see seem to be it eg failing to deal with torchtyping?
Eg:

transformer_lens/components.py:16: error: Skipping analyzing "fancy_einsum": module is installed, but missing library stubs or py.typed marker  [import]             
transformer_lens/components.py:23: error: Skipping analyzing "torchtyping": module is installed, but missing library stubs or py.typed marker  [import]              
transformer_lens/components.py:37: error: Name "batch" is not defined  [name-defined]                                                                                
transformer_lens/components.py:37: error: Name "position" is not defined  [name-defined]                                                                             
transformer_lens/components.py:38: error: Name "batch" is not defined  [name-defined]                                                                                
transformer_lens/components.py:38: error: Name "position" is not defined  [name-defined]                                                                             
transformer_lens/components.py:38: error: Name "d_model" is not defined  [name-defined]                                                                              
transformer_lens/components.py:57: error: Name "batch" is not defined  [name-defined]                                                                                
transformer_lens/components.py:57: error: Name "position" is not defined  [name-defined]                                                                             
transformer_lens/components.py:57: error: Name "d_model" is not defined  [name-defined]                                                                              
transformer_lens/components.py:58: error: Name "batch" is not defined  [name-defined]                                                                                
transformer_lens/components.py:58: error: Name "position" is not defined  [name-defined]                                                                             
transformer_lens/components.py:58: error: Name "d_vocab_out" is not defined  [name-defined]                                                                          
transformer_lens/components.py:79: error: Name "batch" is not defined  [name-defined]                                                                                
transformer_lens/components.py:79: error: Name "position" is not defined  [name-defined]                                                                             
transformer_lens/components.py:80: error: Name "batch" is not defined  [name-defined]                                                                                
transformer_lens/components.py:80: error: Name "position" is not defined  [name-defined]                                                                             
transformer_lens/components.py:80: error: Name "d_model" is not defined  [name-defined]                                                                              
transformer_lens/components.py:118: error: Name "batch" is not defined  [name-defined]                                                                               
transformer_lens/components.py:118: error: Name "position" is not defined  [name-defined]                                                                            
transformer_lens/components.py:118: error: Name "length" is not defined  [name-defined]                                                                              
transformer_lens/components.py:119: error: Name "batch" is not defined  [name-defined]                                                                               
transformer_lens/components.py:119: error: Name "position" is not defined  [name-defined]
transformer_lens/components.py:119: error: Name "length" is not defined  [name-defined]

[dkamm]

Hi guys - just wanted to add that it looks like this might be expected because of mypy limitations (see patrick-kidger/torchtyping#35)

Patrick's first suggestion of defining an object with the strings might be the way to go. It could help solidify the naming conventions and avoid potential confusion like using "num_components" vs "components", etc.

I tried this in components.py and it got rid of the errors

from enum import Enum
class T(Enum):
    BATCH = "batch"
    D_VOCAB = "d_vocab"
    POSITION = "position"
    D_MODEL = "d_model"

# Embed & Unembed
class Embed(nn.Module):

    ...

    def forward(
        self, tokens: TT[T.BATCH, T.POSITION]
    ) -> TT[T.BATCH, T.POSITION, T.D_MODEL]:
        # If A has shape [a, b] and B has shape [c, d], then A[:, B] has shape [a, c, d]
        # B acts as a tensor of indices into the second dimension (so >=0 and <b)
        return self.W_E[tokens, :]

[alan-cooney]

@dkamm great idea! This is a far better approach

[neelnanda-io]

Yeah, that looks great @dkamm, I'd love for you to add that if you're interested.

…

On Mon, 16 Jan 2023 at 10:18, Alan ***@***.***> wrote: @dkamm <https://github.com/dkamm> great idea! This is a far better approach — Reply to this email directly, view it on GitHub <#130 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ASRPNKLNF33XI3KWMTVTWTLWSUN6TANCNFSM6AAAAAATJU6GPA> . You are receiving this because you commented.Message ID: ***@***.***>

[dkamm]

Sure! Will send a PR