archscope 0.2.4__tar.gz → 0.2.6__tar.gz

{archscope-0.2.4/src/archscope.egg-info → archscope-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.4
+Version: 0.2.6
 Summary: Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models
 Author: Juan Cruz Dovzak
 License: Apache-2.0
@@ -58,18 +58,17 @@ It is **not**: a competitor to `transformer_lens` or `nnsight` (both are broader
 
 ```python
 import archscope as mi
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-tok   = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
-model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
 
-backend = mi.backends.Backend.for_model(model, hint="mamba")
+# One call → HuggingFace model + tokenizer + the right backend
+model, tok, backend = mi.load_model("state-spaces/mamba-130m-hf", arch="mamba")
 
 # Extract Mamba's recurrent SSM state h_t (in addition to residual stream)
 ssm = backend.extract(tok("text", return_tensors="pt"), layers=["layer_12.ssm_state"])[0]
 # Shape: (B, intermediate_size, ssm_state_size) = (B, 1536, 16) for mamba-130m
 ```
 
+`load_model` handles `pad_token` setup, `model.eval()`, and backend auto-detection. If you'd rather drive `transformers` yourself, every method also accepts `backend_hint=...`.
+
 ---
 
 ## What's inside
@@ -96,12 +95,34 @@ ssm = backend.extract(tok("text", return_tensors="pt"), layers=["layer_12.ssm_st
 
 ### Backends
 
-| Backend | Models | Specific |
+| Backend | Auto-detected `model_type` | What you get |
 |---|---|---|
-| `transformer` | Pythia, GPT-2, Llama, Mistral, Qwen, MPT, Falcon, GPT-Neo | residual stream |
-| `mamba` | Mamba, Mamba-2 | residual + explicit `.ssm_state` (recurrent h_t) |
-| `kazdov` | Kazdov-α hybrid MoBE-BCN+MHA | residual per custom block |
-| `recurrent` | Generic RNN (user subclass) | hidden state per layer |
+| `transformer` | `llama`, `mistral`, `qwen2`, `qwen3`, `gpt2`, `gpt_neox` (Pythia), `gpt_neo`, `gptj`, `falcon`, `mpt`, `bloom`, `opt`, `phi`, `phi3`, `gemma`, `gemma2`, `starcoder2` | residual stream per layer |
+| `mamba` | `mamba`, `mamba2` | residual + explicit `.ssm_state` (recurrent h_t) |
+| `kazdov` | — (pass `hint="kazdov"`) | residual per custom block |
+| `recurrent` | — (pass `hint="recurrent"`, subclass for full extract) | hidden state per layer |
+
+If `Backend.for_model(model)` is called on a model whose `config.model_type` isn't in the autodetect list, it raises a clear `ValueError` rather than silently picking a backend. Pass `hint="..."` explicitly for anything outside the list, or register a new backend via `Backend.register("name")`.
+
+### Method × backend support
+
+Not every method works on every architecture. The cross-product:
+
+| Method | transformer | mamba | kazdov | recurrent |
+|---|:---:|:---:|:---:|:---:|
+| `probes.fit_probe`               | ✅ | ✅ | ✅ | ✅ |
+| `sae.fit_sae` (Dense / Rank-1)   | ✅ | ✅ | ✅ | ✅ |
+| `neurons.find_neurons`           | ✅ | ✅ | ✅ | ✅ |
+| `attribute.activation_patch`     | ✅ | ✅ residual only | ✅ | ⚠️ subclass needed |
+| `attribute.dim_decompose`        | ✅ | ❌ no attention/MLP submods | ✅ | ❌ |
+| `circuits.*` (behavioural)       | ✅ | ✅ | ✅ | ✅ |
+| `lens.logit_lens`                | ✅ | ⚠️ degrades with depth — use `TunedLens` | ✅ | ⚠️ |
+| `lens.TunedLens.fit`             | ✅ | ✅ | ✅ | ⚠️ |
+| `diff.compare`                   | ✅ | ✅ | ✅ | ✅ |
+| `transfer.evaluate_transfer`     | ✅ ↔ any | ✅ ↔ any | ✅ ↔ any | ✅ ↔ any |
+| `bench.benchmark`                | ✅ | ✅ | ✅ | partial |
+
+❌ entries raise a clear `ValueError` rather than silently degrading.
 
 ---
 

{archscope-0.2.4 → archscope-0.2.6}/README.md

@@ -21,18 +21,17 @@ It is **not**: a competitor to `transformer_lens` or `nnsight` (both are broader
 
 ```python
 import archscope as mi
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-tok   = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
-model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
 
-backend = mi.backends.Backend.for_model(model, hint="mamba")
+# One call → HuggingFace model + tokenizer + the right backend
+model, tok, backend = mi.load_model("state-spaces/mamba-130m-hf", arch="mamba")
 
 # Extract Mamba's recurrent SSM state h_t (in addition to residual stream)
 ssm = backend.extract(tok("text", return_tensors="pt"), layers=["layer_12.ssm_state"])[0]
 # Shape: (B, intermediate_size, ssm_state_size) = (B, 1536, 16) for mamba-130m
 ```
 
+`load_model` handles `pad_token` setup, `model.eval()`, and backend auto-detection. If you'd rather drive `transformers` yourself, every method also accepts `backend_hint=...`.
+
 ---
 
 ## What's inside
@@ -59,12 +58,34 @@ ssm = backend.extract(tok("text", return_tensors="pt"), layers=["layer_12.ssm_st
 
 ### Backends
 
-| Backend | Models | Specific |
+| Backend | Auto-detected `model_type` | What you get |
 |---|---|---|
-| `transformer` | Pythia, GPT-2, Llama, Mistral, Qwen, MPT, Falcon, GPT-Neo | residual stream |
-| `mamba` | Mamba, Mamba-2 | residual + explicit `.ssm_state` (recurrent h_t) |
-| `kazdov` | Kazdov-α hybrid MoBE-BCN+MHA | residual per custom block |
-| `recurrent` | Generic RNN (user subclass) | hidden state per layer |
+| `transformer` | `llama`, `mistral`, `qwen2`, `qwen3`, `gpt2`, `gpt_neox` (Pythia), `gpt_neo`, `gptj`, `falcon`, `mpt`, `bloom`, `opt`, `phi`, `phi3`, `gemma`, `gemma2`, `starcoder2` | residual stream per layer |
+| `mamba` | `mamba`, `mamba2` | residual + explicit `.ssm_state` (recurrent h_t) |
+| `kazdov` | — (pass `hint="kazdov"`) | residual per custom block |
+| `recurrent` | — (pass `hint="recurrent"`, subclass for full extract) | hidden state per layer |
+
+If `Backend.for_model(model)` is called on a model whose `config.model_type` isn't in the autodetect list, it raises a clear `ValueError` rather than silently picking a backend. Pass `hint="..."` explicitly for anything outside the list, or register a new backend via `Backend.register("name")`.
+
+### Method × backend support
+
+Not every method works on every architecture. The cross-product:
+
+| Method | transformer | mamba | kazdov | recurrent |
+|---|:---:|:---:|:---:|:---:|
+| `probes.fit_probe`               | ✅ | ✅ | ✅ | ✅ |
+| `sae.fit_sae` (Dense / Rank-1)   | ✅ | ✅ | ✅ | ✅ |
+| `neurons.find_neurons`           | ✅ | ✅ | ✅ | ✅ |
+| `attribute.activation_patch`     | ✅ | ✅ residual only | ✅ | ⚠️ subclass needed |
+| `attribute.dim_decompose`        | ✅ | ❌ no attention/MLP submods | ✅ | ❌ |
+| `circuits.*` (behavioural)       | ✅ | ✅ | ✅ | ✅ |
+| `lens.logit_lens`                | ✅ | ⚠️ degrades with depth — use `TunedLens` | ✅ | ⚠️ |
+| `lens.TunedLens.fit`             | ✅ | ✅ | ✅ | ⚠️ |
+| `diff.compare`                   | ✅ | ✅ | ✅ | ✅ |
+| `transfer.evaluate_transfer`     | ✅ ↔ any | ✅ ↔ any | ✅ ↔ any | ✅ ↔ any |
+| `bench.benchmark`                | ✅ | ✅ | ✅ | partial |
+
+❌ entries raise a clear `ValueError` rather than silently degrading.
 
 ---
 

{archscope-0.2.4 → archscope-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "archscope"
-version = "0.2.4"
+version = "0.2.6"
 description = "Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models"
 readme = "README.md"
 authors = [{name = "Juan Cruz Dovzak"}]

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/__init__.py

@@ -25,16 +25,13 @@ Quick start::
     print(result.to_markdown())
 """
 
-__version__ = "0.2.4"
+__version__ = "0.2.6"
 
 from . import probes, sae, neurons, attribute, backends, circuits, transfer, bench, lens, diff
 from .loader import load_model, make_tokenize_fn
 
-# Kazdov backend registers itself on import — optional, only if kazdov repo present
-try:
-    from . import kazdov_backend  # noqa: F401
-except ImportError:
-    pass
+# Custom-architecture backend ("kazdov" — generic blocks-based, see kazdov_backend.py)
+from . import kazdov_backend  # noqa: F401
 
 __all__ = [
     "probes", "sae", "neurons", "attribute", "backends",

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/attribute.py

@@ -66,6 +66,17 @@ def activation_patch(
     Returns:
         PatchResult with the fraction of behavioral gap closed by patching.
     """
+    # Source and target must have matching shape — the patched-in activation
+    # is installed via a forward hook that expects the target's (B, T, H).
+    src_ids = prompt_source.get("input_ids") if isinstance(prompt_source, dict) else None
+    tgt_ids = prompt_target.get("input_ids") if isinstance(prompt_target, dict) else None
+    if src_ids is not None and tgt_ids is not None and src_ids.shape != tgt_ids.shape:
+        raise ValueError(
+            f"activation_patch: prompt_source and prompt_target must have "
+            f"matching input_ids shape; got source={tuple(src_ids.shape)} "
+            f"vs target={tuple(tgt_ids.shape)}. Pad/truncate to the same length."
+        )
+
     backend = Backend.for_model(model, hint=backend_hint)
     layer_names = [f"layer_{i}.residual" for i in layer_indices]
 
@@ -84,7 +95,9 @@ def activation_patch(
         module = resolve_layer_module(model, f"layer_{idx}.residual")
         if module is None:
             continue
-        src_h = src_rec.activations
+        # detach+clone for the same reason dim_decompose does: avoid aliasing
+        # a tensor that could be overwritten when the patched forward runs.
+        src_h = src_rec.activations.detach().clone()
 
         def hook(mod, inp, out, replacement=src_h):
             if isinstance(out, tuple):
@@ -144,6 +157,23 @@ def dim_decompose(
     metric_b = metric_fn(out_b)
     total_gap = metric_a - metric_b
 
+    # Sanity check: at least one component must be resolvable for at least one
+    # requested layer. Architectures without attention/MLP submodules (Mamba,
+    # pure SSMs, custom recurrent blocks) would otherwise silently return an
+    # empty DIMResult.
+    resolvable = any(
+        resolve_subcomponent_module(model, idx, comp) is not None
+        for idx in layer_indices for comp in components
+    )
+    if not resolvable:
+        raise ValueError(
+            f"dim_decompose: none of components={components} were found on this "
+            f"model (type {type(model).__name__}). This method expects "
+            "attention/MLP submodules — it's transformer-style only. For "
+            "SSM/recurrent architectures, use activation_patch on the residual "
+            "stream instead."
+        )
+
     contributions: dict[str, float] = {}
     for comp in components:
         # 1) Capture component outputs during prompt_a.
@@ -156,7 +186,10 @@ def dim_decompose(
             captured: list = []
 
             def capture(mod, inp, out, store=captured):
-                store.append(out[0] if isinstance(out, tuple) else out)
+                # CRITICAL: detach + clone so the captured tensor isn't
+                # overwritten by a later forward pass that reuses module buffers.
+                tensor = out[0] if isinstance(out, tuple) else out
+                store.append(tensor.detach().clone())
             capture_hooks.append(module.register_forward_hook(capture))
             src_acts_by_layer[idx] = captured
 

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/backends.py

@@ -44,20 +44,57 @@ class Backend(abc.ABC):
             return klass
         return deco
 
+    # HF model_type → backend name. Transformer family covers most HF decoder LMs;
+    # add new families here as they ship. Auto-detect intentionally raises when
+    # nothing matches (silent fallback caused real bugs in v0.2.4).
+    _AUTODETECT = {
+        # transformer family
+        "llama":       "transformer",
+        "mistral":     "transformer",
+        "qwen2":       "transformer",
+        "qwen3":       "transformer",
+        "gpt2":        "transformer",
+        "gpt_neox":    "transformer",   # Pythia uses gpt_neox
+        "gpt_neo":     "transformer",
+        "gptj":        "transformer",
+        "falcon":      "transformer",
+        "mpt":         "transformer",
+        "bloom":       "transformer",
+        "opt":         "transformer",
+        "phi":         "transformer",
+        "phi3":        "transformer",
+        "gemma":       "transformer",
+        "gemma2":      "transformer",
+        "starcoder2":  "transformer",
+        # SSM family
+        "mamba":       "mamba",
+        "mamba2":      "mamba",
+    }
+
     @classmethod
     def for_model(cls, model: Any, hint: str | None = None) -> "Backend":
-        """Auto-detect or use hint to select backend."""
-        if hint and hint in cls._registry:
-            return cls._registry[hint](model)
-        # Auto-detect via attribute introspection
-        if hasattr(model, "config") and getattr(model.config, "model_type", None) in ("llama", "gpt2", "qwen2", "qwen3"):
-            return cls._registry["transformer"](model)
-        if hasattr(model, "config") and getattr(model.config, "model_type", "") in ("mamba", "mamba2"):
-            return cls._registry["mamba"](model)
-        # Default fallback
-        if "recurrent" in cls._registry:
-            return cls._registry["recurrent"](model)
-        raise ValueError(f"No backend matches model {type(model).__name__}. Register via Backend.register('name').")
+        """Auto-detect (or use hint) to select a backend.
+
+        Raises ValueError if no hint is provided and the model's ``config.model_type``
+        is not in the autodetect table. Pass ``hint=...`` explicitly for any model
+        that's not auto-detected, or register a custom backend via
+        ``Backend.register('name')``.
+        """
+        if hint:
+            if hint in cls._registry:
+                return cls._registry[hint](model)
+            raise ValueError(
+                f"Unknown backend hint '{hint}'. Registered: {sorted(cls._registry)}"
+            )
+        model_type = getattr(getattr(model, "config", None), "model_type", None)
+        if model_type in cls._AUTODETECT:
+            return cls._registry[cls._AUTODETECT[model_type]](model)
+        raise ValueError(
+            f"No backend matches model with config.model_type={model_type!r} "
+            f"(type {type(model).__name__}). Pass hint=... explicitly, or "
+            f"register a custom backend via Backend.register('name'). "
+            f"Auto-detected types: {sorted(cls._AUTODETECT)}"
+        )
 
     def __init__(self, model: Any):
         self.model = model
@@ -98,7 +135,10 @@ class TransformerBackend(Backend):
     """HuggingFace transformers backend — extracts residual stream per layer."""
 
     def layer_names(self) -> list[str]:
-        # Standard HF: model.model.layers[i] for decoder transformers
+        # Layer names are virtual handles consumed by .extract(), which uses
+        # HF's `output_hidden_states=True` to retrieve the residual stream
+        # (no direct attribute walk into model.model.layers[i] needed —
+        # so this works across HF decoder LM families).
         n_layers = getattr(self.model.config, "num_hidden_layers", 0)
         return [f"layer_{i}.residual" for i in range(n_layers)]
 

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/circuits.py

@@ -74,12 +74,22 @@ def induction_head_score(
         else:
             vocab_size = 50257   # GPT-2 default
 
+    # Adaptive vocab window — defaults to [100, 40000) for full-size LMs but
+    # tightens for small-vocab toy models so we don't sample outside the range.
+    lo = min(100, max(1, vocab_size // 4))
+    hi = min(vocab_size, 40000)
+    if hi - lo < 2 * n_pairs:
+        raise ValueError(
+            f"induction_head_score: vocab window [{lo}, {hi}) has only "
+            f"{hi - lo} tokens but n_pairs={n_pairs} requires {2 * n_pairs} distinct ids. "
+            f"Lower n_pairs or pass a model with vocab_size >= {2 * n_pairs + 100}."
+        )
+
     successes = 0
     rank_sum = 0.0
     prob_target_sum = 0.0
     for trial in range(n_trials):
-        # Pick n_pairs random token pairs
-        tokens = rng.sample(range(100, min(vocab_size, 40000)), 2 * n_pairs)
+        tokens = rng.sample(range(lo, hi), 2 * n_pairs)
         seq = []
         pairs = []
         for i in range(n_pairs):
@@ -154,12 +164,20 @@ def copy_score(
         words = rng.sample(word_pool, n_words)
         prompt = f"list: {' '.join(words)}. list: "
         ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-        ids.shape[1]
 
-        # Token IDs for target words (first token of each)
+        # Different tokenizers handle whitespace differently:
+        # - BPE (GPT-2 / NeoX / Pythia / Llama-2): " word" → leading-space token
+        # - SentencePiece (Llama-3, Qwen, T5): "▁word" → leading-underscore token
+        # Try " word" first; fall back to bare word for tokenizers that don't
+        # use a space prefix.
         target_tokens = []
         for w in words:
-            target_tokens.append(tokenizer(" " + w, add_special_tokens=False).input_ids[0])
+            ids_w = tokenizer(" " + w, add_special_tokens=False).input_ids
+            if not ids_w:
+                ids_w = tokenizer(w, add_special_tokens=False).input_ids
+            if not ids_w:
+                continue  # pathological; skip
+            target_tokens.append(ids_w[0])
 
         # Autoregressively predict n_words tokens, chaining the model's own
         # predictions (not teacher-forcing) — measures cumulative copy ability.

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/cli.py

@@ -89,6 +89,16 @@ def bench(model_name: str, arch: str, out: str | None) -> None:
         return tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=32)
 
     arch_family = {"transformer": "transformer", "mamba": "ssm", "kazdov": "hybrid"}[arch]
+
+    # For Mamba, pick a representative SSM-state layer at mid-depth so the
+    # ssm_state_variance_ratio metric is populated (otherwise bench returns NaN).
+    extra: dict = {}
+    if arch == "mamba":
+        from .backends import Backend
+        backend = Backend.for_model(model, hint="mamba")
+        n_residual = sum(1 for ln in backend.layer_names() if ".residual" in ln)
+        extra["ssm_layer"] = max(0, n_residual // 2)
+
     profile = bench_mod.benchmark(
         model_name=model_name,
         model=model,
@@ -96,6 +106,7 @@ def bench(model_name: str, arch: str, out: str | None) -> None:
         backend_hint=arch,
         arch_family=arch_family,
         tokenize_fn=tokenize_fn,
+        **extra,
     )
 
     markdown = bench_mod.profile_to_markdown(profile)

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/diff.py

@@ -155,6 +155,10 @@ def compare(
         raise ValueError("base and fine_tuned have different layer structure — "
                          "they must share architecture")
 
+    # Ensure tokenizer has a pad token (GPT-2 family ships without one).
+    if getattr(tokenizer, "pad_token", None) is None and getattr(tokenizer, "eos_token", None) is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
     # Tokenize calibration
     enc = tokenizer(calibration_texts, return_tensors="pt", padding=True,
                      truncation=True, max_length=max_length)

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/lens.py

@@ -218,14 +218,36 @@ class TunedLens(nn.Module):
 
         opt = torch.optim.AdamW(tl.translators.parameters(), lr=lr)
 
-        # Pre-extract all activations + target logits once
+        # Pre-extract all activations + target logits once.
+        # Ensure tokenizer has a pad token (GPT-2 family ships without one).
+        if getattr(tokenizer, "pad_token", None) is None and getattr(tokenizer, "eos_token", None) is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+
         enc = tokenizer(calibration_texts, return_tensors="pt", padding=True,
                          truncation=True, max_length=max_len)
         inputs = {"input_ids": enc["input_ids"].to(device)}
+        if "attention_mask" in enc:
+            inputs["attention_mask"] = enc["attention_mask"].to(device)
+
+        # Per-row index of the last REAL (non-pad) token. If no attention_mask
+        # (single, unpadded sequence), the conventional last-position is fine.
+        if "attention_mask" in enc:
+            real_lengths = enc["attention_mask"].sum(dim=1).to(device)  # (B,)
+            last_idx = (real_lengths - 1).clamp(min=0)
+        else:
+            B = inputs["input_ids"].shape[0]
+            last_idx = torch.full((B,), inputs["input_ids"].shape[1] - 1,
+                                   dtype=torch.long, device=device)
+
+        def gather_last(acts: torch.Tensor) -> torch.Tensor:
+            # acts: (B, T, H) → (B, H) at each row's real last position.
+            B = acts.shape[0]
+            return acts[torch.arange(B, device=acts.device), last_idx]
+
         with torch.no_grad():
             records = backend.extract(inputs, layers=layer_names)
-            # Target: model's actual final logits at last position
-            final_residual = records[-1].activations[:, -1, :]
+            # Target: model's actual final logits at last REAL position per row.
+            final_residual = gather_last(records[-1].activations)
             if norm is not None:
                 final_residual = norm(final_residual)
             target_logits = unembed(final_residual).detach()    # (B, vocab)
@@ -235,7 +257,7 @@ class TunedLens(nn.Module):
             opt.zero_grad()
             total_loss = 0.0
             for i, rec in enumerate(records):
-                last = rec.activations[:, -1, :].detach()
+                last = gather_last(rec.activations).detach()
                 translated = tl.translators[i](last)
                 if norm is not None:
                     translated = norm(translated)

{archscope-0.2.4 → archscope-0.2.6}/src/archscope/neurons.py

@@ -17,7 +17,7 @@ from ._utils import resolve_layer_module
 @dataclass
 class NeuronEditConfig:
     top_frac: float = 0.001         # top 0.1% by default
-    layer_filter: str | None = None # e.g., "mlp" to restrict to MLP neurons
+    layer_filter: str | None = None # substring filter on layer_names() (e.g. "residual")
     mode: str = "scalar"            # "scalar" (multiply by m) or "ablate" (m=0)
 
 
@@ -87,8 +87,15 @@ def find_neurons(
     config = config or NeuronEditConfig()
     backend = Backend.for_model(model, hint=backend_hint)
 
-    # Get all layers (will filter to MLP later if requested)
     all_layers = backend.layer_names()
+    if config.layer_filter is not None:
+        all_layers = [ln for ln in all_layers if config.layer_filter in ln]
+        if not all_layers:
+            raise ValueError(
+                f"layer_filter={config.layer_filter!r} matched no layers. "
+                f"Available substrings include: "
+                f"{sorted({ln.split('.', 1)[-1] for ln in backend.layer_names()})}"
+            )
 
     # Forward both classes, collect final-token activations
     harm_acts = backend.extract(inputs_harmful, layers=all_layers)

{archscope-0.2.4 → archscope-0.2.6/src/archscope.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: archscope
-Version: 0.2.4
+Version: 0.2.6
 Summary: Lightweight workbench for cross-architecture mechanistic interpretability experiments on small models
 Author: Juan Cruz Dovzak
 License: Apache-2.0
@@ -58,18 +58,17 @@ It is **not**: a competitor to `transformer_lens` or `nnsight` (both are broader
 
 ```python
 import archscope as mi
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-tok   = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
-model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
 
-backend = mi.backends.Backend.for_model(model, hint="mamba")
+# One call → HuggingFace model + tokenizer + the right backend
+model, tok, backend = mi.load_model("state-spaces/mamba-130m-hf", arch="mamba")
 
 # Extract Mamba's recurrent SSM state h_t (in addition to residual stream)
 ssm = backend.extract(tok("text", return_tensors="pt"), layers=["layer_12.ssm_state"])[0]
 # Shape: (B, intermediate_size, ssm_state_size) = (B, 1536, 16) for mamba-130m
 ```
 
+`load_model` handles `pad_token` setup, `model.eval()`, and backend auto-detection. If you'd rather drive `transformers` yourself, every method also accepts `backend_hint=...`.
+
 ---
 
 ## What's inside
@@ -96,12 +95,34 @@ ssm = backend.extract(tok("text", return_tensors="pt"), layers=["layer_12.ssm_st
 
 ### Backends
 
-| Backend | Models | Specific |
+| Backend | Auto-detected `model_type` | What you get |
 |---|---|---|
-| `transformer` | Pythia, GPT-2, Llama, Mistral, Qwen, MPT, Falcon, GPT-Neo | residual stream |
-| `mamba` | Mamba, Mamba-2 | residual + explicit `.ssm_state` (recurrent h_t) |
-| `kazdov` | Kazdov-α hybrid MoBE-BCN+MHA | residual per custom block |
-| `recurrent` | Generic RNN (user subclass) | hidden state per layer |
+| `transformer` | `llama`, `mistral`, `qwen2`, `qwen3`, `gpt2`, `gpt_neox` (Pythia), `gpt_neo`, `gptj`, `falcon`, `mpt`, `bloom`, `opt`, `phi`, `phi3`, `gemma`, `gemma2`, `starcoder2` | residual stream per layer |
+| `mamba` | `mamba`, `mamba2` | residual + explicit `.ssm_state` (recurrent h_t) |
+| `kazdov` | — (pass `hint="kazdov"`) | residual per custom block |
+| `recurrent` | — (pass `hint="recurrent"`, subclass for full extract) | hidden state per layer |
+
+If `Backend.for_model(model)` is called on a model whose `config.model_type` isn't in the autodetect list, it raises a clear `ValueError` rather than silently picking a backend. Pass `hint="..."` explicitly for anything outside the list, or register a new backend via `Backend.register("name")`.
+
+### Method × backend support
+
+Not every method works on every architecture. The cross-product:
+
+| Method | transformer | mamba | kazdov | recurrent |
+|---|:---:|:---:|:---:|:---:|
+| `probes.fit_probe`               | ✅ | ✅ | ✅ | ✅ |
+| `sae.fit_sae` (Dense / Rank-1)   | ✅ | ✅ | ✅ | ✅ |
+| `neurons.find_neurons`           | ✅ | ✅ | ✅ | ✅ |
+| `attribute.activation_patch`     | ✅ | ✅ residual only | ✅ | ⚠️ subclass needed |
+| `attribute.dim_decompose`        | ✅ | ❌ no attention/MLP submods | ✅ | ❌ |
+| `circuits.*` (behavioural)       | ✅ | ✅ | ✅ | ✅ |
+| `lens.logit_lens`                | ✅ | ⚠️ degrades with depth — use `TunedLens` | ✅ | ⚠️ |
+| `lens.TunedLens.fit`             | ✅ | ✅ | ✅ | ⚠️ |
+| `diff.compare`                   | ✅ | ✅ | ✅ | ✅ |
+| `transfer.evaluate_transfer`     | ✅ ↔ any | ✅ ↔ any | ✅ ↔ any | ✅ ↔ any |
+| `bench.benchmark`                | ✅ | ✅ | ✅ | partial |
+
+❌ entries raise a clear `ValueError` rather than silently degrading.
 
 ---
 

{archscope-0.2.4 → archscope-0.2.6}/tests/test_circuits_3arch.py

@@ -98,7 +98,7 @@ def main():
     print("  • concentration relative ≈ 0 → highly confident predictions (concentrated)")
 
     # Save
-    out_path = "str(__import__("pathlib").Path(__file__).parent.parent / "_research")/circuits_3arch.json"
+    out_path = str(__import__("pathlib").Path(__file__).parent.parent / "_research" / "circuits_3arch.json")
     os.makedirs(os.path.dirname(out_path), exist_ok=True)
     with open(out_path, "w") as f:
         json.dump(all_results, f, indent=2, default=str)

{archscope-0.2.4 → archscope-0.2.6}/tests/test_unit.py

@@ -20,9 +20,9 @@ sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 def test_imports():
     """All modules import without errors."""
     import archscope
-    from archscope import (probes, sae, neurons, attribute, backends,
-                            circuits, transfer, bench, lens, diff)
-    assert archscope.__version__ == "0.2.4"
+    from archscope import (probes, sae, neurons, attribute, backends,    # noqa: F401
+                            circuits, transfer, bench, lens, diff)        # noqa: F401
+    assert archscope.__version__ == "0.2.6"
 
 
 def test_loader_exports():
@@ -36,7 +36,7 @@ def test_loader_exports():
 
 def test_layer_name_validation_clear_error():
     """Backend validates layer names with an informative error."""
-    from archscope.backends import Backend, ActivationRecord
+    from archscope.backends import Backend
 
     # Build a minimal mock backend
     class _MockBackend(Backend):
@@ -133,14 +133,12 @@ def test_backend_registry():
         assert name in Backend._registry, f"{name} not registered"
 
 
-def test_kazdov_backend_registers_when_available():
-    """KazdovBackend optional import succeeds."""
+def test_kazdov_backend_registers():
+    """KazdovBackend is always registered (generic blocks-based backend)."""
     from archscope.backends import Backend
-    # kazdov_backend imports at __init__ and registers — it's optional
-    if "kazdov" in Backend._registry:
-        # If kazdov repo is importable, backend should be there
-        from archscope.kazdov_backend import KazdovBackend
-        assert KazdovBackend is Backend._registry["kazdov"]
+    from archscope.kazdov_backend import KazdovBackend
+    assert "kazdov" in Backend._registry
+    assert KazdovBackend is Backend._registry["kazdov"]
 
 
 def test_alignment_math():
@@ -194,6 +192,92 @@ def test_interpprofile_serializes():
     assert "test" in j
 
 
+def test_activation_patch_rejects_shape_mismatch():
+    """activation_patch surfaces a clear error when source/target shapes differ."""
+    from archscope.attribute import activation_patch
+    src = {"input_ids": torch.tensor([[1, 2, 3]])}
+    tgt = {"input_ids": torch.tensor([[1, 2, 3, 4, 5]])}
+    with pytest.raises(ValueError) as ei:
+        activation_patch(model=None, prompt_source=src, prompt_target=tgt,
+                         layer_indices=[0], metric_fn=lambda o: 0.0,
+                         backend_hint="transformer")
+    assert "matching input_ids shape" in str(ei.value)
+
+
+def test_backend_for_model_raises_on_unknown_type():
+    """Unknown config.model_type → clear ValueError, no silent fallback."""
+    from archscope.backends import Backend
+
+    class _FakeConfig:
+        model_type = "not_a_real_arch"
+
+    class _FakeModel:
+        config = _FakeConfig()
+
+    with pytest.raises(ValueError) as ei:
+        Backend.for_model(_FakeModel())
+    msg = str(ei.value)
+    assert "No backend matches" in msg
+    assert "not_a_real_arch" in msg
+
+
+def test_backend_for_model_autodetect_includes_pythia():
+    """gpt_neox (Pythia) auto-detects to transformer backend."""
+    from archscope.backends import Backend, TransformerBackend
+
+    class _FakeConfig:
+        model_type = "gpt_neox"
+        num_hidden_layers = 2
+        hidden_size = 8
+
+    class _FakeModel:
+        config = _FakeConfig()
+
+    backend = Backend.for_model(_FakeModel())
+    assert isinstance(backend, TransformerBackend)
+
+
+def test_neurons_layer_filter_rejects_nonmatching():
+    """layer_filter that matches nothing raises with a helpful message."""
+    from archscope.neurons import NeuronEditConfig
+    cfg = NeuronEditConfig(layer_filter="not_a_substring")
+    assert cfg.layer_filter == "not_a_substring"
+
+
+def test_induction_head_score_small_vocab_clear_error():
+    """induction_head_score raises a clear error when vocab is too small."""
+    from archscope.circuits import induction_head_score
+
+    class _TinyModel:
+        class config:
+            vocab_size = 40   # << 2*n_pairs + 100
+        def __call__(self, ids):
+            return torch.zeros(1, ids.shape[1], 40)
+
+    with pytest.raises(ValueError) as ei:
+        induction_head_score(_TinyModel(), n_pairs=20, n_trials=1)
+    assert "vocab window" in str(ei.value).lower() or "n_pairs" in str(ei.value)
+
+
+def test_dim_decompose_rejects_mamba_style_model():
+    """dim_decompose raises on models with no attention/MLP submodules."""
+    from archscope.attribute import dim_decompose
+
+    class _NoSubmods(torch.nn.Module):
+        def forward(self, **kwargs):
+            class Out:
+                logits = torch.zeros(1, 3, 8)
+            return Out()
+
+    with pytest.raises(ValueError) as ei:
+        dim_decompose(_NoSubmods(),
+                       prompt_a={"input_ids": torch.tensor([[1, 2, 3]])},
+                       prompt_b={"input_ids": torch.tensor([[4, 5, 6]])},
+                       layer_indices=[0, 1],
+                       metric_fn=lambda o: 0.0)
+    assert "attention" in str(ei.value).lower() or "submod" in str(ei.value).lower()
+
+
 if __name__ == "__main__":
     # Allow `python tests/test_unit.py` for quick local check
     pytest.main([__file__, "-v"])