Make VI compatible with JAX backend #7103

ferrine · 2024-01-15T09:39:00Z

Description

Related Issue

Closes BUG: VI can't be used with Jax #7104
Related to FEAT: JAX Conversion for the given SortOp pytensor#595

Checklist

Checked that the pre-commit linting/style checks pass
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

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

codecov · 2024-01-15T09:53:07Z

Codecov Report

All modified and coverable lines are covered by tests ✅

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%)`	⬆️
pymc/variational/approximations.py	`80.09% <100.00%> (-10.41%)`	⬇️

... and 12 files with indirect coverage changes

pymc/sampling/jax.py

ricardoV94 · 2024-01-16T11:27:50Z

pymc/pytensorf.py

@@ -47,6 +46,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 from pytensor.scalar.basic import Cast
+from pytensor.scalar.basic import identity as scalar_identity


You don't need to create a new Elemwise, there's already one defined in tensor.math (or basic), just called tensor_copy

ricardoV94 · 2024-01-17T08:45:06Z

pymc/pytensorf.py

@@ -387,7 +386,7 @@ def hessian_diag(f, vars=None):
        return empty_gradient


-identity = Elemwise(scalar_identity, name="identity")
+identity = tensor_copy


Nitpick just import it directly in the VI module, no need to define it in pytensorf?

It might be used by someone else I assume

I don't think so, but even if we keep we should add a deprecatation warning

ferrine · 2024-01-22T08:05:34Z

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

ricardoV94 · 2024-01-22T10:09:59Z

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 · 2024-03-07T07:08:39Z

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

ricardoV94 · 2024-03-09T19:58:48Z

Any updates for Windows?

I don't think anyone investigated the problem yet

ferrine · 2024-03-09T20:44:43Z

How about marking these tests as xfail then?

ricardoV94 · 2024-03-13T09:36:29Z

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 · 2024-03-17T15:20:59Z

@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 · 2024-03-17T15:25:39Z

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 · 2024-03-27T21:41:31Z

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

fonnesbeck · 2024-03-29T02:48:50Z

tests/sampling/test_jax.py

+def test_vi_sampling_jax(method):
+    with pm.Model() as model:
+        x = pm.Normal("x")
+        pm.fit(10, method=method, fn_kwargs=dict(mode="JAX"))


To be consistent with pm.sample and the nuts_sampler= arg, should we have a dedicated argument for the VI backend instead of kwargs?

I vote yes, this API looks super weird.

What looks weird? This is the compilation mode, would be exactly the same if you wanted to use Numba or JAX for the PyMC nuts sampler or for prior/posterior predictive.

The only thing I would change is the name of fn_kwargs, which is called compile_kwargs I think in those other functions

Wouldn't this be what the user would have to do if they wanted to run VI on JAX?

I don't understand the question, this PR is just doing minor tweaks so the PyMC VI module can compile to JAX. It's not linking to specific JAX VI libraries.

We used this for sample_posterior_predictive for projects just last week, as we were sampling new variables that had heavy matmuls, went down from hours to minutes.

Great idea, should definitely add it there too.

pm.sample is still useful as you can sample discrete variables with JAX this way.

That makes sense, I'm not opposed to adding it there. Maybe we can add a warning that the sampler is still running Python and they likely will want to use nuts_sampler.

This is still doing python loops, it's exactly the same argument you need for pm.sample.

It's different than linking to a JAX VI library, which is what would be equivalent to the nuts_sampler kwarg that Chris mentioned in the first comment

This is still doing python loops, it's exactly the same argument you need for pm.sample.

Oh, I somehow assumed that VI was implemented mostly in PyTensor?

As for this, I'd prefer to focus this PR on backend compatibility and later address possible API changes in a new issue + PR. Agreed that there is inconsistency, we need to resolve that, but this will only defer the push to main with at least some working solution which went through many issues already.

Agreed @ferrine. My only suggestion is to switch fn_kwargs to compile_kwargs which we use in the other sample methods

codecov-commenter · 2024-05-01T09:24:33Z

Codecov Report

All modified and coverable lines are covered by tests ✅

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

❗ 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%)`	⬆️
pymc/variational/approximations.py	`80.09% <100.00%> (-10.78%)`	⬇️

