(v3) Draft PR: add Pyright static typing to hook_points.py #590 (#607)

* Initial Commit (add pyright + test by adding few annotations)

* Slightly more typing added

* more typing

* Additional typing

* Completed typing for hook_points.py file

* todo clarifications

* formatting changes to hook_points.py

* Apply some suggestions from code review

Co-authored-by: Alan <[email protected]>

* Added typing for Literals and changed some assertions to if statements

* formatting

* update to accout for merged code

* small typing issue

* changing hookfunction protocol + more assertions

* change the slice input

* change from isinstance to callable checks

* fix: Update Gemma to reflect upstream HF changes (#596)

* update activation function to tanh approximation

* keep RMSNorm calcs in float32 and match cfg dtype for embedding scaling

* formatting

* keep mypy happy

* formatting

* allow user to force trust_remote_code=true via from_pretrained kwargs (#597)

* change + revert HookFunctionProtocol

* format

* module_output is now just a tensor

* set module ouput to be any type

---------

Co-authored-by: Alan <[email protected]>
Co-authored-by: Bryce Meyer <[email protected]>
Co-authored-by: cmathw <[email protected]>
Co-authored-by: Clement Dumas <[email protected]>

pyproject.toml

@@ -124,3 +124,68 @@
 [build-system]
     build-backend="poetry.core.masonry.api"
     requires=["poetry-core"]
+
+[tool.pyright]
+    # All rules apart from base are shown explicitly below
+    deprecateTypingAliases=true
+    disableBytesTypePromotions=true
+    exclude = [
+        "*/**/*.py",
+        "!/transformer_lens/hook_points.py"
+    ]
+    reportAssertAlwaysTrue=true
+    reportConstantRedefinition=true
+    reportDeprecated=true
+    reportDuplicateImport=true
+    reportFunctionMemberAccess=true
+    reportGeneralTypeIssues=true
+    reportIncompatibleMethodOverride=true
+    reportIncompatibleVariableOverride=true
+    reportIncompleteStub=true
+    reportInconsistentConstructor=true
+    reportInvalidStringEscapeSequence=true
+    reportInvalidStubStatement=true
+    reportInvalidTypeVarUse=true
+    reportMatchNotExhaustive=true
+    reportMissingParameterType=true
+    reportMissingTypeArgument=false
+    reportMissingTypeStubs=false
+    reportOptionalCall=true
+    reportOptionalContextManager=true
+    reportOptionalIterable=true
+    reportOptionalMemberAccess=true
+    reportOptionalOperand=true
+    reportOptionalSubscript=true
+    reportOverlappingOverload=true
+    reportPrivateImportUsage=true
+    reportPrivateUsage=true
+    reportSelfClsParameterName=true
+    reportTypeCommentUsage=true
+    reportTypedDictNotRequiredAccess=true
+    reportUnboundVariable=true
+    reportUnknownArgumentType=false
+    reportUnknownLambdaType=true
+    reportUnknownMemberType=false
+    reportUnknownParameterType=false
+    reportUnknownVariableType=false
+    reportUnnecessaryCast=true
+    reportUnnecessaryComparison=true
+    reportUnnecessaryContains=true
+    reportUnnecessaryIsInstance=true
+    reportUnsupportedDunderAll=true
+    reportUntypedBaseClass=true
+    reportUntypedClassDecorator=true
+    reportUntypedFunctionDecorator=true
+    reportUntypedNamedTuple=true
+    reportUnusedClass=true
+    reportUnusedCoroutine=true
+    reportUnusedExpression=true
+    reportUnusedFunction=true
+    reportUnusedImport=true
+    reportUnusedVariable=true
+    reportWildcardImportFromLibrary=true
+    strictDictionaryInference=true
+    strictListInference=true
+    strictParameterNoneValue=true
+    strictSetInference=true
+