Full static typing for `torch.distributions` #144219

randolf-scholz · 2025-01-05T18:59:21Z

Fixes #144196
Extends #144197 #144106 #144110

Open Problems /// LSP violations

mixture_same_family.py: cdf and log_prob violate LSP (argument named x instead of value).
- suggestion: Imo these kinds of methods should make use of positional-only parameters, at least in base classes.
exp_family.py: LSP problem with _log_normalizer (parent class requires (*natural_params: Tensor) -> Tensor, subclasses implement (a: Tensor, b: Tensor) -> Tensor).
- suggestion: change parent class signature to (natural_params: Tuple[Tensor, ...]) -> Tensor. While this is BC breaking, (a) this is a private method, i.e. implementation detail, and (b) no one other than torch seems to overwrite it
constraints.py: dependent_property: mypy does not apply the same special casing to subclasses of property as it does to property itself, hence the need for type: ignore[assignment] statements.
- affects: relaxed_bernoulli.py, relaxed_categorical.py, logistic_normal.py, log_normal.py, kumaraswamy.py, half_cauchy.py, half_normal.py, inverse_gamma.py, gumbel.py, weibull.py.
- suggestion: consider a construction similar to lazy_property in distributions/utils.
constraints.py public interface not usable as type hints.
- More crisp design would likely have one class per constraints, instead of using a mix of classes and instances.
- suggestion: Add 1 class per constraint in the public interface; this can be subclasses of the existing ones.
- As a workaround, I currently added a bunch of TypeAlias-variants, but that is likely not the best solution.
transforms.py: _InverseTransform.with_cache violates LSP.
- suggestion: change with_cache to return _InverseTransform.
test_distributions.py: One test uses Dist.arg_constraints.get, hence assumes arg_constraints is a class-attribute, but the base class Distribution defines it as a @property.
test_distributions.py: One test uses Dist.support.event_dim, hence assumes support is a class-attribute, but the base class Distribution defines it as a @property.
test_distributions.py: Multiple tests use dist.cdf(float), but the base class annotates cdf(Tensor) -> Tensor.
- suggestion: replace float values with tensors in test, unless floats should be officially supported. Note that floats are nonsensical for multivariate distributions, so supporting it would probably require introducing a subclass for univariate distributions.
test_distributions.py: Multiple tests use dist.log_prob(float), but the base class annotates log_prob(Tensor) -> Tensor.

Notes

__init__.py: use += instead of extends (ruff PYI056)
binomial.py: Allow float arguments in probs and logits (gets used in tests)
constraints.py: made _DependentProperty a generic class, and _DependentProperty.__call__ polymorphic.
constraint_registry.py: Made ConstraintRegistry.register a polymorphic method, checking that the factory is compatible with the constraint.
constraint_registry.py: Needed to add type: ignore comments to functions that try to register multiple different constraints at once.
- maybe split them up?
dirichlet.py: @once_differentiable is untyped, requires type: ignore[misc] comment.
dirichlet.py: ctx: Any could be replaced with ctx: FunctionContext, however, the type lacks the saved_tensors attribute.
distribution.py: Distribution._get_checked_instance Accessing "__init__" on an instance is unsound, requires type: ignore comment.
distribution.py: Changed support from Optional[Constraint] to Constraint (consistent with the existing docstring, and several functions in tests rely on this assumption)
exp_family.py: small update to ExponentialFamily.entropy to fix type error.
independent.py: fixed type bug in Independent.support.
multivariate_normal.py: Added type: ignore comments to _batch_mahalanobis caused by¹.
relaxed_bernoulli.py: Allow float temperature argument (used in tests)
relaxed_categorical.py: Allow float temperature argument (used in tests)
transforms.py: Needed to change ComposeTransform.__init__ signature to accept Sequence[Transform] rather than just list[Transform] (covariance!)
transformed_distribution.py: Needed to change TransformedDistribution.__init__ signature to accept Sequence[Transform] rather than just list[Transform] (covariance!)
transformed_distribution.py: TransformedDistribution.support is problematic, because the parent class defines it as @property but several subclasses define it as an attribute, violating LSP.
von_mises.py: fixed result type being initialized as float instead of Tensor.
von_mises.py: @torch.jit.script_if_tracing is untyped, requires type: ignore[misc] comment.
von_mises.py: Allow float loc and scale (used in tests)

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

torch.Size is not correctly typed, causing mypy to think Size + Size is tuple[int, ...] instead of Size, see https://github.com/pytorch/pytorch/issues/144218. ↩

pytorch-bot · 2025-01-05T18:59:25Z

