Fix type hints Python 3.9 compatibility

[michaelosthege]

What is this PR about?
Trying to fix one of the reasons why our CIs are going ❌.

Closes #6941

Checklist

• Explain important implementation details 👉 Literal | None doesn't work in Python 3.9
• 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

• Fixed a type hint incompatibility of the jax submodule under Python 3.9

---

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

[codecov]

Codecov Report

Merging #6944 (daa8a86) into main (602234b) will decrease coverage by 5.53%.
The diff coverage is n/a.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6944      +/-   ##
==========================================
- Coverage   90.69%   85.17%   -5.53%     
==========================================
  Files         100      100              
  Lines       16833    16833              
==========================================
- Hits        15267    14337     -930     
- Misses       1566     2496     +930     

Files	Coverage Δ
pymc/sampling/jax.py	0.00% <ø> (ø)

... and 21 files with indirect coverage changes

[ricardoV94]

Should we run pre-commit on the oldest supported python version? Would that have caught the problem?

[ricardoV94]

Sounds like we just need a with: https://github.com/pymc-devs/pytensor/blob/36df37981381dad62a24d46f1f21945b3b649c57/.github/workflows/test.yml#L59-L60

[michaelosthege]

Sounds like we just need a with: https://github.com/pymc-devs/pytensor/blob/36df37981381dad62a24d46f1f21945b3b649c57/.github/workflows/test.yml#L59-L60

Looks like we're running pre-commit multiple times then, and here it's already under Python 3.9:

pymc/.github/workflows/pre-commit.yml

Line 31 in 602234b

key: ${{ runner.os }}-py39-conda-${{ env.CACHE_NUMBER }}-${{

[michaelosthege]

jaxlib version 0.4.17 is newer than and incompatible with jax version 0.4.16. Please update your jax and/or jaxlib packages.

also this error again. It shouldn't be our job to pin jax's dependencies 😖

[ricardoV94]

Sounds like we just need a with: https://github.com/pymc-devs/pytensor/blob/36df37981381dad62a24d46f1f21945b3b649c57/.github/workflows/test.yml#L59-L60

Looks like we're running pre-commit multiple times then, and here it's already under Python 3.9:

pymc/.github/workflows/pre-commit.yml

Line 31 in 602234b

key: ${{ runner.os }}-py39-conda-${{ env.CACHE_NUMBER }}-${{

That's just the mypy job which apparently didn't pick the incompatible type hints. I think the pre-commit itself would have

[ricardoV94]

jaxlib version 0.4.17 is newer than and incompatible with jax version 0.4.16. Please update your jax and/or jaxlib packages.

also this error again. It shouldn't be our job to pin jax's dependencies 😖

Is that in our CI?

[michaelosthege]

closing in favor of #6945 to which I cherry-picked the commit