[torchgen] Refactor and simplify `gen_pyi.py` to use Generic TypeAlias (PEP 585) and Union Type (PEP 604) by XuehaiPan · Pull Request #150727 · pytorch/pytorch

XuehaiPan · 2025-04-05T07:49:18Z

#129001 (comment) is the motivation for the whole stack of PRs. In torch/__init__.py, torch._C.Type shadows from typing import Type, and there is no type stub for torch._C.Type in torch/_C/__init__.pyi. So we need to use from typing import Type as _Type. After enabling Generic TypeAlias (PEP 585) in the .pyi type stub files, we can use type instead of typing.Type or from typing import Type as _Type.

Generic TypeAlias (PEP 585): e.g. typing.List[T] -> list[T], typing.Dict[KT, VT] -> dict[KT, VT], typing.Type[T] -> type[T].
Union Type (PEP 604): e.g. Union[X, Y] -> X | Y, Optional[X] -> X | None, Optional[Union[X, Y]] -> X | Y | None.

Note that in .pyi stub files, we do not need from __future__ import annotations. So this PR does not violate issue #117449:

Deprecate and remove usage of from __future__ import annotations in codebase #117449

Stack from ghstack (oldest at bottom):

cc @ezyang @malfet @xuzhao9 @gramster @bhosmer @bdhirsh @kadeng

[ghstack-poisoned]

pytorch-bot · 2025-04-05T07:49:21Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/150727

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Large queue time for macos-m2-15 instances

✅ You can merge normally! (1 Unrelated Failure)

As of commit c49612e with merge base e2ce17c ():

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / cuda12.4-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu) (gh) (#149370)
MISSING REGRESSION TEST

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ghstack-poisoned]

Copilot

Copilot reviewed 3 out of 4 changed files in this pull request and generated 1 comment.

Files not reviewed (1)

torch/_C/return_types.pyi.in: Language not supported

torchgen/api/python.py

[ghstack-poisoned]

Copilot

Copilot reviewed 3 out of 4 changed files in this pull request and generated no comments.

Files not reviewed (1)

torch/_C/return_types.pyi.in: Language not supported

Comments suppressed due to low confidence (1)

torch/_subclasses/fake_tensor.py:1021

Using 'type: ignore[assignment]' suppresses potential type mismatches between the return type of suggest_memory_format and the expected type. Consider either updating the return type of suggest_memory_format or explicitly casting memory_format to ensure type consistency.

memory_format = None  # type: ignore[assignment]

XuehaiPan · 2025-04-05T08:37:46Z

torchgen/api/python.py

+def format_function_signature(
+    name: str, arguments: Iterable[str] = (), return_type: str | None = None
+) -> str:
+    if not isinstance(arguments, (list, tuple)):
+        arguments = tuple(arguments)
+    return_type = f" -> {return_type}" if return_type is not None else ""
+
+    sig = f"def {name}({', '.join(arguments)}){return_type}: ..."
+    if len(sig) <= 80 or len(arguments) == 0 or tuple(arguments) == ("self",):
+        return sig
+
+    return "\n".join(
+        (
+            f"def {name}(",
+            *(f"    {arg}," for arg in arguments),
+            f"){return_type}: ...",
+        )
+    )


This PR provides a standard way to create a formatted signature. Previously, we did it in a handcrafted way.

[ghstack-poisoned]

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: 2b00421 Pull Request resolved: pytorch#150727

[ghstack-poisoned]

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: 312b22a Pull Request resolved: pytorch#150727

[ghstack-poisoned]

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: 87a2ba3 Pull Request resolved: pytorch#150727

aorenste

Unfortunately PEP604 is python 3.10+ and we have to support python 3.9 - so we can't convert to it yet.

[ghstack-poisoned]

XuehaiPan · 2025-05-14T15:26:15Z

Unfortunately PEP604 is python 3.10+ and we have to support python 3.9 - so we can't convert to it yet.

