Add attention_mask argument to loss_fn() and lm_cross_entropy_loss() and adjust the cross entropy calculation to ignore masked (padding) tokens.

Author: UFO-101

pyproject.toml

@@ -3,7 +3,7 @@
     description="An implementation of transformers tailored for mechanistic interpretability."
     license="MIT"
     name="transformer-lens"
-    packages=[{include="transformer_lens"}]
+    packages=[{include="transformer_lens"}, {include="transformer_lens/py.typed"}]
     readme="README.md"
     # Version is automatically set by the pipeline on release
     version="0.0.0"