[typing] Add missing type annotations to torch.nn.init module by janumiko · Pull Request #154504 · pytorch/pytorch

janumiko · 2025-05-28T12:45:32Z

Summary

Adds missing type annotations to torch.nn.init and removes # mypy: allow-untyped-defs since all functions are now properly typed.

Changes

Added missing type annotations to initialization functions in the module.
Added missing typing imports: Any, Callable, Union
Removed # mypy: allow-untyped-defs comment
Create Literal types for kaiming initialization mode and nonlinearity.
Created __all__

Why

Better IDE support, catches type errors earlier, and brings the module up to PyTorch's typing standards. No runtime changes - purely additive typing improvements.

Tested with existing test suite and lintrunner.

pytorch-bot · 2025-05-28T12:45:35Z

🔗 Helpful Links

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

📄 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.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 810fdcb with merge base f58143b ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (trunk failure)
dynamo/test_error_messages.py::GraphBreakMessagesTest::test_unsupported_context

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

janumiko · 2025-05-28T12:46:43Z

@pytorchbot label "topic: not user facing"

Skylion007 · 2025-05-28T14:08:30Z

Suggested change

def _make_deprecate(meth: Callable[..., Tensor]) -> Callable[..., Tensor]:

new_name = meth.__name__

old_name = new_name[:-1]

def deprecated_init(*args, **kwargs):

def deprecated_init(*args: tuple[Any], **kwargs: dict[str, Any]) -> Tensor:

def _make_deprecate(meth: Callable[_P, _R]) -> Callable[_P, _R]:

new_name = meth.__name__

old_name = new_name[:-1]

def deprecated_init(*args : _P.args, **kwargs: _P.kwargs) -> _R:

Where _R is a typing.TypeVar and _P is a typing_extensions.ParamSpec

Besides, this should all be replaced with typing_extensions.deprecated anyway

Thanks for the feedback!
Updated the typing to use *args: _P.args, **kwargs: _P.kwargs) -> _R as requested.

Good point about replacing this with @deprecated, I'll work on that in the next PR and keep this one contained to typing.

Skylion007 · 2025-05-28T16:26:29Z

Suggested change

def _calculate_correct_fan(tensor: Tensor, mode: str) -> int:

def _calculate_correct_fan(tensor: Tensor, mode: Literal["fan_in", "fan_out"]) -> int:

Fixed, created FanMode type alias for Literal instead because it appears in the code multiple times.
Did the same for the nonlinearity parameter typing, due to over 10 possible values.

janumiko · 2025-05-28T21:35:03Z

Added __all__ due to failing CI tests.

janumiko · 2025-05-28T21:47:55Z

@pytorchbot label "suppress-bc-linter"

Add supress-bc-linter due to false-positive about changing str to Literal in kaiming init.

Warning: Function kaiming_uniform_: mode changed from str to _FanMode
fetch github.event.pull_request.base.sha
  Warning: Function kaiming_uniform_: nonlinearity changed from str to _NonlinearityType
  Warning: Function kaiming_normal_: mode changed from str to _FanMode
  Warning: Function kaiming_normal_: nonlinearity changed from str to _NonlinearityType
Error: Process completed with exit code 1.

https://github.com/pytorch/pytorch/actions/runs/15306112609/job/43064631148

janumiko · 2025-05-29T10:50:21Z

@pytorchbot merge

pytorch-bot · 2025-05-29T10:50:26Z

Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra.

…dule

janumiko · 2025-05-31T12:46:49Z

Fixed linting issues that were causing CI failures. The only remaining CI error seems to be from the BC linter, which I believe is a false positive.

janumiko · 2025-05-31T18:40:36Z

@pytorchbot merge

pytorch-bot · 2025-05-31T18:40:41Z

Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra.

Skylion007 · 2025-06-03T14:51:14Z

@pytorchbot merge

pytorchmergebot · 2025-06-03T14:53:45Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…h#154504) ## Summary Adds missing type annotations to `torch.nn.init` and removes `# mypy: allow-untyped-defs` since all functions are now properly typed. ## Changes - Added missing type annotations to initialization functions in the module. - Added missing typing imports: `Any`, `Callable`, `Union` - Removed `# mypy: allow-untyped-defs` comment - Create Literal types for kaiming initialization mode and nonlinearity. - Created `__all__` ## Why Better IDE support, catches type errors earlier, and brings the module up to PyTorch's typing standards. No runtime changes - purely additive typing improvements. Tested with existing test suite and lintrunner. Pull Request resolved: pytorch#154504 Approved by: https://github.com/Skylion007