It's safe to use PEP604 because pyi stub files are not runnable at runtime. In .pyi files, we also do not need from __future__ import annotations.

aorenste · 2025-05-14T15:40:35Z

Unfortunately PEP604 is python 3.10+ and we have to support python 3.9 - so we can't convert to it yet.

It's safe to use PEP604 because pyi stub files are not runnable at runtime. In .pyi files, we also do not need from __future__ import annotations.

That's a good point.

aorenste

n.b.: I spot-checked the diff in the output .pyi files and they look reasonable as well.

aorenste · 2025-05-14T15:41:39Z

tools/pyi/gen_pyi.py

+
+At the moment, this module only handles type stubs for torch and
+torch.Tensor.  It should eventually be expanded to cover all functions
+which come are autogenerated.


nit (I know this is just moved from below - but might as well fix if you have to update a new version):

Suggested change

which come are autogenerated.

which are autogenerated.

aorenste · 2025-05-14T15:54:28Z

torch/_subclasses/fake_tensor.py

    Extract the TensorMetadata of a tensor.
    """
-    memory_format: Optional[torch.memory_format] = suggest_memory_format(t)
+    memory_format = suggest_memory_format(t)


Why this change? It seems that memory_format can validly be None and it also seems unrelated to the rest of this PR...

aorenste · 2025-05-14T15:57:43Z

torchgen/api/python.py

+) -> str:
+    if not isinstance(arguments, (list, tuple)):
+        arguments = tuple(arguments)
+    return_type = f" -> {return_type}" if return_type is not None else ""


Does return_type is None mean "unknown return type" or "returns None"? It seems to be the former - because here if it's None I think it's implicitly Any (because it's missing)... Maybe a quick comment to clarify.

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: e9a854c Pull Request resolved: pytorch#150727

pytorchmergebot · 2025-05-15T07:03:19Z

Starting merge as part of PR stack under #150728

…nd Union Type (PEP 604) (#150728) #129001 (comment) is the motivation for the whole stack of PRs. In `torch/__init__.py`, `torch._C.Type` shadows `from typing import Type`, and there is no type stub for `torch._C.Type` in `torch/_C/__init__.pyi`. So we need to use `from typing import Type as _Type`. After enabling [Generic TypeAlias (PEP 585)](https://peps.python.org/pep-0585) in the `.pyi` type stub files, we can use `type` instead of `typing.Type` or `from typing import Type as _Type`. ------ - [Generic TypeAlias (PEP 585)](https://peps.python.org/pep-0585): e.g. `typing.List[T] -> list[T]`, `typing.Dict[KT, VT] -> dict[KT, VT]`, `typing.Type[T] -> type[T]`. - [Union Type (PEP 604)](https://peps.python.org/pep-0604): e.g. `Union[X, Y] -> X | Y`, `Optional[X] -> X | None`, `Optional[Union[X, Y]] -> X | Y | None`. Note that in `.pyi` stub files, we do not need `from __future__ import annotations`. So this PR does not violate issue #117449: - #117449 ------ Pull Request resolved: #150728 Approved by: https://github.com/cyyever, https://github.com/aorenste ghstack dependencies: #150726, #150727

Update

f29778b

[ghstack-poisoned]

pytorch-bot bot added the ciflow/inductor label Apr 5, 2025

pytorchbot added the open source label Apr 5, 2025

XuehaiPan added 2 commits April 5, 2025 16:06

Update

dc83767

[ghstack-poisoned]

Update

06aedd8

[ghstack-poisoned]

XuehaiPan changed the title ~~Refactor and simplify gen_pyi.py to use Generic TypeAlias (PEP 585) and Union Type (PEP 604)~~ [torchgen] Refactor and simplify gen_pyi.py to use Generic TypeAlias (PEP 585) and Union Type (PEP 604) Apr 5, 2025

XuehaiPan requested a review from Copilot April 5, 2025 08:21

Copilot AI reviewed Apr 5, 2025

View reviewed changes

torchgen/api/python.py Outdated Show resolved Hide resolved

Update

4475b14

[ghstack-poisoned]

XuehaiPan requested a review from Copilot April 5, 2025 08:25

Copilot AI reviewed Apr 5, 2025

View reviewed changes

XuehaiPan requested review from Skylion007, ezyang and zou3519 April 5, 2025 08:27

XuehaiPan commented Apr 5, 2025

View reviewed changes

XuehaiPan added 2 commits May 2, 2025 16:33

Update

b758bf6

[ghstack-poisoned]

Update

0a5fbb6

[ghstack-poisoned]

Update

360bded

[ghstack-poisoned]

Update

8f0a658

[ghstack-poisoned]

Update

474024e

[ghstack-poisoned]

Update

c366aaf

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 2, 2025

[torchgen] Refactor and simplify gen_pyi.py to use Generic TypeAlia…

f2a83ba

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: 2b00421 Pull Request resolved: pytorch#150727

Update

1dbe070

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 10, 2025

[torchgen] Refactor and simplify gen_pyi.py to use Generic TypeAlia…

548ca76

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: 312b22a Pull Request resolved: pytorch#150727

Update

e3b9bdc

[ghstack-poisoned]

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 14, 2025

[torchgen] Refactor and simplify gen_pyi.py to use Generic TypeAlia…

887ef4c

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: 87a2ba3 Pull Request resolved: pytorch#150727

aorenste requested changes May 14, 2025

View reviewed changes

Update

c49612e

[ghstack-poisoned]

XuehaiPan requested a review from aorenste May 14, 2025 15:26

aorenste approved these changes May 14, 2025

View reviewed changes

XuehaiPan added a commit to XuehaiPan/pytorch that referenced this pull request May 14, 2025

[torchgen] Refactor and simplify gen_pyi.py to use Generic TypeAlia…

0dae705

…s (PEP 585) and Union Type (PEP 604) ghstack-source-id: e9a854c Pull Request resolved: pytorch#150727

pytorchmergebot closed this in f7a5aa1 May 15, 2025

pytorchmergebot added the Merged label May 15, 2025

github-actions bot deleted the gh/XuehaiPan/264/head branch June 18, 2025 02:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[torchgen] Refactor and simplify `gen_pyi.py` to use Generic TypeAlias (PEP 585) and Union Type (PEP 604)#150727

[torchgen] Refactor and simplify `gen_pyi.py` to use Generic TypeAlias (PEP 585) and Union Type (PEP 604)#150727
XuehaiPan wants to merge 22 commits intogh/XuehaiPan/264/basefrom
gh/XuehaiPan/264/head

XuehaiPan commented Apr 5, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Apr 5, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

XuehaiPan Apr 5, 2025

Uh oh!

aorenste left a comment

Uh oh!

XuehaiPan commented May 14, 2025

Uh oh!

aorenste commented May 14, 2025

Uh oh!

aorenste left a comment

Uh oh!

aorenste May 14, 2025

Uh oh!

aorenste May 14, 2025

Uh oh!

aorenste May 14, 2025

Uh oh!

pytorchmergebot commented May 15, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

XuehaiPan commented Apr 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/150727

❗ 1 Active SEVs

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

XuehaiPan Apr 5, 2025

Choose a reason for hiding this comment

Uh oh!

aorenste left a comment

Choose a reason for hiding this comment

Uh oh!

XuehaiPan commented May 14, 2025

Uh oh!

aorenste commented May 14, 2025

Uh oh!

aorenste left a comment

Choose a reason for hiding this comment

Uh oh!

aorenste May 14, 2025

Choose a reason for hiding this comment

Uh oh!

aorenste May 14, 2025

Choose a reason for hiding this comment

Uh oh!

aorenste May 14, 2025

Choose a reason for hiding this comment

Uh oh!

pytorchmergebot commented May 15, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

XuehaiPan commented Apr 5, 2025 •

edited

Loading

pytorch-bot bot commented Apr 5, 2025 •

edited

Loading