[`typings`] Automatically type decorator return types as `tuple | X` by tomaarsen · Pull Request #43446 · huggingface/transformers

tomaarsen · 2026-01-23T12:58:02Z

What does this PR do?

This PR updates the can_return_tuple and ~~check_model_inputs~~ capture_outputs (see #43446 (comment)) typings such that:

from transformers.utils import can_return_tuple
from transformers.utils.generic import check_model_inputs
try:
    from typing import reveal_type
except ImportError:
    from typing_extensions import reveal_type


@can_return_tuple
def my_func(foo: int, bar: int, **kwargs) -> dict[str, int]:
    return {
        "sum": foo + bar,
        "product": foo * bar,
    }

result = my_func(1, 2)
reveal_type(result)  # tuple[Unknown, ...] | dict[str, int]


@check_model_inputs
def my_second_func(foo: int, bar: int, **kwargs) -> dict[str, int | float]:
    return {
        "minus": float(foo - bar),
        "div": foo / bar,
    }

result2 = my_second_func(1, 2)
reveal_type(result2)  # tuple[Unknown, ...] | dict[str, int | float]

In short: functions and method that use can_return_tuple or check_model_inputs must currently be typed like tuple | X, but can now be typed like X as well.

I also had Copilot help me write a check_decorator_return_types.py file that ensures that functions decorated with can_return_tuple or check_model_inputs:

Have an explicit, non-None return annotation.
Are not annotated with a union that already includes tuple.

The latter can be auto-fixed with python utils/check_decorator_return_types.py --fix_and_overwrite, and the former must be manually fixed. This checking script is included in make fix-repo and the CI. The script only runs in about ~5 seconds for me.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case: [BC] Update get_(text|image|audio|video)_features methods to return BaseModelOutputWithPooling #42564 (comment)
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

cc @zucchini-nlp @vasqu @tarekziade

Tom Aarsen

…rn type This updates the typings of these two functions, so that a wrapped function that has return type X is automatically typed as `tuple | X`.

It verifies that users don't use e.g. `tuple | BaseModelOutputWithPooling` return typings anymore, as they should use `BaseModelOutputWithPooling` instead then. It also makes sure that a typing is used

This is the main actual change, the rest is just typings

And remove return_dict from altclip and clap: the can_return_tuple should take care of it fully

HuggingFaceDocBuilderDev · 2026-01-23T13:07:12Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

tomaarsen · 2026-01-23T13:07:25Z

src/transformers/utils/generic.py

+P = ParamSpec("P")
+T = TypeVar("T")


Reviewer note: the ParamSpec and TypeVar as used here match Python's own typing documentation exactly: https://docs.python.org/3.10/library/typing.html#typing.ParamSpec

They also use an example of a decorator, so the implementation here is quite standard/common.

Doesn't hurt to add the doc link as comment to the wrappers we use it at below - it's like 1:1 matching the example :D

Very true! Done in dbd43d6 & 3fd9571

tomaarsen · 2026-01-23T13:09:05Z

src/transformers/models/blt/modular_blt.py


        self.post_init()

-    @check_model_inputs


Reviewer note: There are only 2 actual code changes under src/transformers/models, and this is one of them. This class does not return a BaseModelOutput, and the other related private classes (BltLocalEncoder, BltGlobalTransformer) don't use this @check_model_inputs either.

ah I see, so it is an issue with how the model was implemented. So many tiny inconsistencies 😅

Actually, they shouldn't inherit from PreTrained model at all it seems like 😅 it's not a high priority model so fine for now

tomaarsen · 2026-01-23T13:10:43Z

src/transformers/models/altclip/modeling_altclip.py

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

+    @can_return_tuple


Reviewer note: There are only 2 actual code changes under src/transformers/models, and this is one of them. This class copies from CLAP, which uses @can_return_tuple, but this class did not. I've added it here.

I also updated this class and the CLAP variant to remove the return_dict = return_dict if return_dict is not None else self.config.use_return_dict line, as that was just dead code.

nice! Ideally would be great to start using check_model_inputs for pretrained models. Though it might require more manual work than current state of PR

utils/check_decorator_return_types.py

zucchini-nlp

Great work! I have just a few questions about the auto-checker

zucchini-nlp · 2026-01-23T13:40:38Z

src/transformers/models/altclip/modeling_altclip.py

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

+    @can_return_tuple


nice! Ideally would be great to start using check_model_inputs for pretrained models. Though it might require more manual work than current state of PR

zucchini-nlp · 2026-01-23T13:42:20Z

src/transformers/models/blt/modeling_blt.py


        self.post_init()

-    @check_model_inputs


accidentally deleted?

zucchini-nlp · 2026-01-23T13:43:44Z

src/transformers/models/blt/modular_blt.py


        self.post_init()

-    @check_model_inputs


ah I see, so it is an issue with how the model was implemented. So many tiny inconsistencies 😅

utils/check_decorator_return_types.py

vasqu

Lgtm

Maybe we should start splitting into more than fix-repo again @Cyrilvallez 👀

vasqu · 2026-01-23T13:45:48Z

src/transformers/models/blt/modular_blt.py


        self.post_init()

-    @check_model_inputs


Actually, they shouldn't inherit from PreTrained model at all it seems like 😅 it's not a high priority model so fine for now

vasqu · 2026-01-23T13:54:51Z

src/transformers/utils/generic.py

+P = ParamSpec("P")
+T = TypeVar("T")


Doesn't hurt to add the doc link as comment to the wrappers we use it at below - it's like 1:1 matching the example :D

utils/check_decorator_return_types.py

This reverts commit f71a072.

Without extra code; this time

github-actions · 2026-01-23T15:08:12Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43446&sha=3fd957

tomaarsen · 2026-02-20T10:02:44Z

I updated this PR with @Cyrilvallez' latest changes regarding model_check_inputs -> capture_outputs, so that methods wrapped with capture_outputs are also nicely typed:

from transformers.utils import can_return_tuple
from transformers.utils.output_capturing import capture_outputs

try:
    from typing import reveal_type
except ImportError:
    from typing_extensions import reveal_type


@can_return_tuple
def my_func(foo: int, bar: int, **kwargs) -> dict[str, int]:
    return {
        "sum": foo + bar,
        "product": foo * bar,
    }


result = my_func(1, 2)
reveal_type(result)  # tuple[Unknown, ...] | dict[str, int]


@capture_outputs
def my_second_func(foo: int, bar: int, **kwargs) -> dict[str, int | float]:
    return {
        "minus": float(foo - bar),
        "div": foo / bar,
    }


result2 = my_second_func(1, 2)
reveal_type(result2)  # tuple[Unknown, ...] | dict[str, int | float]

ty check demo.py

info[revealed-type]: Revealed type
  --> demo_can_return_tuple_hint.py:19:13
   |
18 | result = my_func(1, 2)
19 | reveal_type(result)  # tuple[Unknown, ...] | dict[str, int]
   |             ^^^^^^ `tuple[Unknown, ...] | dict[str, int]`
   |

info[revealed-type]: Revealed type
  --> demo_can_return_tuple_hint.py:31:13
   |
30 | result2 = my_second_func(1, 2)
31 | reveal_type(result2)  # tuple[Unknown, ...] | dict[str, int | float]
   |             ^^^^^^^ `tuple[Unknown, ...] | dict[str, int | float]`
   |

Found 2 diagnostics

This should help cut down some bloat from our modeling/modular files without considerably complicating any of the decorators. Ideally, we can get this merged somewhat quickly, as the merge conflicts start quickly too.

also cc @molbap as you're also working on decorator-related cleanup in #43590 right now

Tom Aarsen

github-actions · 2026-03-02T12:09:01Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: afmoe, aimv2, albert, align, altclip, aria, audioflamingo3, aya_vision, bert, bert_generation, blip, blip_2, blt, bridgetower, bros, camembert

tomaarsen added 7 commits January 23, 2026 13:36

Automatically update can_return_tuple/check_model_inputs wrapped retu…

bffa963

…rn type This updates the typings of these two functions, so that a wrapped function that has return type X is automatically typed as `tuple | X`.

Add check_decorator_return_types.py script to check against 'tuple | X'

6762069

It verifies that users don't use e.g. `tuple | BaseModelOutputWithPooling` return typings anymore, as they should use `BaseModelOutputWithPooling` instead then. It also makes sure that a typing is used

Run checking script & fix existing typings

fa9caf7

Fix blt by removing check_model_inputs; matching other classes

abaca9c

This is the main actual change, the rest is just typings

Use can_return_tuple on altclip

3652612

And remove return_dict from altclip and clap: the can_return_tuple should take care of it fully

Add docstring for check_model_inputs / can_return_tuple

9d9666a

Import ParamSpec via generic as we're in Python 3.10+

99d5deb

tomaarsen commented Jan 23, 2026

View reviewed changes

Add check_decorator_return_types to fix-repo & check-repo

6ad9269

tarekziade reviewed Jan 23, 2026

View reviewed changes

utils/check_decorator_return_types.py Outdated Show resolved Hide resolved

utils/check_decorator_return_types.py Outdated Show resolved Hide resolved

zucchini-nlp reviewed Jan 23, 2026

View reviewed changes

tomaarsen added 2 commits January 23, 2026 15:35

Remove some dead code

a9afbdb

Update _decorator_name to not check for our targets yet

aa97fc6

vasqu approved these changes Jan 23, 2026

View reviewed changes

tomaarsen added 8 commits January 23, 2026 15:41

Use src/transformers/models path instead

f71a072

Revert "Use src/transformers/models path instead"

63155b6

This reverts commit f71a072.

Use src/transformers/models path instead

864a6a7

Without extra code; this time

Simplify _is_none_annotation

456ee32

Explain why super() is skipped

3b0808e

Add comments for ParamSpec/TypeVar linking to Python docs

dbd43d6

Make style

688be4b

Move decorator typing comments around

3fd9571

Merge branch 'main' into feat/auto_decorator_return_typing

155fad7

Sebmono mentioned this pull request Feb 1, 2026

[PR] [typings] Automatically type decorator return types as tuple | X Sandgarden-Demo/transformers#31

Open

5 tasks

zucchini-nlp mentioned this pull request Feb 16, 2026

Remove many output_attentions and other traced outputs on 100+ models #43590

Open

Merge branch 'main' into feat/auto_decorator_return_typing

ea9df1d

Rerun utils/check_decorator_return_types.py

bc81a60

tomaarsen requested a review from Cyrilvallez February 20, 2026 09:39

Update incorrect typings on modular classes that inherit decorators

b05bc5d

Merge branch 'main' into feat/auto_decorator_return_typing

e3da0ce

Conversation

tomaarsen commented Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Jan 23, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tomaarsen Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tomaarsen Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tomaarsen Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jan 23, 2026

Uh oh!

tomaarsen commented Feb 20, 2026

Uh oh!

github-actions bot commented Mar 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tomaarsen commented Jan 23, 2026 •

edited

Loading

tomaarsen Jan 23, 2026 •

edited

Loading

tomaarsen Jan 23, 2026 •

edited

Loading

tomaarsen Jan 23, 2026 •

edited

Loading