Improve testing

docs/index.html

@@ -23,11 +23,12 @@ <h1>Saev</h1>
       It also includes an interactive webapp for looking through a trained SAE's features.
     </p>
     <p>
-      <a href="saev">API reference docs</a> are available, as well as the <a href="https://github.com/samuelstevens/saev">source code on GitHub</a>.
+      API reference docs are available below, as well as the <a href="https://github.com/samuelstevens/saev">source code on GitHub</a>.
     <p>
     <p>
       <a href="https://github.com/samuelstevens/saev/blob/main/logbook.md">My logbook</a> is a set of notes that might also be useful.
     </p>
+    <h2>Package Docs</h2>
     <dl id="http-server-module-list">
       <div class="flex">
         <dt><a href="saev">saev</a></dt>

docs/saev/test_config.html

@@ -41,6 +41,12 @@ <h2 class="section-title" id="header-functions">Functions</h2>
 <dd>
 <div class="desc"></div>
 </dd>
+<dt id="saev.test_config.test_expand_multiple"><code class="name flex">
+<span>def <span class="ident">test_expand_multiple</span></span>(<span>)</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
 <dt id="saev.test_config.test_expand_nested"><code class="name flex">
 <span>def <span class="ident">test_expand_nested</span></span>(<span>)</span>
 </code></dt>
@@ -83,6 +89,7 @@ <h2 class="section-title" id="header-functions">Functions</h2>
 <li><h3><a href="#header-functions">Functions</a></h3>
 <ul class="">
 <li><code><a title="saev.test_config.test_expand" href="#saev.test_config.test_expand">test_expand</a></code></li>
+<li><code><a title="saev.test_config.test_expand_multiple" href="#saev.test_config.test_expand_multiple">test_expand_multiple</a></code></li>
 <li><code><a title="saev.test_config.test_expand_nested" href="#saev.test_config.test_expand_nested">test_expand_nested</a></code></li>
 <li><code><a title="saev.test_config.test_expand_nested_and_unnested" href="#saev.test_config.test_expand_nested_and_unnested">test_expand_nested_and_unnested</a></code></li>
 <li><code><a title="saev.test_config.test_expand_nested_and_unnested_backwards" href="#saev.test_config.test_expand_nested_and_unnested_backwards">test_expand_nested_and_unnested_backwards</a></code></li>

justfile

@@ -2,7 +2,7 @@ docs: lint
     uv run pdoc3 --force --html --output-dir docs --config latex_math=True saev probing
 
 test: lint
-    uv run pytest saev probing
+    uv run pytest --cov saev -n auto saev probing
 
 lint: fmt
     fd -e py | xargs ruff check

logbook.md

@@ -821,8 +821,16 @@ There are several changes
 * Removed b_dec re-init
 * Scaled mean activation norm to approximately sqrt(d_vit)
 * Subtracted approximate mean activation to center activations
-* Various code changes :(
+* Various code changes to the way the activations are recorded :(.
 
 So I will re-record ViT-B/16 activations, then re-add the b_dec re-init, ignore the scale_norm and scale_mean and *pray* that it works again.
 Then I will re-add the original changes, debugging what went wrong at each step.
 It might just be a learning rate thing.
+
+# 11/20/2024
+
+What exactly are my next steps?
+
+1. Reproduce original workflow with ViT-B/16. There's some bug introduced and I need to fix it. Add tests for it.
+2. Propose a trend in sparse autoencoders between CLIP and DINOv2 with at least 3 examples.
+3. Ensure that your probing methodology works.

pyproject.toml

@@ -33,6 +33,8 @@ dev-dependencies = [
     "hypothesis-torch>=0.8.4",
     "pdoc3>=0.11.1",
     "pytest>=8.3.3",
+    "pytest-xdist>=3.6.1",
+    "pytest-cov>=6.0.0",
 ]
 
 [tool.pytest.ini_options]

saev/test_activations.py

@@ -16,7 +16,7 @@ def test_dataloader_batches():
     cfg = config.Activations(
         model_ckpt="ViT-B-32/openai",
         d_vit=768,
-        n_layers=3,
+        layers=[-2, -1],
         n_patches_per_img=49,
         vit_batch_size=8,
     )
@@ -38,11 +38,9 @@ def test_shard_writer_and_dataset_e2e():
         cfg = config.Activations(
             model_org="timm",
             model_ckpt="hf_hub:timm/test_vit3.r160_in1k",
-            width=160,
-            height=160,
             d_vit=96,
             n_patches_per_img=100,
-            n_layers=2,
+            layers=[-2, -1],
             vit_batch_size=8,
             n_workers=8,
             dump_to=tmpdir,
@@ -52,7 +50,11 @@ def test_shard_writer_and_dataset_e2e():
         writer = activations.ShardWriter(cfg)
         dataset = activations.Dataset(
             config.DataLoad(
-                shard_root=activations.get_acts_dir(cfg), patches="cls", layer=-1
+                shard_root=activations.get_acts_dir(cfg),
+                patches="cls",
+                layer=-1,
+                scale_mean=False,
+                scale_norm=False,
             )
         )
 
@@ -62,8 +64,9 @@ def test_shard_writer_and_dataset_e2e():
             _, cache = vit(batch["image"])
             writer[i : i + len(cache)] = cache
             i += len(cache)
-            assert cache.shape == (cfg.vit_batch_size, 2, 101, 96)
-            torch.testing.assert_close(
-                cache[:, -1, 0], dataset[batch["index"].numpy()][0]
-            )
+            assert cache.shape == (cfg.vit_batch_size, len(cfg.layers), 101, 96)
+
+            acts, _, _ = zip(*[dataset[i.item()] for i in batch["index"]])
+            from_dataset = torch.stack(acts)
+            torch.testing.assert_close(cache[:, -1, 0], from_dataset)
             print(f"Batch {b} matched.")