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 · 2024-07-09T08:41:39Z

rebased the old PR to see if any issues got resolved

ricardoV94 · 2024-07-09T09:09:39Z

pymc/pytensorf.py

-
-scalar_identity = IdentityOp(scalar.upgrade_to_float, name="scalar_identity")
-identity = Elemwise(scalar_identity, name="identity")
+identity = tensor_copy


Just do from pytensor... import tensor_copy as identity?

ricardoV94 · 2024-07-09T09:11:45Z

tests/sampling/test_jax.py

+
+TypeError: The broadcast pattern of the output of scan
+(Matrix(float64, shape=(?, 1))) is inconsistent with the one provided in `output_info`
+(Vector(float64, shape=(?,))). The output on axis 0 is `True`, but it is `False` on axis


This is actually something wrong, if the number of dimensions of a recurring output is different from the initial state. The difference between None and 1 is more annoying but this one looks like an error.

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 force-pushed the variational-jax branch from 30a2d73 to 994da6c Compare May 1, 2024 09:23

ferrine added 6 commits July 9, 2024 08:30

add dispatch for identity Op, use static shapes for parameters

581118a

add test

5544852

add expected fail to remember about svgd

44a75b0

replace pytensorf identity with pytensor identity

6474f9f

use tensor copy instead of identity

5dfd869

update pytensor version, make xfail more elaborate

829c6cb

ferrine force-pushed the variational-jax branch from 994da6c to 829c6cb Compare July 9, 2024 08:40

ricardoV94 reviewed Jul 9, 2024

View reviewed changes

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

Loading

codecov bot commented Jan 15, 2024 •

edited

Loading

ricardoV94 Jan 16, 2024

ricardoV94 Jan 17, 2024

ferrine Jan 17, 2024

ricardoV94 Jan 17, 2024 •

edited

Loading

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

Loading

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

Loading

twiecki Apr 1, 2024

ferrine May 1, 2024

ricardoV94 May 1, 2024

codecov-commenter commented May 1, 2024

ferrine commented May 1, 2024

ferrine commented Jul 9, 2024

ricardoV94 Jul 9, 2024

ricardoV94 Jul 9, 2024 •

edited

Loading

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 Loading

Description

Related Issue

Checklist

Type of change

codecov bot commented Jan 15, 2024 • edited Loading

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 Loading

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 Loading

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 Loading

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 Jul 9, 2024

Choose a reason for hiding this comment

ricardoV94 Jul 9, 2024 • edited Loading

Choose a reason for hiding this comment

ferrine commented Jan 15, 2024 •

edited by twiecki

Loading

codecov bot commented Jan 15, 2024 •

edited

Loading

ricardoV94 Jan 17, 2024 •

edited

Loading

ferrine commented Mar 9, 2024 •

edited

Loading

ricardoV94 Apr 1, 2024 •

edited

Loading

ricardoV94 Jul 9, 2024 •

edited

Loading