fix(bridge): make native TransformerBridge state_dict()/load_state_dict() true inverses

[LightWork666]

Fixes #1587.

The bug

On a native TransformerBridge, state_dict() returns TL-renamed keys (embed.weight, blocks.0.attn.q.weight, ...), but load_state_dict() only knew how to match raw native parameter names. So bridge.load_state_dict(bridge.state_dict()) silently did nothing — no error, no warning, params just stayed whatever they were before. Worse, strict=True was silently downgraded to strict=False whenever the key counts didn't line up, so there was no way to even notice the round trip had failed.

The fix

load_state_dict now builds the inverse of the TL-key renaming that state_dict() applies (_tl_key_to_actual_keys), so it can map TL-format keys back to the underlying native parameter paths before handing them to the wrapped model's own load_state_dict. One wrinkle: some bridge components expose the same underlying parameter through more than one attribute path — e.g. GPT-2's q/k/v are views into the wrapped module's combined c_attn weight, reachable both through a block-level shortcut and through the nested _original_component chain. Writing to only one of those paths leaves the model's actual forward pass untouched even though state_dict() looks fine, so the mapping keeps every alias for a given TL key, not just the first one found.

The silent strict=True → False downgrade is gone. Missing/unexpected keys are now computed properly (scoped to what the TL state dict actually needs) and raise a real RuntimeError under strict=True, matching how HookedTransformer already behaves.

The raw-key loading path used by tracr (make_tracr_transformer_bridge_state_dict) still works — raw native keys are matched directly before falling through to the TL-key path.

Testing

• New tests in tests/unit/model_bridge/test_state_dict_round_trip.py: round-trip actually overwrites params (not a no-op), strict=True raises on both missing and unexpected keys, strict=False doesn't raise on a partial dict, the tracr raw-key path still loads, and a real GPT-2 case checking forward-pass logits match after a zero-and-reload cycle (this last one is what caught the aliasing issue above — a naive key-rename fix passes the round-trip-key-equality check but still produces different logits, because it never touches the aliased storage).
• uv run mypy . — clean.
• Ran the full tests/unit/model_bridge/ suite (207 files) in isolated batches; everything passes except a handful of pre-existing generate()/KV-cache crashes already tracked as an upstream PyTorch/HF bug on Apple Silicon in tests/QUARANTINES.md — confirmed via git stash that those reproduce identically on unmodified dev-4.x, unrelated to this change.

Note on #1595

I noticed after finishing this that #1595 is already open for the same issue, taking a related but different approach (it detects aliasing by comparing live tensor identity rather than by TL-key-name collisions, which is arguably a more principled check). I'm submitting this anyway since it was already done and transparency seemed better than not mentioning it. Happy to have the maintainers pick whichever they prefer, or close this if #1595 is the better fix.