Fix bug in NUTS variable assignment

[aloctavodia]

This is a bug I found while working on bringing BART to V4 #4914. NUTS should use the variables passed to the sampler and not get them from the model. @ricardoV94 recommend having this fix in its own PR and separated from the BART PR.

[codecov]

Codecov Report

Merging #4952 (6249cf5) into main (3fb5807) will increase coverage by 0.05%.
The diff coverage is 75.00%.

@@            Coverage Diff             @@
##             main    #4952      +/-   ##
==========================================
+ Coverage   74.01%   74.06%   +0.05%     
==========================================
  Files          86       86              
  Lines       13860    13862       +2     
==========================================
+ Hits        10258    10267       +9     
+ Misses       3602     3595       -7     

Impacted Files	Coverage Δ
pymc3/parallel_sampling.py	87.29% <50.00%> (-0.93%)	⬇️
pymc3/step_methods/hmc/base_hmc.py	91.05% <100.00%> (ø)
pymc3/step_methods/hmc/integration.py	90.38% <0.00%> (-1.93%)	⬇️
pymc3/model.py	83.45% <0.00%> (-0.15%)	⬇️
pymc3/tests/models.py	87.76% <0.00%> (+8.63%)	⬆️

[ricardoV94]

We should add a minimal test. This seems to trigger it locally:

import pymc3 as pm
with pm.Model() as m:
    x = pm.Normal("x", 0, 1)
    y = pm.Normal("y", 0, 1)
    pm.sample(step=pm.Metropolis([m.rvs_to_values[x]]), chains=1, tune=10, draws=10)

[ricardoV94]

Getting this error on the Windows run it test_parallel_sampling.py::test_explicit_sample: AttributeError: 'functools.partial' object has no attribute '__name__'

https://github.com/pymc-devs/pymc3/pull/4952/checks?check_run_id=3368890696#step:7:75

@Spaak, @michaelosthege do you happen to have any idea? These tests are all related to multiprocessing...

[Spaak]

Getting this error on the Windows run it test_parallel_sampling.py::test_explicit_sample: AttributeError: 'functools.partial' object has no attribute '__name__'

https://github.com/pymc-devs/pymc3/pull/4952/checks?check_run_id=3368890696#step:7:75

@Spaak, @michaelosthege do you happen to have any idea? These tests are all related to multiprocessing...

Yes this is definitely caused by pickling (probably a step method), required when multiprocessing uses spawn (which is the default on Windows; POSIX OSes use fork). Default pickle chokes on functools.partial objects, which is why we switched to cloudpickle (#4858). Could be I missed one/a few instances of pickling that still uses default pickle? I can have a look later today.

[ricardoV94]

Could be I missed one/a few instances of pickling that still uses default pickle? I can have a look later today.

You might have missed it because the tests were marked as xfail for other reasons (the one addressed in this PR)

[Spaak]

You might have missed it because the test was marked as an x-fail for other reasons (the one addressed in this PR)

Yes that's probably what happened. I'll work on a fix now, which I'll file as a separate PR (rather than push commits here) to keep the commit history on main clean (I don't yet know how much code I'll have to touch).

[ricardoV94]

You might have missed it because the test was marked as an x-fail for other reasons (the one addressed in this PR)

Yes that's probably what happened. I'll work on a fix now, which I'll file as a separate PR (rather than push commits here) to keep the commit history on main clean (I don't yet know how much code I'll have to touch).

Sure, but if you find it easier, feel free to force push a clean history to this PR.

The bugfix that is revelant to this PR is really just those two lines that @aloctavodia changed in pymc3/step_methods/hmc/base_hmc.py in the first commit

[Spaak]

Apologies; pre-commit has made me lazy and I forgot I was working in an environment that I hadn't installed pre-commit in yet. (So I pushed a commit here that failed pre-commit CI.) I'll force-push the same changes (but now properly blacked/isorted etc.) here in a minute. (So no separate PR; it turned out to be a minor change.)

[ricardoV94]

@Spaak any chances something similar was needed for this other failing test?

https://github.com/pymc-devs/pymc3/blob/e6ded51174f13ba0033c3cf51c6f3ed0ddc7c3b6/pymc3/tests/test_parallel_sampling.py#L88-L118

[Spaak]

@ricardoV94 I did make the same changes to test_abort, but it's still failing for another reason, so I left the xfail.

[ricardoV94]

@ricardoV94 I did make the same changes to test_abort, but it's still failing for another reason, so I left the xfail.

Didn't notice. Thanks.

One of the changed test seems to have failed though

[ricardoV94]

Oh well... that test_abort is still giving problems. Going to simply skip it until we figure it out. I did not change anything in my force push: https://github.com/pymc-devs/pymc3/compare/e721e11df92604d448601e97bbd0ad7677bd7735..b504f0c416c91c4f6b3334110b4e2703a1aaaa92

[Spaak]

Hmm after your force push the test_abort is xpassing: https://github.com/pymc-devs/pymc3/pull/4952/checks?check_run_id=3370627851 which was still xfail before. Did you rebase onto main or so causing it now to succeed..?

[ricardoV94]

It fails locally for me (regardless of floatX), so I think it is just a weird fluke

[Spaak]

Going to simply skip it until we figure it out.

Yep good idea

[ricardoV94]

Thanks @aloctavodia and @Spaak