Make VI compatible with JAX backend

[ferrine]

Description

Related Issue

• [x] Closes #7104
• [x] Related to https://github.com/pymc-devs/pytensor/issues/595

Checklist

• [x] Checked that the pre-commit linting/style checks pass
• [x] Included tests that prove the fix is effective or that the new feature works
• [ ] Added necessary documentation (docstrings and/or example notebooks)
• [ ] If you are a pro: each commit corresponds to a relevant logical change

Type of change

• [x] New feature / enhancement
• [x] Bug fix
• [ ] Documentation
• [ ] Maintenance
• [ ] Other (please specify):

---

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

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 91.11%. Comparing base (a06081e) to head (30a2d73).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7103      +/-   ##
==========================================
- Coverage   91.87%   91.11%   -0.76%     
==========================================
  Files         100      100              
  Lines       16874    16858      -16     
==========================================
- Hits        15503    15361     -142     
- Misses       1371     1497     +126     

Files	Coverage Δ
pymc/pytensorf.py	91.46% <100.00%> (+0.16%)	:arrow_up:
pymc/variational/approximations.py	80.09% <100.00%> (-10.41%)	:arrow_down:

... and 12 files with indirect coverage changes

[ferrine]

Windows tests seem to be very weird and can't reproduce it on a Linux machine, is shape inference platform dependent?

[ricardoV94]

windows behaves differently with regard to integers. Default type is int32, which sometimes causes problems due to some rewrite or check that doesn't expect that (shape in PyTensor is supposed to be int64)

Just a guess from previous experiences. I can have a look on my windows machine next week

[ferrine]

I see one of the issues got resolved with sort op recently. Any updates for Windows?

[ricardoV94]

Any updates for Windows?

I don't think anyone investigated the problem yet

[ferrine]

How about marking these tests as xfail then?

[ricardoV94]

How about marking these tests as xfail then?

Let me or someone investigate on a Windows machine. Seems like an important failure on Windows. In the meantime you can rebase and pin PyMC to the next PyTensor version to see if the current xfail can be removed?

[ferrine]

@ricardoV94 updated the dependency on pytensor and commented on one of the xfails in the tests. Hope windows tests get resolved with newer pytensor

[ferrine]

In addition, mypy started to complain about pytensor

[pymc/sampling/forward.py]
pymc/sampling/forward.py:201: error: No overload variant of "general_toposort" matches argument types "list[Variable[Any, Any]]", "Callable[[Any], Any]"
pymc/sampling/forward.py:201: note: Possible overload variants:
pymc/sampling/forward.py:201: note:     def [T <: Node] general_toposort(outputs: Iterable[T], deps: None, compute_deps_cache: Callable[[T], Union[OrderedSet, list[T], None]], deps_cache: Optional[dict[T, list[T]]], clients: Optional[dict[T, list[T]]]) -> list[T]
pymc/sampling/forward.py:201: note:     def [T <: Node] general_toposort(outputs: Iterable[T], deps: Callable[[T], Union[OrderedSet, list[T]]], compute_deps_cache: None, deps_cache: None, clients: Optional[dict[T, list[T]]]) -> list[T]

[ricardoV94]

@ferrine feel free to rebase, we have already bumped the dependency on main

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 91.11%. Comparing base (0ad689c) to head (30a2d73).

:exclamation: Current head 30a2d73 differs from pull request most recent head 994da6c. Consider uploading reports for the commit 994da6c to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7103      +/-   ##
==========================================
- Coverage   92.34%   91.11%   -1.23%     
==========================================
  Files         102      100       -2     
  Lines       17032    16858     -174     
==========================================
- Hits        15728    15361     -367     
- Misses       1304     1497     +193     

Files	Coverage Δ
pymc/pytensorf.py	91.46% <100.00%> (+0.23%)	:arrow_up:
pymc/variational/approximations.py	80.09% <100.00%> (-10.78%)	:arrow_down:

... and 60 files with indirect coverage changes

[ferrine]

@ferrine feel free to rebase, we have already bumped the dependency on main

Just rebased, let's see how it goes

[ferrine]

rebased the old PR to see if any issues got resolved