Skip to content

Require evaluation mode for attribution patching - #1788

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/attribution-patching-require-eval
Sep 15, 2026
Merged

jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
emerardd:fix/attribution-patching-require-eval

Conversation

@emerardd

Copy link
Copy Markdown
Contributor

Description

Require attribution_patch() to run with the model and every submodule in evaluation mode, preventing separate clean and corrupt forwards from mixing independent dropout realizations into false causal effects.

This extracts the all-submodule state validation already used by JacobianLens.fit() into a shared private helper. The helper also inspects a Bridge's deliberately unregistered original_model, reports the offending module paths, and rejects invalid state without mutating caller-owned training flags. Input-shape validation still runs first, while the state guard runs before either attribution forward.

Regression coverage includes root training mode, an eval root with only a nested dropout in training mode, hidden original_model state, preservation of training flags, and exact-zero attribution for identical inputs in evaluation mode.

Fixes #1787

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes; only the affected unit-test surface was run (142 passed)
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

Local validation

pytest tests/unit/tools/test_attribution_patching.py tests/unit/tools/test_jacobian_lens.py -q
142 passed, 3 warnings in 27.65s

pycln --check --all: passed
isort --check-only: passed
black --check: passed
mypy: Success: no issues found in 399 source files
git diff --check: passed

@jlarson4

Copy link
Copy Markdown
Collaborator

Looks good, thanks @emerardd

@jlarson4
jlarson4 merged commit ab42d85 into TransformerLensOrg:dev Sep 15, 2026
26 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants