Type overloading for return_inferencedata in sample

[thomasaarholt]

What is this PR about?
This PR adds typing support which lets a type checker (like mypy or vscode's pylance) infer that when when passing pm.sample(return_inferencedata=True), then the return type is InferenceData, and when return_inferencedata=False, it is MultiTrace. This helps later on when referring to methods or properties that only exist on one of the two types.

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 👇

Maintenance

Added typing support to let a type checker (like mypy or vscode's pylance) infer that when passing pm.sample(return_inferencedata=True), then the return type is InferenceData, and when return_inferencedata=False, it is MultiTrace.

---

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

[codecov]

Codecov Report

Merging #6709 (c3b9b20) into main (a617bf2) will decrease coverage by 1.86%.
The diff coverage is 71.42%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6709      +/-   ##
==========================================
- Coverage   89.57%   87.71%   -1.86%     
==========================================
  Files          95       95              
  Lines       16131    16141      +10     
==========================================
- Hits        14449    14158     -291     
- Misses       1682     1983     +301     

Impacted Files	Coverage Δ
pymc/sampling/mcmc.py	89.81% <71.42%> (-0.32%)	⬇️

... and 12 files with indirect coverage changes

[thomasaarholt]

Mypy is failing with

pymc/sampling/mcmc.py:333: error: Overloaded function signatures 1 and 2 overlap with incompatible return types

Pyright (pyright pymc/sampling/mcmc.py) does not raise an error, and I'm inclined to believe that it is more accurate than mypy for this. If pymc-devs disagree, I'm happy to take any advice.

[thomasaarholt]

I fixed mypy, looks like it couldn't infer the same that pyright could.

[michaelosthege]

I never knew how to do this, so thanks for showing!