Skip to content

Preserve alias hook replacements across no-op callbacks - #1722

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
koriyoshi2041:rios/fix-alias-hook-lineage
Aug 25, 2026
Merged

Preserve alias hook replacements across no-op callbacks#1722
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
koriyoshi2041:rios/fix-alias-hook-lineage

Conversation

@koriyoshi2041

Copy link
Copy Markdown
Contributor

Problem

When a hook runs under multiple alias names, each callback sees the previous callback's replacement. However, if a later alias callback returns None, HookPoint.full_hook returns that final None and silently discards an earlier replacement. This affects both forward activations and backward gradients.

Fix

Track whether any alias callback changed the value and return the accumulated result after all aliases have run. The existing all-no-op behavior remains unchanged.

Test

  • uv run pytest tests/unit/test_hook_points.py -q — 31 passed
  • uv run black --check transformer_lens/hook_points.py tests/unit/test_hook_points.py
  • uv run isort --check-only transformer_lens/hook_points.py tests/unit/test_hook_points.py
  • uv run mypy transformer_lens — 388 source files clean
  • git diff --check

The regressions cover forward replacements, conversion reversion, and backward gradient replacements.

Risk

Low. The change is limited to alias dispatch and only changes the returned value when an earlier alias callback supplied a replacement. Canonical hooks and all-None alias chains retain their current behavior.

@jlarson4

Copy link
Copy Markdown
Collaborator

Looks great! Thanks @koriyoshi2041 for putting this together

@jlarson4
jlarson4 merged commit 2c639ab into TransformerLensOrg:dev Aug 25, 2026
25 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