... and 60 files with indirect coverage changes

ferrine · 2024-05-01T09:24:45Z

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

Just rebased, let's see how it goes

ferrine changed the title ~~add dispatch for identity Op, use static shapes for parameters~~ VI: add dispatch for identity Op, use static shapes for parameters Jan 15, 2024

ferrine added jax VI Variational Inference labels Jan 15, 2024

ricardoV94 reviewed Jan 15, 2024

View reviewed changes

pymc/sampling/jax.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jan 16, 2024

View reviewed changes

ricardoV94 changed the title ~~VI: add dispatch for identity Op, use static shapes for parameters~~ Make VI compatible with JAX backend Jan 16, 2024

ricardoV94 reviewed Jan 17, 2024

View reviewed changes

ferrine mentioned this pull request Jan 17, 2024

add alias for tensor_copy pymc-devs/pytensor#604

Open

ferrine force-pushed the variational-jax branch from 2330568 to 30a2d73 Compare March 17, 2024 15:19

fonnesbeck reviewed Mar 29, 2024

View reviewed changes

ferrine added 6 commits May 1, 2024 08:46

add dispatch for identity Op, use static shapes for parameters

bb28c39

add test

79fc937

add expected fail to remember about svgd

8c98a1f

replace pytensorf identity with pytensor identity

3d04679

use tensor copy instead of identity

28cde7d

update pytensor version, make xfail more elaborate

994da6c

ferrine force-pushed the variational-jax branch from 30a2d73 to 994da6c Compare May 1, 2024 09:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make VI compatible with JAX backend #7103

Make VI compatible with JAX backend #7103

ferrine commented Jan 15, 2024 •

edited by twiecki

codecov bot commented Jan 15, 2024 •

edited

ricardoV94 Jan 16, 2024

ricardoV94 Jan 17, 2024

ferrine Jan 17, 2024

ricardoV94 Jan 17, 2024 •

edited

ferrine commented Jan 22, 2024

ricardoV94 commented Jan 22, 2024

ferrine commented Mar 7, 2024

ricardoV94 commented Mar 9, 2024

ferrine commented Mar 9, 2024 •

edited

ricardoV94 commented Mar 13, 2024

ferrine commented Mar 17, 2024

ferrine commented Mar 17, 2024

ricardoV94 commented Mar 27, 2024

fonnesbeck Mar 29, 2024

twiecki Apr 1, 2024

ricardoV94 Apr 1, 2024

twiecki Apr 1, 2024

ricardoV94 Apr 1, 2024

twiecki Apr 1, 2024

ricardoV94 Apr 1, 2024 •

edited

twiecki Apr 1, 2024

ferrine May 1, 2024

ricardoV94 May 1, 2024

codecov-commenter commented May 1, 2024

ferrine commented May 1, 2024

Make VI compatible with JAX backend #7103

Are you sure you want to change the base?

Make VI compatible with JAX backend #7103

Conversation

ferrine commented Jan 15, 2024 • edited by twiecki

Description

Related Issue

Checklist

Type of change

codecov bot commented Jan 15, 2024 • edited

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Jan 17, 2024 • edited

Choose a reason for hiding this comment

ferrine commented Jan 22, 2024

ricardoV94 commented Jan 22, 2024

ferrine commented Mar 7, 2024

ricardoV94 commented Mar 9, 2024

ferrine commented Mar 9, 2024 • edited

ricardoV94 commented Mar 13, 2024

ferrine commented Mar 17, 2024

ferrine commented Mar 17, 2024

ricardoV94 commented Mar 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Apr 1, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov-commenter commented May 1, 2024

Codecov Report

ferrine commented May 1, 2024

ferrine commented Jan 15, 2024 •

edited by twiecki

codecov bot commented Jan 15, 2024 •

edited

ricardoV94 Jan 17, 2024 •

edited

ferrine commented Mar 9, 2024 •

edited

ricardoV94 Apr 1, 2024 •

edited