[ska2704]

Description

stop_at_layer=n_layers doesn't stop. Blocks raise the stop on entry and no block follows the last one, so the model runs to completion and returns logits, with ln_final and unembed in the cache. A layer sweep gets a d_vocab tensor on its last step and nothing raises. Out of range values do the same, while start_at_layer already rejects them.

n = bridge.cfg.n_layers
bridge(tokens, stop_at_layer=n).shape      # [batch, pos, d_vocab], expected d_model
bridge(tokens, stop_at_layer=n + 5).shape  # same

stop_at_layer=n_layers now returns the residual leaving the last block, which is what HookedTransformer returned (blocks[:stop_at_layer], then return residual). Anything outside [-n_layers, n_layers] raises ValueError, like start_at_layer.

The stop is a forward hook on the last block module, installed by forward() for the length of the call. It isn't on hook_out because raising in a HookPoint's forward hook loses that point's backward hooks, and OpenAI GPT's block returns a list, which never reaches hook_out. The negative index handling was copied in three places and is now one helper.

Type of change

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

Testing

New tests in test_bridge_stop_at_layer.py on the default and compatibility bridges, plus one in test_mamba_adapter.py.

• Against unmodified dev: 12 of the new tests fail, and the 4 controls pass. With the change: test_bridge_stop_at_layer.py 36 passed, and TestMambaStopAtLayer passes.
• tests/unit/model_bridge -m "not slow" plus the stop, start and incl_bwd integration files: 4903 passed, 36 skipped, 3 xfailed.
• Tiny from_config models across the adapter factory: all 73 that boot return the residual at stop_at_layer=n_layers on all three entry points. 90 didn't boot this way or were skipped as too large.
• make check-format and uv run mypy . are clean.

The new test_openai_gpt_adapter.py test wasn't run locally because the weights aren't cached (a tiny OpenAI GPT model was checked instead), and the full integration and acceptance tiers are left to CI.

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