Implement and switch to lazy initval evaluation framework #4983

michaelosthege · 2021-09-04T18:32:17Z

Changes

The initial values are now evaluated lazily, all at once. Model.initial_values and Distribution(initval=...) can now take None, ndarray, Variable, "moment" or "prior".
Model.initial_values are now managed by RV instead of value var tensor
init_nuts signature changes
pm.sample() start kwarg change of name and meaning
Model.update_start_vals was removed (with informative error message)
All new initial point logic was moved to it's own module at initial_point.py

ToDo

Add additional tests for the new re-size flexibility that the lazy evaluation allows for. (test_initvals.py::TestInitvalEvaluation::test_initval_resizing)
~~Mention changes in the RELEASE-NOTES.md~~ Doing this in the HackMD document

Closes #4924
Closes #4484

codecov · 2021-09-04T19:18:31Z

Codecov Report

Merging #4983 (4e40756) into main (ab1178b) will increase coverage by 16.76%.
The diff coverage is 96.11%.

❗ Current head 4e40756 differs from pull request most recent head ff3b7f7. Consider uploading reports for the commit ff3b7f7 to get more accurate results

@@             Coverage Diff             @@
##             main    #4983       +/-   ##
===========================================
+ Coverage   61.45%   78.21%   +16.76%     
===========================================
  Files         130      131        +1     
  Lines       24461    24525       +64     
===========================================
+ Hits        15033    19183     +4150     
+ Misses       9428     5342     -4086

Impacted Files	Coverage Δ
pymc/variational/approximations.py	`29.83% <0.00%> (+0.85%)`	⬆️
pymc/sampling.py	`87.60% <90.00%> (+0.71%)`	⬆️
pymc/model.py	`83.18% <92.30%> (+0.26%)`	⬆️
pymc/tests/test_sampling.py	`97.47% <93.33%> (+0.18%)`	⬆️
pymc/distributions/continuous.py	`95.83% <100.00%> (-0.10%)`	⬇️
pymc/distributions/distribution.py	`94.47% <100.00%> (+0.38%)`	⬆️
pymc/initvalues.py	`100.00% <100.00%> (ø)`
pymc/tests/test_distributions.py	`96.34% <100.00%> (+0.43%)`	⬆️
pymc/tests/test_distributions_random.py	`86.28% <100.00%> (ø)`
pymc/tests/test_initvals.py	`100.00% <100.00%> (+8.53%)`	⬆️
... and 61 more

michaelosthege · 2021-09-04T20:29:28Z

Could the remaining test failure pymc3/tests/test_ndarray_backend.py::TestSaveLoad::test_sample_posterior_predictive be related to seeds/rng of the initial values?

pymc3/model.py

michaelosthege · 2021-09-05T13:23:18Z

I added two more tests - dependent initvals and resizing works.

But something is odd with the random state. This is from the failing test_ndarray_backend.py::TestSaveLoad::test_sample_posterior_predictive:

That backend is deprecated and the test doesn't seem to be very targeted. It looks like the MultiTrace constructor accesses Model.initial_point and this random states used for the initval evaluation actually do interfer with the models rng_seeder.
@ricardoV94 can you take a look at the evaluation method?

Also what exactly do we want in terms of rng for initial value evaluation? Should it be independent, optionally with their own rng, or should it re-use the models rng?

aseyboldt · 2021-09-21T21:50:27Z

~~The remaining test failures seem to be a result of #5007, but hopefully that would be fixed by #4887, so maybe we wait for that to be merged?~~

UPDATE: I marked them as XFAIL to unblock this PR.

pymc3/distributions/discrete.py

pymc3/distributions/distribution.py

pymc3/model.py

ricardoV94 · 2021-09-22T09:20:02Z

Good time to update the brittle model.update_start_vals

https://github.com/pymc-devs/pymc3/blob/4f8ad5d2772c8fd98b403d7c1d141e8f8394ba84/pymc3/model.py#L1594

There's more context here (with a different pseudo-API): #4924 (comment)

twiecki · 2021-10-13T20:37:36Z

LGTM. What happens when an init point is -inf?

