Fix/1076 logit lens layer norm

[evcyen]

Description

Fixes #1076

Issue: Logit lens with accumulated_resid(..., apply_ln=True) gave wrong probabilities and entropy (early layers looked almost uniform instead of peaked).

Root cause: apply_ln_to_stack used the final layer’s scale (ln_final.hook_scale) for every layer. Early layers were normalized with a scale that didn’t match their own residual.

Fix: apply_ln_to_stack now supports recompute_ln. When recompute_ln=True and the target is the final layer, each slice is normalized with model.ln_final using statistics from that slice. accumulated_resid(..., apply_ln=True) calls apply_ln_to_stack(..., recompute_ln=True) for the final layer so logit lens uses per-layer normalization. Other callers keep recompute_ln=False to keep the single cached scale.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

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