Fix mypy checks and jax installation in CI pipeline

[michaelosthege]

What is this PR about?
Fixing the ❌ mypy job in our CI pipelines.

We were on an outdated mypy==0.990 which failed to even run correctly.
I upgraded to the currently latest version and proceeded to fix the errors it found.

In addition, I also fixed some more typing errors, thereby removing three files from the list of expected failures.

The jax installation was fixed by pinning jax and jaxlib to the same version 0.4.16 which is the jax version that got installed in recent CI runs (probably due to a pin in an upstream package?).

Closes #6940
Closes #6941

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 👇

Bugfixes

• Several type hints got fixed

Maintenance

• Fixed mypy job
• Fixed jax installation in CI environments

---

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

[codecov]

Codecov Report

Merging #6945 (3957f16) into main (602234b) will increase coverage by 0.00%.
The diff coverage is 90.62%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #6945   +/-   ##
=======================================
  Coverage   90.69%   90.70%           
=======================================
  Files         100      100           
  Lines       16833    16845   +12     
=======================================
+ Hits        15267    15279   +12     
  Misses       1566     1566           

Files	Coverage Δ
pymc/gp/hsgp_approx.py	92.38% <100.00%> (+0.07%)	⬆️
pymc/model/core.py	91.29% <100.00%> (ø)
pymc/model_graph.py	80.74% <100.00%> (+0.31%)	⬆️
pymc/sampling/jax.py	0.00% <ø> (ø)
pymc/util.py	80.50% <100.00%> (+0.08%)	⬆️
pymc/model/transform/basic.py	90.47% <75.00%> (-4.53%)	⬇️
pymc/variational/opvi.py	87.38% <75.00%> (-0.19%)	⬇️

... and 1 file with indirect coverage changes

[michaelosthege]

Ah nice, the mypy upgrade helps. Now it just fails because mypy has things to complain about.

I'll try to fix these locally

[michaelosthege]

This damn jax stuff 🤐

Going to give it its own environment yml now

[ricardoV94]

This damn jax stuff 🤐

Going to give it its own environment yml now

Great, one more :D

[ricardoV94]

Thanks @michaelosthege!