ricardoV94 · 2021-10-13T20:46:17Z

LGTM. What happens when an init point is -inf?

What do you mean? The model logp at the initial point? If so, same as before

michaelosthege · 2021-10-13T20:58:53Z

Both pre-commit and docs worked on master. But that should be easy to fix(?)

ricardoV94 · 2021-10-14T08:43:36Z

Both pre-commit and docs worked on master. But that should be easy to fix(?)

Rebased after #5070. Should be fixed now

Related to pymc-devs#4924

With this commit "moment" or "prior" become legal initvals. Furthermore rv.tag.test_value is no longer assigned or used for initvals. The tolerance on test_mle_jacobian was eased to account for non- deterministic starting points of the optimization.

…odel This function can also handle variable specific jittering and user defined overrides The pm.sampling module was adapted to use the new functionality. This changed the signature of `init_nuts`: + `start` kwarg becomes `initvals` + `initvals` are required to be complete for all chains + `seeds` can now be specified for all chains

The test relied on monkey patching the jitter so that the model initial logp would fail predictably. This does not seem to be possible with the new numpy random generators, so a different test strategy has to be developed

To unblock this PR/branch from the aeppl integration.

michaelosthege

I can't approve since I'm the original author.

That one nitpick comment I made shouldn't hold us back from anything

michaelosthege · 2021-10-14T14:37:48Z

pymc/initial_point.py

