hook_result & Hook Aliases issues

[jlarson4]

Description

AttentionBridge aliases hook_result to hook_out. But hook_out fires on the HF attention output which has shape [batch, pos, d_model] (3D, already projected by W_O). No conversion reshapes it into the expected [batch, pos, n_heads, d_head] (4D per-head format). This resolves that bug.

Ensure the cache always stores under the alias name (since that's what hook.name returns and what get_act_name() generates), or make ActivationCache.getitem resolve aliases transparently. This resolves a second, related bug

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
• I have not rewritten tests relating to key interfaces which would affect backward compatibility