[typing] Add type hints to `init` methods in `torch.distributions`. #144197

randolf-scholz · 2025-01-04T18:20:12Z

Fixes #144196
Extends #144106 and #144110

Open Problems:

Annotating with numbers.Number is a bad idea, should consider using float, SupportsFloat or some Procotol. [typing] Add type hints to __init__ methods in torch.distributions. #144197 (comment)

Notes

beta.py: needed to add type: ignore since broadcast_all is untyped.
categorical.py: converted else branches of mutually exclusive arguments to if branch¹.
~~dirichlet.py: replaced axis with dim arguments.~~ Dirichlet.mode: use dim= instead of axis= #144402

gemoetric.py: converted else branches of mutually exclusive arguments to if branch¹.

~~independent.py: fixed bug in Independent.__init__ where tuple[int, ...] could be passed to Distribution.__init__ instead of torch.Size.~~ EDIT: turns out the bug is related to typing of torch.Size. Improve static typing for torch.Size #144218

independent.py: made Independent a generic class of its base distribution.

multivariate_normal.py: converted else branches of mutually exclusive arguments to if branch¹.

relaxed_bernoulli.py: added class-level type hint for base_dist.

relaxed_categorical.py: added class-level type hint for base_dist.

~~transforms.py: Added missing argument to docstring of ReshapeTransform~~ ReshapeTransform: added missing argument in docstring #144401

~~transforms.py: Fixed bug in AffineTransform.sign (could return Tensor instead of int).~~ Fix AffineTransform.sign #144400

transforms.py: Added type: ignore comments to AffineTransform.log_abs_det_jacobian²; replaced torch.abs(scale) with scale.abs().

transforms.py: Added type: ignore comments to AffineTransform.__eq__².

transforms.py: Fixed type hint on CumulativeDistributionTransform.domain. Note that this is still an LSP violation, because Transform.domain is defined as Constraint, but Distribution.domain is defined as Optional[Constraint].

skipped: constraints.py, constraints_registry.py, kl.py, utils.py, exp_family.py, __init__.py.

Remark

TransformedDistribution: __init__ uses the check if reinterpreted_batch_ndims > 0:, which can lead to the creation of Independent distributions with only 1 component. This results in awkward code like base_dist.base_dist in LogisticNormal.

import torch
from torch.distributions import *
b1 = Normal(torch.tensor([0.0]), torch.tensor([1.0]))
b2 = MultivariateNormal(torch.tensor([0.0]), torch.eye(1))
t = StickBreakingTransform()
d1 = TransformedDistribution(b1, t)
d2 = TransformedDistribution(b2, t)
print(d1.base_dist)  # Independent with 1 dimension
print(d2.base_dist)  # MultivariateNormal

One could consider changing this to if reinterpreted_batch_ndims > 1:.

cc @fritzo @neerajprad @alicanb @nikitaved @ezyang @malfet @xuzhao9 @gramster