+    rvs_to_jitter : set
+        The random variables for which jitter should be added.
+    """
+    # TODO: implement this


Created an issue for this: #5077

michaelosthege · 2021-10-14T14:43:12Z

pymc/initial_point.py

+    def find_rng_nodes(variables):
+        return [
+            node
+            for node in graph_inputs(variables)
+            if isinstance(
+                node,
+                (
+                    at.random.var.RandomStateSharedVariable,
+                    at.random.var.RandomGeneratorSharedVariable,
+                ),
+            )
+        ]


Can be extracted

twiecki · 2021-10-14T17:11:04Z

🥳

michaelosthege added the enhancements label Sep 4, 2021

michaelosthege marked this pull request as draft September 4, 2021 18:40

michaelosthege force-pushed the lazy-initval-evaluation branch from e119af0 to ec38672 Compare September 4, 2021 19:00

michaelosthege force-pushed the lazy-initval-evaluation branch from ec38672 to 2beeb04 Compare September 4, 2021 19:33

ricardoV94 reviewed Sep 5, 2021

View reviewed changes

pymc3/model.py Outdated Show resolved Hide resolved

michaelosthege force-pushed the lazy-initval-evaluation branch from 2beeb04 to b371760 Compare September 5, 2021 13:18

michaelosthege force-pushed the lazy-initval-evaluation branch 2 times, most recently from 94764d5 to 756acec Compare September 5, 2021 21:59

michaelosthege added help wanted needs info Additional information required pytensor labels Sep 13, 2021

michaelosthege force-pushed the lazy-initval-evaluation branch 3 times, most recently from 59f6c3c to 5f2ee98 Compare September 21, 2021 17:48

aseyboldt force-pushed the lazy-initval-evaluation branch 3 times, most recently from 4b25e3d to f5aa61a Compare September 21, 2021 20:58

aseyboldt mentioned this pull request Sep 21, 2021

Change default initial_point method to "moment" #5009

Closed

michaelosthege mentioned this pull request Sep 22, 2021

Restore support for default initvals #4867

Closed

ricardoV94 reviewed Sep 22, 2021

View reviewed changes

michaelosthege mentioned this pull request Sep 23, 2021

Fix Flat/HalfFlat initial value dimensionality #4996

Closed

michaelosthege force-pushed the lazy-initval-evaluation branch 3 times, most recently from f8c6562 to 40aa4fd Compare September 24, 2021 16:25

michaelosthege added this to the vNext (4.0.0) milestone Sep 27, 2021

ricardoV94 force-pushed the lazy-initval-evaluation branch 2 times, most recently from 65294ff to d28f897 Compare October 14, 2021 08:34

ricardoV94 mentioned this pull request Oct 14, 2021

Disable sampling of transformed variables from prior_predictive #5076

Closed

ricardoV94 force-pushed the lazy-initval-evaluation branch 3 times, most recently from ea49486 to 1f6e8f9 Compare October 14, 2021 09:05

ricardoV94 and others added 10 commits October 14, 2021 13:52

Update aesara dependency version

eff7e60

Evaluate initial values lazily

d2bb41c

Related to pymc-devs#4924

Manage initial values by RV var instead of RV value var

bf26f86

Stop caching initial points and wrap function creation

5ad9fa8

Deprecate Model.update_start_vals method

816c31d

Skip test_init_jitter

da05ab4

The test relied on monkey patching the jitter so that the model initial logp would fail predictably. This does not seem to be possible with the new numpy random generators, so a different test strategy has to be developed

Remove XFAIL mark on passing tests

6d678f4

XFAIL tests that depend on pymc-devs#5007

ff3b7f7

To unblock this PR/branch from the aeppl integration.

ricardoV94 force-pushed the lazy-initval-evaluation branch from 1f6e8f9 to ff3b7f7 Compare October 14, 2021 11:52

michaelosthege mentioned this pull request Oct 14, 2021

Implement a standardized interface for filtering RVs to jitter #5077

Closed

michaelosthege commented Oct 14, 2021

View reviewed changes

twiecki merged commit 3b763a1 into pymc-devs:main Oct 14, 2021

ricardoV94 mentioned this pull request Oct 14, 2021

Bring back distribution moments #5078

Closed

51 tasks

ferrine mentioned this pull request Oct 22, 2021

recompute_test_point fails after pickling #5090

Closed

michaelosthege mentioned this pull request Nov 10, 2021

Recursion problems in make_initial_point_expression #5168

Closed

michaelosthege deleted the lazy-initval-evaluation branch November 10, 2021 12:09

ricardoV94 mentioned this pull request Dec 8, 2021

Shape mismatch when using pm.Data #3631

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement and switch to lazy initval evaluation framework #4983

Implement and switch to lazy initval evaluation framework #4983

michaelosthege commented Sep 4, 2021 •

edited by ricardoV94

Loading

codecov bot commented Sep 4, 2021 •

edited

Loading

michaelosthege commented Sep 4, 2021

michaelosthege commented Sep 5, 2021

aseyboldt commented Sep 21, 2021 •

edited by michaelosthege

Loading

ricardoV94 commented Sep 22, 2021 •

edited

Loading

twiecki commented Oct 13, 2021

ricardoV94 commented Oct 13, 2021

michaelosthege commented Oct 13, 2021

ricardoV94 commented Oct 14, 2021

michaelosthege left a comment

michaelosthege Oct 14, 2021

michaelosthege Oct 14, 2021

twiecki commented Oct 14, 2021

Implement and switch to lazy initval evaluation framework #4983

Implement and switch to lazy initval evaluation framework #4983

Conversation

michaelosthege commented Sep 4, 2021 • edited by ricardoV94 Loading

Changes

ToDo

codecov bot commented Sep 4, 2021 • edited Loading

Codecov Report

michaelosthege commented Sep 4, 2021

michaelosthege commented Sep 5, 2021

aseyboldt commented Sep 21, 2021 • edited by michaelosthege Loading

ricardoV94 commented Sep 22, 2021 • edited Loading

twiecki commented Oct 13, 2021

ricardoV94 commented Oct 13, 2021

michaelosthege commented Oct 13, 2021

ricardoV94 commented Oct 14, 2021

michaelosthege left a comment

Choose a reason for hiding this comment

michaelosthege Oct 14, 2021

Choose a reason for hiding this comment

michaelosthege Oct 14, 2021

Choose a reason for hiding this comment

twiecki commented Oct 14, 2021

michaelosthege commented Sep 4, 2021 •

edited by ricardoV94

Loading

codecov bot commented Sep 4, 2021 •

edited

Loading

aseyboldt commented Sep 21, 2021 •

edited by michaelosthege

Loading

ricardoV94 commented Sep 22, 2021 •

edited

Loading