Skip to content

fix(model-bridge): restore native state dict round trips - #1591

Open
mikemikimike wants to merge 1 commit into
TransformerLensOrg:dev-4.xfrom
mikemikimike:fix/1587-state-dict-roundtrip
Open

fix(model-bridge): restore native state dict round trips#1591
mikemikimike wants to merge 1 commit into
TransformerLensOrg:dev-4.xfrom
mikemikimike:fix/1587-state-dict-roundtrip

Conversation

@mikemikimike

Copy link
Copy Markdown

Summary

  • Restore TransformerBridge.state_dict() round trips through the native/raw model.
  • Map each TransformerLens key back to all underlying aliased raw keys before loading.
  • Preserve direct clean/native key loading for Tracr-style state dicts.
  • Respect strict=True instead of silently downgrading strict loads.
  • Add regression coverage for round trips, strict missing-key failures, and raw native keys.

Fixes #1587

Validation

  • PYTHONUTF8=1 uv run pytest tests/unit -m "not slow" -n 4 -q — 5,165 passed, 57 skipped, 10 xfailed
  • PYTHONUTF8=1 uv run pytest tests/unit/model_bridge/test_boot_native.py -q — 28 passed
  • uv run mypy . — passed, 421 source files
  • uv run pycln --check ... — passed
  • uv run isort --check-only ... — passed
  • uv run black --check ... — passed
  • git diff --check — passed

The focused tests/integration/model_bridge run was attempted locally but exceeded the 15-minute execution limit without producing a test assertion failure; the unit suite and targeted regression suite completed successfully.

@jlarson4

jlarson4 commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Hi @mikemikimike! There were 3 different PRs all submitted around the same time, providing different solutions to issue #1587. I have opted to take #1598 as our primary solution, but I do want to keep your round-trip tests because they are very strong. Once #1598 lands, would you be willing to rebase to dev-4.x and trim your PR to those tests so I can merge them as is?

@mikemikimike
mikemikimike force-pushed the fix/1587-state-dict-roundtrip branch from 66f5f55 to 2573da4 Compare August 4, 2026 16:03
@mikemikimike
mikemikimike force-pushed the fix/1587-state-dict-roundtrip branch from 2573da4 to 6a12132 Compare August 4, 2026 16:04
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