Add type hints to dist() for discrete distributions

[johanbog]

What is this PR about?
Add type hints to the dist() class methods in the discrete distributions.
Issue #5358

Discussion points:

• I am not sure if the type hints are correct for the cutpoints parameter in OrderedProbit and OrderedLogistic
• I believe that the type hints should also include list[Union[int, float]] (or list[int], in the cases where ints are expected), but that would require changing the implementation of typing for the continuous distributions as well, so I can do that in a separate PR if it should be done.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

N/A

New features

N/A

Bugfixes

N/A

Documentation

N/A

Maintenance

N/A

---

📚 Documentation preview 📚: https://pymc--6937.org.readthedocs.build/en/6937/

[welcome]

💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.83%. Comparing base (1249c86) to head (fe39189).
Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #6937   +/-   ##
=======================================
  Coverage   92.83%   92.83%           
=======================================
  Files         106      106           
  Lines       17669    17675    +6     
=======================================
+ Hits        16403    16409    +6     
  Misses       1266     1266           

Files with missing lines	Coverage Δ
pymc/distributions/discrete.py	99.41% <100.00%> (+0.01%)	⬆️

[thomasaarholt]

I originally made some comments on this PR, which I then turned into a feature request that I am willing to work on, over at #7122.

[thomasaarholt]

Flyby review here, not a maintainer (though keen on pymc and type checking).
With the exception of one bit, this looks good to me.

[thomasaarholt]

The mypy error is going to be a pain to fix. Here is a minimum repro of the issue.

class Dist:
    def dist(self, *args: int):
        return 1

class Foo(Dist):
    def dist(self, foo: int, bar: int):
        return 1

run.py:6: error: Signature of "dist" incompatible with supertype "Dist"  [override]
run.py:6: note:      Superclass:
run.py:6: note:          def dist(self, *args: int) -> Any
run.py:6: note:      Subclass:
run.py:6: note:          def dist(self, foo: int, bar: int) -> Any

The issue is that in the Distribution class we have:

    @classmethod
    def dist(
        cls,
        dist_params, # <-- main issue is here
        *,
        shape: Shape | None = None,
        **kwargs,
    ) -> TensorVariable:

For the Binomial class we have multiple parameters where the superclass (Distribution) has one, and that one isn't even an *args. But that wouldn't work anyway.

    @classmethod
    def dist(
        cls,
        n: DISCRETE_DIST_PARAMETER_TYPES,
        p: DIST_PARAMETER_TYPES | None = None,
        logit_p: DIST_PARAMETER_TYPES | None = None,
        *args,
        **kwargs,
    ):

I am tempted to suggest that this PR won't be merged and that we focus on replacing the class-based system with a function-based one as per #5308.