T5Gemma Support in TransformerLens

Status: Merged
Period: Merged upstream on 2026-06-23
Confidentiality: Public. Public merged contribution. I've kept host-project ownership and adapter scope visible.

Integrated a new encoder-decoder architecture across model mapping, conversion, hooks, registration, typing, and regression tests.

I added T5Gemma support to TransformerLens through a dedicated architecture adapter and decoder bridge, configuration and key translation, runtime registration, conversion alignment, and interpretability hook semantics. The contribution merged upstream with focused adapter, regression, registry, and static-typing evidence.

Type: Open Source
Categories: Open Source
AI & Machine Learning
Technologies: Python
PyTorch
Hugging Face Transformers
TransformerLens
pytest
mypy

Context

TransformerLens users and maintainers who need consistent T5Gemma conversion and interpretability behavior.

The work is a focused adapter contribution inside an established community project.
Full checkpoint parity was unavailable locally because accessible T5Gemma checkpoints were gated.

Problem

Supporting T5Gemma required mapping an encoder-decoder architecture, parameters, configuration, registration, conversion, and hook semantics into an existing shared framework abstraction.

My contribution

I implemented the architecture adapter, T5Gemma decoder bridge, residual hook points, configuration mapping, key translation, conversion alignment, runtime registration, registry metadata, and focused tests.

Approach

Reuse TransformerLens bridge conventions, model the encoder-decoder stages explicitly, translate representative keys, register the model through existing paths, and protect conversion and hook behavior with focused and nearby regression tests.

Outcomes

Merged upstream.
The PR reported 102 focused adapter tests, 44 nearby regression tests, four registry-sync tests, and clean mypy output across 287 source files.

Limitations and current status