🔗 Helpful Links

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

📄 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 2a8e2ec with merge base aec3ef1 ():
💚 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-05T20:58:47Z

@pytorchbot label "module: typing"

randolf-scholz · 2025-01-05T20:58:55Z

@pytorchbot label "module: distributions"

randolf-scholz · 2025-01-05T20:59:05Z

@pytorchbot label "release notes: python_frontend"

torch/distributions/__init__.py

randolf-scholz

8000

More accurate would be using some sort of polymorphic mapping like

class _KL_REGISTRY_TYPE(Protocol):
    def __iter__(self) -> tuple[type[Distribution], type[Distribution]]: ...
    def __getitem__(self, key: tuple[type[P], type[Q]], / ) -> _KL[P, Q]: ...
    def __setitem__(self, key: tuple[type[P], type[Q]], value: _KL[P, Q], /) -> None: ...
    def __delitem__(self, key: tuple[type[Distribution], type[Distribution]], /) -> None: ...
    def clear(self) -> None: ...

but likely this overcomplicates things unnecessarily

torch/distributions/uniform.py

Skylion007 · 2025-05-19T17:36:38Z

@randolf-scholz Still interested in merging this?

randolf-scholz · 2025-05-22T09:48:17Z

@Skylion007 Yes, this was quite a bit of work, and it would be a shame if it goes to waste...

As I wrote in my last comment and in the OP, there are a few remaining open problems that mostly stem from LSP violations. It would be good to get some feedback on these.

Also, @malfet suggested in the other PR (the one adding __init__ signatures), which was also rather large, to split it up into multiple PRs to lessen the review burden. However, I am not sure how to do that reasonably. Is there a tool to automate this? Maybe git-explode?

randolf-scholz · 2025-05-22T14:18:25Z

@Skylion007 I rebased onto main and squashed some of my commits

Skylion007 · 2025-05-22T17:10:18Z

torch/distributions/constraints.py

