Add CircuitLens feature-circuit attribution - #131
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #131 +/- ##
==========================================
- Coverage 98.36% 97.56% -0.81%
==========================================
Files 53 54 +1
Lines 4091 4305 +214
==========================================
+ Hits 4024 4200 +176
- Misses 67 105 +38 ☔ View full report in Codecov by Harness. |
There was a problem hiding this comment.
2 issues found across 6 files
Confidence score: 3/5
- In
src/tdhook/attribution/circuit_lens.py, mismatchedpatternandvalueshead selections can combine data from different heads while reporting the pattern index, producing incorrect attribution results; require matching head selections before positional decomposition. - In
src/tdhook/attribution/circuit_lens.py,target_position=-1attributes the final query but returns an invalid negative token index, which can mislead downstream consumers; reject negative positions consistently withAttentionSitevalidation.
You’re at about 99% of the monthly reviewed-line limit. You may want to disable incremental reviews to conserve quota. Reviews will continue until that limit is exceeded. If you need help avoiding interruptions, please contact contact@cubic.dev.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="src/tdhook/attribution/circuit_lens.py">
<violation number="1" location="src/tdhook/attribution/circuit_lens.py:202">
P3: When callers use `target_position=-1`, the helper attributes the final query while returning an invalid negative token index. Reject negative positions here, matching `AttentionSite` validation.</violation>
<violation number="2" location="src/tdhook/attribution/circuit_lens.py:316">
P2: When `pattern` and `values` select different head indices, this pairs one head’s pattern with another head’s values and reports the pattern index. Require matching head selections before positional decomposition.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| _captured(gradient, "attention output gradient"), | ||
| layer=site.layer, | ||
| target_position=site.target_position, | ||
| head_indices=site.pattern.indices, |
There was a problem hiding this comment.
P2: When pattern and values select different head indices, this pairs one head’s pattern with another head’s values and reports the pattern index. Require matching head selections before positional decomposition.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/tdhook/attribution/circuit_lens.py, line 316:
<comment>When `pattern` and `values` select different head indices, this pairs one head’s pattern with another head’s values and reports the pattern index. Require matching head selections before positional decomposition.</comment>
<file context>
@@ -0,0 +1,439 @@
+ _captured(gradient, "attention output gradient"),
+ layer=site.layer,
+ target_position=site.target_position,
+ head_indices=site.pattern.indices,
+ )
+ )
</file context>
| raise ValueError("output_weight head and head_dim axes do not match values") | ||
| if output_gradient.shape != (queries, output_weight.shape[2]): | ||
| raise ValueError("output_gradient must have shape [query, model]") | ||
| if target_position >= queries: |
There was a problem hiding this comment.
P3: When callers use target_position=-1, the helper attributes the final query while returning an invalid negative token index. Reject negative positions here, matching AttentionSite validation.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/tdhook/attribution/circuit_lens.py, line 202:
<comment>When callers use `target_position=-1`, the helper attributes the final query while returning an invalid negative token index. Reject negative positions here, matching `AttentionSite` validation.</comment>
<file context>
@@ -0,0 +1,439 @@
+ raise ValueError("output_weight head and head_dim axes do not match values")
+ if output_gradient.shape != (queries, output_weight.shape[2]):
+ raise ValueError("output_gradient must have shape [query, model]")
+ if target_position >= queries:
+ raise IndexError("target_position is outside the attention query axis")
+ if head_indices is None:
</file context>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7fa3842d8d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _captured(activations, "upstream activation"), | ||
| _captured(gradients, "upstream gradient"), |
There was a problem hiding this comment.
Match repeated feature captures by invocation
When an upstream module is invoked more than once and multiple invocations contribute to the target, activation captures are recorded in forward order while gradient captures are recorded in reverse order. Reading .value for both therefore multiplies the last invocation's activation by the first invocation's gradient, silently corrupting feature scores. Pair CapturedTarget.values in the appropriate order or reject sites reached multiple times.
Useful? React with 👍 / 👎.
| output = model(*model_args, **kwargs) | ||
| live_target = _site_scalar(_captured(target_activation, "target feature"), target_feature) | ||
| retain_graph = bool(logit_indices) | ||
| live_target.backward(retain_graph=retain_graph) |
There was a problem hiding this comment.
Preserve caller-owned input gradients
When a tensor in model_args is a leaf with requires_grad=True, this backward pass—and each later logit backward pass—accumulates attribution gradients into its existing .grad. Only parameter gradients are saved and restored, so invoking this helper during training or input optimization silently contaminates caller-owned input gradients even though model gradients are restored.
Useful? React with 👍 / 👎.
Summary
TargetandHookSessionactivation/gradient capture APIs, including graph-retaining activation capture and restoration of existing parameter gradientstransformer-lensandcircuit-tracerin an optionalcircuit-lensintegration extra rather than TDHook core dependenciesValidation
uv run pre-commit run --all-filesuv run pytest tests -q(645 passed, 1 skipped)uv buildCloses #125