[typing] Add type hints to @property and @lazy_property in torch.distributions. (pytorch#144110)

Fixes pytorch#76772, pytorch#144196
Extends pytorch#144106

- added type annotations to `lazy_property`.
- added type annotation to all `@property` and `@lazy_property` inside `torch.distributions` module.
- added simply type-check unit test to ensure type inference is working.
- replaced deprecated annotations like `typing.List` with the corresponding counterpart.
- simplified `torch.Tensor` hints with plain `Tensor`, otherwise signatures can become very verbose.

Pull Request resolved: pytorch#144110
Approved by: https://github.com/Skylion007

Author: randolf-scholz

test/typing/pass/distributions.py

@@ -0,0 +1,11 @@
+from typing_extensions import assert_type
+
+import torch
+from torch import distributions, Tensor
+
+
+dist = distributions.Normal(0, 1)
+assert_type(dist.mean, Tensor)
+
+dist = distributions.MultivariateNormal(torch.zeros(2), torch.eye(2))
+assert_type(dist.covariance_matrix, Tensor)

torch/distributions/bernoulli.py

@@ -2,7 +2,7 @@
 from numbers import Number
 
 import torch
-from torch import nan
+from torch import nan, Tensor
 from torch.distributions import constraints
 from torch.distributions.exp_family import ExponentialFamily
 from torch.distributions.utils import (
@@ -76,29 +76,29 @@ def _new(self, *args, **kwargs):
         return self._param.new(*args, **kwargs)
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         return self.probs
 
     @property
-    def mode(self):
+    def mode(self) -> Tensor:
         mode = (self.probs >= 0.5).to(self.probs)
         mode[self.probs == 0.5] = nan
         return mode
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         return self.probs * (1 - self.probs)
 
     @lazy_property
-    def logits(self):
+    def logits(self) -> Tensor:
         return probs_to_logits(self.probs, is_binary=True)
 
     @lazy_property
-    def probs(self):
+    def probs(self) -> Tensor:
         return logits_to_probs(self.logits, is_binary=True)
 
     @property
-    def param_shape(self):
+    def param_shape(self) -> torch.Size:
         return self._param.size()
 
     def sample(self, sample_shape=torch.Size()):
@@ -125,7 +125,7 @@ def enumerate_support(self, expand=True):
         return values
 
     @property
-    def _natural_params(self):
+    def _natural_params(self) -> tuple[Tensor]:
         return (torch.logit(self.probs),)
 
     def _log_normalizer(self, x):

torch/distributions/beta.py

@@ -2,6 +2,7 @@
 from numbers import Number, Real
 
 import torch
+from torch import Tensor
 from torch.distributions import constraints
 from torch.distributions.dirichlet import Dirichlet
 from torch.distributions.exp_family import ExponentialFamily
@@ -62,19 +63,19 @@ def expand(self, batch_shape, _instance=None):
         return new
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         return self.concentration1 / (self.concentration1 + self.concentration0)
 
     @property
-    def mode(self):
+    def mode(self) -> Tensor:
         return self._dirichlet.mode[..., 0]
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         total = self.concentration1 + self.concentration0
         return self.concentration1 * self.concentration0 / (total.pow(2) * (total + 1))
 
-    def rsample(self, sample_shape: _size = ()) -> torch.Tensor:
+    def rsample(self, sample_shape: _size = ()) -> Tensor:
         return self._dirichlet.rsample(sample_shape).select(-1, 0)
 
     def log_prob(self, value):
@@ -87,23 +88,23 @@ def entropy(self):
         return self._dirichlet.entropy()
 
     @property
-    def concentration1(self):
+    def concentration1(self) -> Tensor:
         result = self._dirichlet.concentration[..., 0]
         if isinstance(result, Number):
             return torch.tensor([result])
         else:
             return result
 
     @property
-    def concentration0(self):
+    def concentration0(self) -> Tensor:
         result = self._dirichlet.concentration[..., 1]
         if isinstance(result, Number):
             return torch.tensor([result])
         else:
             return result
 
     @property
-    def _natural_params(self):
+    def _natural_params(self) -> tuple[Tensor, Tensor]:
         return (self.concentration1, self.concentration0)
 
     def _log_normalizer(self, x, y):

torch/distributions/binomial.py

@@ -1,5 +1,6 @@
 # mypy: allow-untyped-defs
 import torch
+from torch import Tensor
 from torch.distributions import constraints
 from torch.distributions.distribution import Distribution
 from torch.distributions.utils import (
@@ -92,27 +93,27 @@ def support(self):
         return constraints.integer_interval(0, self.total_count)
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         return self.total_count * self.probs
 
     @property
-    def mode(self):
+    def mode(self) -> Tensor:
         return ((self.total_count + 1) * self.probs).floor().clamp(max=self.total_count)
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         return self.total_count * self.probs * (1 - self.probs)
 
     @lazy_property
-    def logits(self):
+    def logits(self) -> Tensor:
         return probs_to_logits(self.probs, is_binary=True)
 
     @lazy_property
-    def probs(self):
+    def probs(self) -> Tensor:
         return logits_to_probs(self.logits, is_binary=True)
 
     @property
-    def param_shape(self):
+    def param_shape(self) -> torch.Size:
         return self._param.size()
 
     def sample(self, sample_shape=torch.Size()):

torch/distributions/categorical.py

@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
 import torch
-from torch import nan
+from torch import nan, Tensor
 from torch.distributions import constraints
 from torch.distributions.distribution import Distribution
 from torch.distributions.utils import lazy_property, logits_to_probs, probs_to_logits
@@ -94,19 +94,19 @@ def support(self):
         return constraints.integer_interval(0, self._num_events - 1)
 
     @lazy_property
-    def logits(self):
+    def logits(self) -> Tensor:
         return probs_to_logits(self.probs)
 
     @lazy_property
-    def probs(self):
+    def probs(self) -> Tensor:
         return logits_to_probs(self.logits)
 
     @property
-    def param_shape(self):
+    def param_shape(self) -> torch.Size:
         return self._param.size()
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         return torch.full(
             self._extended_shape(),
             nan,
@@ -115,11 +115,11 @@ def mean(self):
         )
 
     @property
-    def mode(self):
-        return self.probs.argmax(axis=-1)
+    def mode(self) -> Tensor:
+        return self.probs.argmax(dim=-1)
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         return torch.full(
             self._extended_shape(),
             nan,

torch/distributions/cauchy.py

@@ -3,7 +3,7 @@
 from numbers import Number
 
 import torch
-from torch import inf, nan
+from torch import inf, nan, Tensor
 from torch.distributions import constraints
 from torch.distributions.distribution import Distribution
 from torch.distributions.utils import broadcast_all
@@ -52,22 +52,22 @@ def expand(self, batch_shape, _instance=None):
         return new
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         return torch.full(
             self._extended_shape(), nan, dtype=self.loc.dtype, device=self.loc.device
         )
 
     @property
-    def mode(self):
+    def mode(self) -> Tensor:
         return self.loc
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         return torch.full(
             self._extended_shape(), inf, dtype=self.loc.dtype, device=self.loc.device
         )
 
-    def rsample(self, sample_shape: _size = torch.Size()) -> torch.Tensor:
+    def rsample(self, sample_shape: _size = torch.Size()) -> Tensor:
         shape = self._extended_shape(sample_shape)
         eps = self.loc.new(shape).cauchy_()
         return self.loc + eps * self.scale

torch/distributions/chi2.py

@@ -1,4 +1,5 @@
 # mypy: allow-untyped-defs
+from torch import Tensor
 from torch.distributions import constraints
 from torch.distributions.gamma import Gamma
 
@@ -31,5 +32,5 @@ def expand(self, batch_shape, _instance=None):
         return super().expand(batch_shape, new)
 
     @property
-    def df(self):
+    def df(self) -> Tensor:
         return self.concentration * 2

torch/distributions/constraints.py

@@ -119,13 +119,13 @@ def __init__(self, *, is_discrete=NotImplemented, event_dim=NotImplemented):
         super().__init__()
 
     @property
-    def is_discrete(self):
+    def is_discrete(self) -> bool:  # type: ignore[override]
         if self._is_discrete is NotImplemented:
             raise NotImplementedError(".is_discrete cannot be determined statically")
         return self._is_discrete
 
     @property
-    def event_dim(self):
+    def event_dim(self) -> int:  # type: ignore[override]
         if self._event_dim is NotImplemented:
             raise NotImplementedError(".event_dim cannot be determined statically")
         return self._event_dim
@@ -233,11 +233,11 @@ def __init__(self, base_constraint, reinterpreted_batch_ndims):
         super().__init__()
 
     @property
-    def is_discrete(self):
+    def is_discrete(self) -> bool:  # type: ignore[override]
         return self.base_constraint.is_discrete
 
     @property
-    def event_dim(self):
+    def event_dim(self) -> int:  # type: ignore[override]
         return self.base_constraint.event_dim + self.reinterpreted_batch_ndims
 
     def check(self, value):
@@ -599,11 +599,11 @@ def __init__(self, cseq, dim=0, lengths=None):
         super().__init__()
 
     @property
-    def is_discrete(self):
+    def is_discrete(self) -> bool:  # type: ignore[override]
         return any(c.is_discrete for c in self.cseq)
 
     @property
-    def event_dim(self):
+    def event_dim(self) -> int:  # type: ignore[override]
         return max(c.event_dim for c in self.cseq)
 
     def check(self, value):
@@ -631,11 +631,11 @@ def __init__(self, cseq, dim=0):
         super().__init__()
 
     @property
-    def is_discrete(self):
+    def is_discrete(self) -> bool:  # type: ignore[override]
         return any(c.is_discrete for c in self.cseq)
 
     @property
-    def event_dim(self):
+    def event_dim(self) -> int:  # type: ignore[override]
         dim = max(c.event_dim for c in self.cseq)
         if self.dim + dim < 0:
             dim += 1

torch/distributions/continuous_bernoulli.py

@@ -3,6 +3,7 @@
 from numbers import Number
 
 import torch
+from torch import Tensor
 from torch.distributions import constraints
 from torch.distributions.exp_family import ExponentialFamily
 from torch.distributions.utils import (
@@ -52,7 +53,7 @@ class ContinuousBernoulli(ExponentialFamily):
 
     def __init__(
         self, probs=None, logits=None, lims=(0.499, 0.501), validate_args=None
-    ):
+    ) -> None:
         if (probs is None) == (logits is None):
             raise ValueError(
                 "Either `probs` or `logits` must be specified, but not both."
@@ -127,7 +128,7 @@ def _cont_bern_log_norm(self):
         return torch.where(self._outside_unstable_region(), log_norm, taylor)
 
     @property
-    def mean(self):
+    def mean(self) -> Tensor:
         cut_probs = self._cut_probs()
         mus = cut_probs / (2.0 * cut_probs - 1.0) + 1.0 / (
             torch.log1p(-cut_probs) - torch.log(cut_probs)
@@ -137,11 +138,11 @@ def mean(self):
         return torch.where(self._outside_unstable_region(), mus, taylor)
 
     @property
-    def stddev(self):
+    def stddev(self) -> Tensor:
         return torch.sqrt(self.variance)
 
     @property
-    def variance(self):
+    def variance(self) -> Tensor:
         cut_probs = self._cut_probs()
         vars = cut_probs * (cut_probs - 1.0) / torch.pow(
             1.0 - 2.0 * cut_probs, 2
@@ -151,15 +152,15 @@ def variance(self):
         return torch.where(self._outside_unstable_region(), vars, taylor)
 
     @lazy_property
-    def logits(self):
+    def logits(self) -> Tensor:
         return probs_to_logits(self.probs, is_binary=True)
 
     @lazy_property
-    def probs(self):
+    def probs(self) -> Tensor:
         return clamp_probs(logits_to_probs(self.logits, is_binary=True))
 
     @property
-    def param_shape(self):
+    def param_shape(self) -> torch.Size:
         return self._param.size()
 
     def sample(self, sample_shape=torch.Size()):
@@ -168,7 +169,7 @@ def sample(self, sample_shape=torch.Size()):
         with torch.no_grad():
             return self.icdf(u)
 
-    def rsample(self, sample_shape: _size = torch.Size()) -> torch.Tensor:
+    def rsample(self, sample_shape: _size = torch.Size()) -> Tensor:
         shape = self._extended_shape(sample_shape)
         u = torch.rand(shape, dtype=self.probs.dtype, device=self.probs.device)
         return self.icdf(u)
@@ -220,7 +221,7 @@ def entropy(self):
         )
 
     @property
-    def _natural_params(self):
+    def _natural_params(self) -> tuple[Tensor]:
         return (self.logits,)
 
     def _log_normalizer(self, x):