Otherwise, we would have to add a bunch of type: ignore comments to make mypy happy, as it isn't able to perform the type narrowing. Ideally, such code should be replaced with structural pattern matching once support for Python 3.9 is dropped. ↩ ↩² ↩³
Usage of isinstance(value, numbers.Real) leads to problems with static typing, as the numbers module is not supported by mypy (see https://github.com/python/mypy/issues/3186). This results in us having to add type-ignore comments in several places ↩ ↩²

… to 'Tensor'

pytorch-bot · 2025-01-04T18:20:15Z

🔗 Helpful Links

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

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

✅ No Failures

As of commit aafae2d with merge base 87a63a9 ():
💚 Looks good so far! There are no failures yet. 💚

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

randolf-scholz · 2025-01-04T18:21:30Z

@pytorchbot label "module: typing"
@pytorchbot label "module: distributions"

randolf-scholz · 2025-01-04T18:22:02Z

@pytorchbot label "module: distributions"

…init

randolf-scholz · 2025-01-04T18:47:50Z

@pytorchbot label "release notes: python_frontend"

malfet · 2025-01-15T15:48:49Z

@randolf-scholz but I feel like 50+% of this change could be landed separately even with mypy suppression for the whole file, say something like torch/distributions/weibull.py which will make incremental reviews faster.

randolf-scholz · 2025-01-15T15:56:44Z

@randolf-scholz but I feel like 50+% of this change could be landed separately even with mypy suppression for the whole file, say something like torch/distributions/weibull.py which will make incremental reviews faster.

One can do that, but I think it would make most sense to first fix the numbers.Number issue and then update the annotations after that. Otherwise, we may end up with incorrect annotations or having to update them again.

randolf-scholz · 2025-01-28T16:06:09Z

I added annotations to broadcast_all in this PR, since this function is used in virtually all __init__ methods, this required adding a few assert statements to make mypy happy when mutually exclusive arguments are used.

randolf-scholz · 2025-01-28T16:08:06Z

torch/distributions/utils.py

-def broadcast_all(*values):
+# FIXME: Use (*values: *Ts) -> tuple[Tensor for T in Ts] if Mapping-Type is ever added.
+#   See https://github.com/python/typing/issues/1216#issuecomment-2126153831
+def broadcast_all(*values: Union[Tensor, Number]) -> tuple[Tensor, ...]:


Currently, the annotation does not cover the case of __torch_function__.

This is somewhat outside the scope of this PR, I think possibly a Protocol type should be added to torch.types that covers the concrete semantics.

torch/distributions/transforms.py

randolf-scholz · 2025-02-04T16:33:05Z

One nice thing to also add would be overloads for __init__ methods with mutually exclusive arguments like

pytorch/torch/distributions/bernoulli.py

Line 43 in 87a63a9

def __init__(self, probs=None, logits=None, validate_args=None):

However, the current API does not make use of keyword-only arguments for some reason. The following would be much nicer, imo:

@overload
def __init__(self, probs: Tensor, *, validate_args: Optional[bool]=None) -> None: ...
@overload
def __init__(self, *, logits: Tensor, validate_args: Optional[bool]=None) -> None: ...
def __init__(self, probs: Optional[Tensor]=None, *, logits: Optional[Tensor]=None, validate_args: Optional[bool]=None) -> None:

However, this would require some backward compatibility breaking changes, for instance the argument order in Binomial.__init__ would need to be changed. But maybe a possibility in the future.

github-actions · 2025-04-05T16:37:53Z

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a 9E88 mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

Skylion007 · 2025-04-06T16:59:34Z

@randolf-scholz Mind doing a rebase? @malfet Think this is blocked by your review

malfet · 2025-04-06T17:42:33Z

@pytorchbot rebase

pytorchmergebot · 2025-04-06T17:44:05Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2025-04-06T17:44:08Z

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict pull/144197/head returned non-zero exit code 1

Rebasing (1/31)
Auto-merging torch/distributions/utils.py
CONFLICT (content): Merge conflict in torch/distributions/utils.py
error: could not apply ef2329f5cc9... added type hints to lazy_property
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
hint: Disable this message with "git config set advice.mergeConflict false"
Could not apply ef2329f5cc9... added type hints to lazy_property

Raised by https://github.com/pytorch/pytorch/actions/runs/14295053243

malfet · 2025-04-06T17:48:45Z

@pytorchbot merge -f "Rebase fails, let's see if merge will be fine"

pytorchmergebot · 2025-04-06T17:50:22Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

…`. (pytorch#144197) Fixes pytorch#144196 Extends pytorch#144106 and pytorch#144110 ## Open Problems: - [ ] Annotating with `numbers.Number` is a bad idea, should consider using `float`, `SupportsFloat` or some `Procotol`. pytorch#144197 (comment) # Notes - `beta.py`: needed to add `type: ignore` since `broadcast_all` is untyped. - `categorical.py`: converted `else` branches of mutually exclusive arguments to `if` branch[^2]. - ~~`dirichlet.py`: replaced `axis` with `dim` arguments.~~ pytorch#144402 - `gemoetric.py`: converted `else` branches of mutually exclusive arguments to `if` branch[^2]. - ~~`independent.py`: fixed bug in `Independent.__init__` where `tuple[int, ...]` could be passed to `Distribution.__init__` instead of `torch.Size`.~~ **EDIT:** turns out the bug is related to typing of `torch.Size`. pytorch#144218 - `independent.py`: made `Independent` a generic class of its base distribution. - `multivariate_normal.py`: converted `else` branches of mutually exclusive arguments to `if` branch[^2]. - `relaxed_bernoulli.py`: added class-level type hint for `base_dist`. - `relaxed_categorical.py`: added class-level type hint for `base_dist`. - ~~`transforms.py`: Added missing argument to docstring of `ReshapeTransform`~~ pytorch#144401 - ~~`transforms.py`: Fixed bug in `AffineTransform.sign` (could return `Tensor` instead of `int`).~~ pytorch#144400 - `transforms.py`: Added `type: ignore` comments to `AffineTransform.log_abs_det_jacobian`[^1]; replaced `torch.abs(scale)` with `scale.abs()`. - `transforms.py`: Added `type: ignore` comments to `AffineTransform.__eq__`[^1]. - `transforms.py`: Fixed type hint on `CumulativeDistributionTransform.domain`. Note that this is still an LSP violation, because `Transform.domain` is defined as `Constraint`, but `Distribution.domain` is defined as `Optional[Constraint]`. - skipped: `constraints.py`, `constraints_registry.py`, `kl.py`, `utils.py`, `exp_family.py`, `__init__.py`. ## Remark `TransformedDistribution`: `__init__` uses the check `if reinterpreted_batch_ndims > 0:`, which can lead to the creation of `Independent` distributions with only 1 component. This results in awkward code like `base_dist.base_dist` in `LogisticNormal`. ```python import torch from torch.distributions import * b1 = Normal(torch.tensor([0.0]), torch.tensor([1.0])) b2 = MultivariateNormal(torch.tensor([0.0]), torch.eye(1)) t = StickBreakingTransform() d1 = TransformedDistribution(b1, t) d2 = TransformedDistribution(b2, t) print(d1.base_dist) # Independent with 1 dimension print(d2.base_dist) # MultivariateNormal ``` One could consider changing this to `if reinterpreted_batch_ndims > 1:`. [^1]: Usage of `isinstance(value, numbers.Real)` leads to problems with static typing, as the `numbers` module is not supported by `mypy` (see <python/mypy#3186>). This results in us having to add type-ignore comments in several places [^2]: Otherwise, we would have to add a bunch of `type: ignore` comments to make `mypy` happy, as it isn't able to perform the type narrowing. Ideally, such code should be replaced with structural pattern matching once support for Python 3.9 is dropped. Pull Request resolved: pytorch#144197 Approved by: https://github.com/malfet Co-authored-by: Aaron Gokaslan <aaronGokaslan@gmail.com>

…`. (pytorch#144197) Fixes pytorch#144196 Extends pytorch#144106 and pytorch#144110 ## Open Problems: - [ ] Annotating with `numbers.Number` is a bad idea, should consider using `float`, `SupportsFloat` or some `Procotol`. pytorch#144197 (comment) # Notes - `beta.py`: needed to add `type: ignore` since `broadcast_all` is untyped. - `categorical.py`: converted `else` branches of mutually exclusive arguments to `if` branch[^2]. - ~~` 9A82 dirichlet.py`: replaced `axis` with `dim` arguments.~~ pytorch#144402 - `gemoetric.py`: converted `else` branches of mutually exclusive arguments to `if` branch[^2]. - ~~`independent.py`: fixed bug in `Independent.__init__` where `tuple[int, ...]` could be passed to `Distribution.__init__` instead of `torch.Size`.~~ **EDIT:** turns out the bug is related to typing of `torch.Size`. pytorch#144218 - `independent.py`: made `Independent` a generic class of its base distribution. - `multivariate_normal.py`: converted `else` branches of mutually exclusive arguments to `if` branch[^2]. - `relaxed_bernoulli.py`: added class-level type hint for `base_dist`. - `relaxed_categorical.py`: added class-level type hint for `base_dist`. - ~~`transforms.py`: Added missing argument to docstring of `ReshapeTransform`~~ pytorch#144401 - ~~`transforms.py`: Fixed bug in `AffineTransform.sign` (could return `Tensor` instead of `int`).~~ pytorch#144400 - `transforms.py`: Added `type: ignore` comments to `AffineTransform.log_abs_det_jacobian`[^1]; replaced `torch.abs(scale)` with `scale.abs()`. - `transforms.py`: Added `type: ignore` comments to `AffineTransform.__eq__`[^1]. - `transforms.py`: Fixed type hint on `CumulativeDistributionTransform.domain`. Note that this is still an LSP violation, because `Transform.domain` is defined as `Constraint`, but `Distribution.domain` is defined as `Optional[Constraint]`. - skipped: `constraints.py`, `constraints_registry.py`, `kl.py`, `utils.py`, `exp_family.py`, `__init__.py`. ## Remark `TransformedDistribution`: `__init__` uses the check `if reinterpreted_batch_ndims > 0:`, which can lead to the creation of `Independent` distributions with only 1 component. This results in awkward code like `base_dist.base_dist` in `LogisticNormal`. ```python import torch from torch.distributions import * b1 = Normal(torch.tensor([0.0]), torch.tensor([1.0])) b2 = MultivariateNormal(torch.tensor([0.0]), torch.eye(1)) t = StickBreakingTransform() d1 = TransformedDistribution(b1, t) d2 = TransformedDistribution(b2, t) print(d1.base_dist) # Independent with 1 dimension print(d2.base_dist) # MultivariateNormal ``` One could consider changing this to `if reinterpreted_batch_ndims > 1:`. [^1]: Usage of `isinstance(value, numbers.Real)` leads to problems with static typing, as the `numbers` module is not supported by `mypy` (see <python/mypy#3186>). This results in us having to add type-ignore comments in several places [^2]: Otherwise, we would have to add a bunch of `type: ignore` comments to make `mypy` happy, as it isn't able to perform the type narrowing. Ideally, such code should be replaced with structural pattern matching once support for Python 3.9 is dropped. Pull Request resolved: pytorch#144197 Approved by: https://github.com/malfet Co-authored-by: Aaron Gokaslan <aaronGokaslan@gmail.com>

randolf-scholz added 11 commits January 2, 2025 23:05

added type hints to lazy_property

ef2329f

added type hints to lazy_properties & simplified 'torch.Tensor' hints…

0937a2a

… to 'Tensor'

fixed typing errors

61366fe

replace 'Dict' with 'dict'

1fc1bd9

replace 'Tuple' with 'tuple'

aa3c769

replace 'Type' with 'type'

dac5f4d

fixed typing of transform._inv

a9f72e1

fixed circular import

9e9a1bb

added unit test

092be0b

lintrunner fixes

def92d5

added type hints for __init__ signatures

6e54911

pytorch-bot bot added the module: typing Related to mypy type annotations label Jan 4, 2025

pytorch-bot bot added the module: distributions Related to torch.distributions label Jan 4, 2025

randolf-scholz mentioned this pull request Jan 8000 4, 2025

[typing] Add static type hints to torch.distributions. #144196

Closed

6 tasks

randolf-scholz added 2 commits January 4, 2025 19:29

replaced Any with more concrete type in Distribution.support

d568418

Merge branch 'annotate_dist_properties' into distributions_type_hint_…

6c5c878

…init

pytorchbot added the open source label Jan 4, 2025

randolf-scholz added 2 commits January 4, 2025 19:40

added second typevar to lazy_property

d250c91

Merge branch 'annotate_dist_properties' into distributions_type_hint_…

14c0db1

…init

pytorch-bot bot added the release notes: python_frontend python frontend release notes category label Jan 4, 2025

randolf-scholz added 6 commits January 5, 2025 14:59

use typing.Union for 3.9 compat

c5a4d02

wishart.py: use Tensor instead of torch.Tensor annotation

64d9269

merged

b2852f6

use typing.Union for 3.9 compat

9ca3df5

added missing type hints

dc1010d

added missing type hints and fixed type error

70b0634

randolf-scholz added 2 commits January 28, 2025 10:33

Merge branch 'main' into distributions_type_hint_init

9caeb08

lintrunner fixes

3924079

randolf-scholz commented Jan 28, 2025

View reviewed changes

use assert logits is not None instead of elif branch

212325e

randolf-scholz force-pushed the distributions_type_hint_init branch from fd46589 to 212325e Compare January 28, 2025 16:11

removed spurious newlines

0116bfc

randolf-scholz commented Jan 28, 2025

View reviewed changes

torch/distributions/transforms.py Outdated Show resolved Hide resolved

randolf-scholz requested a review from malfet January 28, 2025 16:32

Merge branch 'main' into distributions_type_hint_init

aafae2d

randolf-scholz requested a review from Skylion007 February 4, 2025 16:33

github-actions bot added the Stale label Apr 5, 2025

malfet approved these changes Apr 6, 2025

View reviewed changes

pytorchmergebot added the merging label Apr 6, 2025

pytorchmergebot added the Merged label Apr 6, 2025

pytorchmergebot closed this in 6c38b9b Apr 6, 2025

pytorchmergebot removed the merging label Apr 6, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[typing] Add type hints to `init` methods in `torch.distributions`. #144197

[typing] Add type hints to `init` methods in `torch.distributions`. #144197

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[typing] Add type hints to __init__ methods in torch.distributions. #144197

[typing] Add type hints to __init__ methods in torch.distributions. #144197

Uh oh!

Conversation

Uh oh!

Open Problems:

Notes

Remark

Footnotes

Uh oh!

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Merge started

Uh oh!

Uh oh!

[typing] Add type hints to `init` methods in `torch.distributions`. #144197

[typing] Add type hints to `init` methods in `torch.distributions`. #144197