Improve ActivationCache docs

.vscode/cspell.json

@@ -6,8 +6,10 @@
         "alonso",
         "arange",
         "argmax",
+        "autodiff",
         "autoregressive",
         "barez",
+        "Beartype",
         "belrose",
         "bertsimas",
         "biderman",
@@ -18,6 +20,7 @@
         "circuitsvis",
         "Codespaces",
         "colab",
+        "collectstart",
         "colour",
         "conmy",
         "cooney",
@@ -35,6 +38,7 @@
         "evals",
         "fazl",
         "firstpage",
+        "fspath",
         "furo",
         "garriga",
         "gelu",
@@ -49,6 +53,7 @@
         "interp",
         "interpretability",
         "ioannis",
+        "ipynb",
         "isort",
         "janiak",
         "jaxtyping",

docs/source/conf.py

@@ -39,7 +39,7 @@
 
 napoleon_include_init_with_doc = True
 napoleon_use_admonition_for_notes = True
-napoleon_custom_sections = ["Motivation:"]
+napoleon_custom_sections = ["Motivation:", "Warning:"]
 
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

docs/source/content/contributing.md

@@ -62,6 +62,10 @@ You should follow this order:
 
 A description of what the function/class does, including as much detail as is necessary to fully understand it.
 
+Warning:
+
+Any warnings to the user (e.g. common pitfalls).
+
 Examples:
 
 Include any examples here. They will be checked with doctest.
@@ -109,7 +113,7 @@ You can use LaTeX, but note that as you're placing this in python strings the ba
 must be repeated (i.e. `\\`). You can write LaTeX inline, or in "display mode".
 
 ```reStructuredText
-.. math:: (a + b)^2 = a^2 + 2ab + b^2
+:math:`(a + b)^2 = a^2 + 2ab + b^2`
 ```
 
 ```reStructuredText

makefile

@@ -15,10 +15,10 @@ acceptance-test:
 	poetry run pytest --cov=transformer_lens/ --cov-report=term-missing --cov-branch tests/acceptance
 
 docstring-test:
-	poetry run pytest transformer_lens/ --doctest-modules --doctest-plus
+	poetry run pytest transformer_lens/
 
 notebook-test:
-	poetry run pytest demos/Exploratory_Analysis_Demo.ipynb --nbval
+	poetry run pytest demos/Exploratory_Analysis_Demo.ipynb
 
 test:
 	make unit-test

pyproject.toml

@@ -87,7 +87,11 @@ filterwarnings = [
     # More info: https://numpy.org/doc/stable/reference/distutils.html#module-numpy.distutils
     "ignore:distutils Version classes are deprecated:DeprecationWarning"
 ]
-addopts = "--jaxtyping-packages=transformer_lens,beartype.beartype"
+addopts = """--jaxtyping-packages=transformer_lens,beartype.beartype \ 
+-W ignore::beartype.roar.BeartypeDecorHintPep585DeprecationWarning \ 
+--deselect transformer_lens/utils.py::test_prompt \
+--doctest-modules --doctest-plus \
+--nbval"""
 
 [tool.isort]
 profile = "black"