janumiko requested review from albanD, jbschlosser and mikaylagawarecki as code owners May 28, 2025 12:45

pytorchbot added the open source label May 28, 2025

pytorch-bot Bot added the topic: not user facing topic category label May 28, 2025

janumiko force-pushed the fix/type-hints-nn-init branch 2 times, most recently from 00bbf1a to fd1d2ee Compare May 28, 2025 13:16

Skylion007 reviewed May 28, 2025

View reviewed changes

janumiko force-pushed the fix/type-hints-nn-init branch 2 times, most recently from cfe9803 to 5916600 Compare May 28, 2025 14:51

Skylion007 reviewed May 28, 2025

View reviewed changes

janumiko force-pushed the fix/type-hints-nn-init branch 2 times, most recently from eeaf9b1 to 4c9f011 Compare May 28, 2025 17:01

Skylion007 approved these changes May 28, 2025

View reviewed changes

janumiko force-pushed the fix/type-hints-nn-init branch 2 times, most recently from 3f780af to eb9fb8d Compare May 28, 2025 21:34

pytorch-bot Bot added the suppress-bc-linter Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) label May 28, 2025

[typing] Add __all__ and missing type annotations to torch.nn.init mo…

810fdcb

…dule

janumiko force-pushed the fix/type-hints-nn-init branch from eb9fb8d to 810fdcb Compare May 30, 2025 23:00

pytorch-bot Bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 3, 2025

pytorchmergebot added the merging label Jun 3, 2025

pytorchmergebot added the Merged label Jun 3, 2025

pytorchmergebot closed this in 31405a6 Jun 3, 2025

pytorchmergebot removed the merging label Jun 3, 2025

randolf-scholz mentioned this pull request Oct 16, 2025

[Meta] Type Hints in Pytorch #165628

Open

-def _make_deprecate(meth: Callable[..., Tensor]) -> Callable[..., Tensor]:
-    new_name = meth.__name__
-    old_name = new_name[:-1]
-    def deprecated_init(*args, **kwargs):
-    def deprecated_init(*args: tuple[Any], **kwargs: dict[str, Any]) -> Tensor:
+def _make_deprecate(meth: Callable[_P, _R]) -> Callable[_P, _R]:
+    new_name = meth.__name__
+    old_name = new_name[:-1]
+    def deprecated_init(*args : _P.args, **kwargs: _P.kwargs) -> _R:

	def _calculate_correct_fan(tensor: Tensor, mode: str) -> int:
	def _calculate_correct_fan(tensor: Tensor, mode: Literal["fan_in", "fan_out"]) -> int:

Conversation

janumiko commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Why

Uh oh!

pytorch-bot Bot commented May 28, 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/154504

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

janumiko commented May 28, 2025

Uh oh!

Skylion007 May 28, 2025 • 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!

Skylion007 May 28, 2025

Choose a reason for hiding this comment

Uh oh!

janumiko May 28, 2025 • 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!

Skylion007 May 28, 2025

Choose a reason for hiding this comment

Uh oh!

janumiko May 28, 2025 • 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!

janumiko commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

janumiko commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

janumiko commented May 29, 2025

Uh oh!

pytorch-bot Bot commented May 29, 2025

Uh oh!

janumiko commented May 31, 2025

Uh oh!

janumiko commented May 31, 2025

Uh oh!

pytorch-bot Bot commented May 31, 2025

Uh oh!

Skylion007 commented Jun 3, 2025

Uh oh!

pytorchmergebot commented Jun 3, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

janumiko commented May 28, 2025 •

edited

Loading

pytorch-bot Bot commented May 28, 2025 •

edited

Loading

Skylion007 May 28, 2025 •

edited

Loading

janumiko May 28, 2025 •

edited

Loading

janumiko May 28, 2025 •

edited

Loading

janumiko commented May 28, 2025 •

edited

Loading

janumiko commented May 28, 2025 •

edited

Loading