TransformerLensOrg · jlarson4 · Feb 17, 2026 · Jan 26, 2026 · Jan 26, 2026 · Jan 26, 2026
diff --git a/demos/LIT_Integration_Demo.ipynb b/demos/LIT_Integration_Demo.ipynb
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -41,6 +41,11 @@
         typeguard="^4.2"
         typing-extensions="*"
         wandb=">=0.13.5"
+        # Optional dependencies
+        lit-nlp = {version = "^1.3", optional = true, python = ">=3.9"}
+
+    [tool.poetry.extras]
+        lit = ["lit-nlp"]
 
     [tool.poetry.group]
         [tool.poetry.group.dev.dependencies]

diff --git a/tests/acceptance/test_hooked_encoder_decoder.py b/tests/acceptance/test_hooked_encoder_decoder.py
@@ -173,7 +173,7 @@ def test_decoder_attention(our_model, huggingface_model, hello_world_tokens):
     input_len = hello_world_tokens.shape[1]
     cache_position = torch.arange(input_len)
     huggingface_attn_out = huggingface_attn(embed_out, cache_position=cache_position)[0]
-    assert_close(our_attn_out, huggingface_attn_out, rtol=3e-4, atol=1e-5)
+    assert_close(our_attn_out, huggingface_attn_out, rtol=5e-4, atol=1e-5)
 
 
 def test_attention_layer(our_model, huggingface_model, hello_world_tokens):