Skip to content

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

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

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.

@mikemikimike
mikemikimike deleted the fix/1587-state-dict-roundtrip branch August 2, 2026 02:16
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.

1 participant