@@ -206,16 +214,19 @@ def support(self):

    def __init__(
        self,
-        fn: Optional[Callable[..., Any]] = None,
+        fn: Optional[Callable[[T], R]] = None,


Why not ParamSpec here? T isn't used elsewhere?

ParamSpec does not make sense here, since this is a property-decorator, and properties are usually not supposed to take arguments. (note that T is actually the type the property gets bound to)

torch/distributions/transformed_distribution.py

test/distributions/test_distributions.py

Skylion007

I have some nits, but this is definitely an improvement for what was there before.

Skylion007 · 2025-05-25T15:23:49Z

torch/distributions/kl.py

 __all__ = ["register_kl", "kl_divergence"]

+P = TypeVar("P", bound=Distribution)


We want these types to be public?

The codebase seems to be inconsistent with respect to that, but more often than not uses private variables. Personally, I strongly prefer non-private, because it makes hints that show up for instance with pylance more readable. Moreover, when support for 3.11 is dropped and PEP 695 is used there is really no reason anymore to use an underscore prefix.

But for the purposes of this PR, I am fine with changing it if necessary.

We're the ones setting a lot of the standards for type naming!

I personally like the _.

But "The names, they convey nothing!" :-D

DistributionT1 = TypeVar("DistributionT1", bound=Distribution) DistributionT2 = TypeVar("DistributionT2", bound=Distribution) DistributionT3 = TypeVar("DistributionT3", bound=Distribution) DistributionT4= TypeVar("DistributionT3", bound=Distribution)

and

DistributionBinaryFunc: Callable[[DistributionT1, DistributionT2], Tensor]

Skylion007 · 2025-05-25T15:25:42Z

torch/distributions/kl.py


-def register_kl(type_p, type_q):
+def register_kl(type_p: type[P], type_q: type[Q]) -> _KL_Decorator:


Hold on, when does P, and Q are different from P2, Q2, this feels like the PERFECT place for a static type check tha the function you are decorating matches here?

That would be nice, but I do not see how to do it currently, I think it requires HKTs. That's because below you have some cases where the same function gets decorated multiple times like

@register_kl(Normal, Beta) @register_kl(Normal, ContinuousBernoulli) @register_kl(Normal, Exponential) @register_kl(Normal, Gamma) @register_kl(Normal, Pareto) @register_kl(Normal, Uniform) def _kl_normal_infinity( p: Normal, q: Union[Beta, ContinuousBernoulli, Exponential, Gamma, Pareto, Uniform] ) -> Tensor: return _infinite_like(p.loc)

So, the fist decorator must return Callable[[Normal, Union[Beta, ContinuousBernoulli, Exponential, Gamma, Pareto, Uniform]], Tensor], and not just Callable[[Normal, Uniform], Tensor], otherwise the next decorator will cause a type error.

What would be ideal would be something like

class _KL_Decorator[P, Q](Protocol): def __call__[P2 :> P, Q2 :> Q](self, arg: _KL[P2, Q2], /) -> _KL[P2, Q2]: ... def register_kl(type_p: type[P], type_q: type[Q]) -> _KL_Decorator[P, Q]:

So that for instance the first @register_kl(Normal, Uniform) would produce a

class _KL_Decorator[Normal, Uniform](Protocol): def __call__[P2: Normal, Q2: Uniform](self, arg: _KL[P2, Q2], /) -> _KL[P2, Q2]: ...

Because then when this gets applied to Callable[[Normal, Uniform | Pareto |...], Tensor], it gives back Callable[[Normal, Uniform | Pareto |...], Tensor].

But this requires HKTs, which are currently not available in the python typing system.

Actually with mypy==1.15 similar issues crop up in the constraint_registry.py file, it seems one also needs to loosen the Factory type hint in a similar manner.

Skylion007 · 2025-05-25T15:27:16Z

torch/distributions/negative_binomial.py

@@ -83,7 +80,7 @@ def expand(self, batch_shape, _instance=None):
        new._validate_args = self._validate_args
        return new

-    def _new(self, *args, **kwargs):
+    def _new(self, *args: Any, **kwargs: Any) -> Tensor:


You might be able to do a typevaruple for *Args here

Possibly, but all this function does is call TensorBase.new which currently has an *args: Any overload. Really what would be needed here is the ability to reference the signature of another function, which is currently not a feature supported by the type system.

torch/distributions/transforms.py

torch/distributions/uniform.py

+ removed '## mypy: allow-untyped-defs' comments

+ removed unused 'type: ignore' comments

Co-authored-by: Aaron Gokaslan <aaronGokaslan@gmail.com>

Skylion007 · 2025-05-28T14:13:32Z

torch/distributions/constraints.py

@@ -735,3 +759,34 @@ def check(self, value):
 positive_definite = _PositiveDefinite()
 cat = _Cat
 stack = _Stack
+
+# Type aliases.


These are now all reexported and demand doc strings I think?

rec · 2025-05-29T11:17:04Z

@skylion: Sigh, the delta is greater than 2k lines, and this makes the "sanity check" test fail.

2025-05-27T18:27:27.4803708Z Your PR is 3049 LOC which is more than the 2000 maximum
2025-05-27T18:27:27.4804568Z allowed within PyTorch infra. PLease make sure to split up
2025-05-27T18:27:27.4805366Z your PR into smaller pieces that can be reviewed.
2025-05-27T18:27:27.4806075Z If you think that this rule should not apply to your PR,
2025-05-27T18:27:27.4806753Z please contact @albanD or @seemethere.

It's generally pretty easy to split typing pull requests.

Starting a full review now, don't split before that! 🙂

rec

Whew, that's a big one!

I read every line though.

Thanks for doing this all!

rec · 2025-05-29T11:14:38Z

test/distributions/test_constraints.py

+
+def build_constraint(
+    constraint_fn: Union[C, type[C]],
+    args: tuple,


What, we can just do that as a synonym for tuple[Any, ...]? I had gotten the impression that this wouldn't work with mypy?

This is a test, is it even being type checked at all?

Currently, the lintrunner is ignoring these, but I do check them locally because the runtime code helps debugging the annotations.

$ mypy torch/distributions/ test/distributions/ --warn-unused-ignores test/distributions/test_distributions.py:3677:19: error: Argument 1 to "cdf" of "TransformedDistribution" has incompatible type "float"; expected "Tensor" [arg-type] test/distributions/test_distributions.py:3687:19: error: Argument 1 to "cdf" of "TransformedDistribution" has incompatible type "float"; expected "Tensor" [arg-type] test/distributions/test_distributions.py:3697:19: error: Argument 1 to "cdf" of "TransformedDistribution" has incompatible type "float"; expected "Tensor" [arg-type] test/distributions/test_distributions.py:3707:19: error: Argument 1 to "cdf" of "TransformedDistribution" has incompatible type "float"; expected "Tensor" [arg-type] test/distributions/test_distributions.py:5225:41: error: Argument 1 to "log_prob" of "Gamma" has incompatible type "int"; expected "Tensor" [arg-type] test/distributions/test_distributions.py:5251:40: error: Argument 1 to "log_prob" of "Gamma" has incompatible type "int"; expected "Tensor" [arg-type] Found 6 errors in 1 file (checked 52 source files)

test/distributions/test_distributions.py

rec · 2025-05-29T11:19:03Z

test/distributions/test_distributions.py

    """
    Create
179B
s a pair of distributions `Dist` initialized to test each element of
    param with each other.
    """
    params1 = [torch.tensor([p] * len(p)) for p in params]
    params2 = [p.transpose(0, 1) for p in params1]
-    return Dist(*params1), Dist(*params2)
+    return Dist(*params1), Dist(*params2)  # type: ignore[arg-type]


Wait, why does this fail? It should understand that both params1 and params2 are Sequence[Tensor] and not have an issue?

mypy infers D as Distribution and Distribution.__init__ only expects batch_shape, event_shape and validate_args.

What we could do is make these arguments keyword-only in Distribution.__init__, then the error goes away. Probably a good idea from a design POV, but backward incompatible!

But I think maybe this is better for a follow-up PR.

rec · 2025-05-29T11:20:22Z

test/distributions/test_distributions.py

@@ -1266,7 +1287,9 @@ def _check_forward_ad(self, fn):
                torch.count_nonzero(fwAD.unpack_dual(dual_out).tangent).item(), 0
            )

-    def _check_log_prob(self, dist, asset_fn):
+    def _check_log_prob(
+        self, dist: Distribution, asset_fn: Callable[Concatenate[int, ...], None]


so... cool... I had no idea you could do that, it's obvious only in hindsight.

Defining partial function signatures can be really handy, I wish this was even better supported (for instance when writing Callback Protocols)

test/distributions/test_distributions.py

rec · 2025-05-29T12:27:38Z

torch/distributions/kl.py

 __all__ = ["register_kl", "kl_divergence"]

+P = TypeVar("P", bound=Distribution)


We're the ones setting a lot of the standards for type naming!

I personally like the _.

But "The names, they convey nothing!" :-D

DistributionT1 = TypeVar("DistributionT1", bound=Distribution) DistributionT2 = TypeVar("DistributionT2", bound=Distribution) DistributionT3 = TypeVar("DistributionT3", bound=Distribution) DistributionT4= TypeVar("DistributionT3", bound=Distribution)

and

DistributionBinaryFunc: Callable[[DistributionT1, DistributionT2], Tensor]

torch/distributions/kl.py

torch/distributions/negative_binomial.py

rec · 2025-05-29T12:31:30Z

torch/distributions/pareto.py

@@ -28,7 +30,14 @@ class Pareto(TransformedDistribution):
        alpha (float or Tensor): Shape parameter of the distribution
    """

-    arg_constraints = {"alpha": constraints.positive, "scale": constraints.positive}
+    arg_constraints: ClassVar[dict[str, Constraint]] = {


This could conceivably a TypedDict?

Yes, that would be a nice enhancement. Thinking about it, it should probably even be Final[ClassVar[SomeTypedDict]], but that would require running mypy with --python-version 3.13. (currently 3.11 in mypy.ini)

rec · 2025-05-29T12:32:35Z

torch/distributions/transforms.py

        if self._cache_size == cache_size:
            return self
        if type(self).__init__ is Transform.__init__:
            return type(self)(cache_size=cache_size)
        raise NotImplementedError(f"{type(self)}.with_cache is not implemented")

-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:


Why object over Any?

Any is for gradual types, really no good reason to use it here. The signature mirrors that of object.__eq__

randolf-scholz · 2025-05-29T15:45:18Z

@rec I implemented most of your suggestions.

So, splitting it up I think it would make most sense to first make a PR for constraints and constraints_registry, and
then distributions and transforms, since the annotations for other modules depend on these.

Regarding constraints, I added all these type-aliases at the bottom. The CI complained that these are not listed in __all__, so should I just add them to __all__?

Alternatively, one could just make the classes they are pointing to public, but it's not my decision to make.

rec · 2025-05-29T17:02:36Z

Well, this was extremely educational, with at least one head-slapper.

You really covered everything in your response and I also agree with your plan to split.

Regarding the type-aliases in constraints, IIRC linting requires them to be in __all__ exactly if they do not start with _, so it's up to you.

I think it's good to go!

Fixes #144196 Part of #144219 Pull Request resolved: #154712 Approved by: https://github.com/Skylion007

Fixes pytorch#144196 Part of pytorch#144219 Pull Request resolved: pytorch#154712 Approved by: https://github.com/Skylion007

This was referenced Jan 5, 2025

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

Closed

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

Closed

pytorchbot added the open source label Jan 5, 2025

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

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

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

randolf-scholz marked this pull request as ready for review January 5, 2025 22:05

cpuhrsch requested a review from Skylion007 January 7, 2025 06:33

cpuhrsch added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jan 7, 2025

malfet reviewed Jan 8, 2025

View reviewed changes

torch/distributions/__init__.py Show resolved Hide resolved

randolf-scholz commented Jan 28, 2025

View reviewed changes

randolf-scholz commented Apr 7, 2025

View reviewed changes

torch/distributions/uniform.py Show resolved Hide resolved

This comment was marked as outdated.

Sign in to view

Skylion007 requested a review from malfet May 19, 2025 17:36

randolf-scholz force-pushed the distributions_full_typing branch 2 times, most recently from fbd5731 to 1ffc7c1 Compare May 22, 2025 14:17

Skylion007 reviewed May 22, 2025

View reviewed changes

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

randolf-scholz commented May 23, 2025

View reviewed changes

test/distributions/test_distributions.py Show resolved Hide resolved

Skylion007 approved these changes May 25, 2025

View reviewed changes

randolf-scholz force-pushed the distributions_full_typing branch from 1b65b97 to e625cb4 Compare May 25, 2025 18:35

This comment was marked as outdated.

Sign in to view

torch/distributions/: type annotations

418ca68

+ removed '## mypy: allow-untyped-defs' comments

randolf-scholz and others added 11 commits May 26, 2025 19:53

test/distributions/: type annotations

ac03f47

lintrunner fixes

5faecd4

+ removed unused 'type: ignore' comments

__init__.py: use += instead of extends (ruff PYI056)

b8e2ee2

Update torch/distributions/transforms.py

dfe3d38

Co-authored-by: Aaron Gokaslan <aaronGokaslan@gmail.com>

allow tensor input in interval-constraints (batched usage)

6f45f31

updated constraint_registry

f6bcf98

rebase

1c9edd9

added type-ignore to test_support_attributes

da458d3

make arg_constraints a ClassVar when parameter independent

38ff8f1

Made a ClassVar for distributions where it is parameter independent.

480a5a0

added some missing attribute annotations + formatting

73f160e

randolf-scholz force-pushed the distributions_full_typing branch from e625cb4 to 73f160e Compare May 26, 2025 19:19

Skylion007 reviewed May 28, 2025

View reviewed changes

Skylion007 requested review from rec and benjaminglass1 May 28, 2025 19:45

rec reviewed May 29, 2025

View reviewed changes

randolf-scholz added 2 commits May 29, 2025 17:38

Merge branch 'pytorch:main' into distributions_full_typing

43aec29

added suggestions from review

8e5c0cd

randolf-scholz added 4 commits May 30, 2025 11:22

constraints.py: added attribute annotations

1bc21d1

constraints.py: made independent/mixturesamefamily generic

c50c246

constraints.py: added type-aliases to __all__

65c6b15

Merge branch 'pytorch:main' into distributions_full_typing

2a8e2ec

This was referenced May 30, 2025

Type hints for distributions/constraints #154711

Open

Type hints for distributions/utils #154712

Closed

pytorchmergebot pushed a commit that referenced this pull request May 30, 2025

Type hints for distributions/utils (#154712)

ba3f91a

Fixes #144196 Part of #144219 Pull Request resolved: #154712 Approved by: https://github.com/Skylion007

randolf-scholz mentioned this pull request Jun 1, 2025

distributions/constraints type annotations + public classes + some refactoring #154827

Draft

nWEIdia pushed a commit to nWEIdia/pytorch that referenced this pull request Jun 2, 2025

Type hints for distributions/utils (pytorch#154712)

aad9275

Fixes pytorch#144196 Part of pytorch#144219 Pull Request resolved: pytorch#154712 Approved by: https://github.com/Skylion007

		__all__ = ["register_kl", "kl_divergence"]

		P = TypeVar("P", bound=Distribution)


		def register_kl(type_p, type_q):
		def register_kl(type_p: type[P], type_q: type[Q]) -> _KL_Decorator:

Full static typing for torch.distributions #144219

Are you sure you want to change the base?

Full static typing for torch.distributions #144219

Uh oh!

Conversation

Uh oh!

Open Problems /// LSP violations

Notes

Footnotes

Uh oh!

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

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!

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!

This comment was marked as outdated.

Choose a reason for hiding this comment

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Full static typing for `torch.distributions` #144219

Full static typing for `torch.distributions` #144219