From 3544ae1cd4eb7ad0ed2c76125a29d1a364fc3511 Mon Sep 17 00:00:00 2001 From: aloctavodia Date: Thu, 19 Aug 2021 08:21:36 +0300 Subject: [PATCH 1/4] bug fix assigned variable nuts --- pymc3/step_methods/hmc/base_hmc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index 915c6d4b6a..c5e9603a90 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/pymc3/step_methods/hmc/base_hmc.py @@ -101,8 +101,8 @@ def __init__( # XXX: If the dimensions of these terms change, the step size # dimension-scaling should change as well, no? test_point = self._model.initial_point - continuous_vars = [test_point[v.name] for v in self._model.cont_vars] - size = sum(v.size for v in continuous_vars) + nuts_vars = [test_point[v.name] for v in vars] + size = sum(v.size for v in nuts_vars) self.step_size = step_scale / (size ** 0.25) self.step_adapt = step_sizes.DualAverageAdaptation( From 8e8a4771d4290d343e3da5b8aad4cb49e9dc987f Mon Sep 17 00:00:00 2001 From: aloctavodia Date: Thu, 19 Aug 2021 08:45:00 +0300 Subject: [PATCH 2/4] Fix failing tests in `test_parallel_sampling` --- pymc3/tests/test_parallel_sampling.py | 36 +++++++++++---------------- 1 file changed, 15 insertions(+), 21 deletions(-) diff --git a/pymc3/tests/test_parallel_sampling.py b/pymc3/tests/test_parallel_sampling.py index d58604b93e..e3074b9711 100644 --- a/pymc3/tests/test_parallel_sampling.py +++ b/pymc3/tests/test_parallel_sampling.py @@ -25,6 +25,8 @@ import pymc3 as pm import pymc3.parallel_sampling as ps +from pymc3.aesaraf import floatX + def test_context(): with pm.Model(): @@ -83,15 +85,13 @@ def test_remote_pipe_closed(): pm.sample(step=step, mp_ctx="spawn", tune=2, draws=2, cores=2, chains=2) -@pytest.mark.xfail( - reason="Possibly the same issue described in https://github.com/pymc-devs/pymc3/pull/4701" -) +@pytest.mark.xfail(reason="Unclear") def test_abort(): with pm.Model() as model: a = pm.Normal("a", shape=1) - pm.HalfNormal("b") - step1 = pm.NUTS([a]) - step2 = pm.Metropolis([model["b_log__"]]) + b = pm.HalfNormal("b") + step1 = pm.NUTS([model.rvs_to_values[a]]) + step2 = pm.Metropolis([model.rvs_to_values[b]]) step = pm.CompoundStep([step1, step2]) @@ -104,7 +104,7 @@ def test_abort(): chain=3, seed=1, mp_ctx=ctx, - start={"a": np.array([1.0]), "b_log__": np.array(2.0)}, + start={"a": floatX(np.array([1.0])), "b_log__": floatX(np.array(2.0))}, step_method_pickled=None, ) proc.start() @@ -118,15 +118,12 @@ def test_abort(): proc.join() -@pytest.mark.xfail( - reason="Possibly the same issue described in https://github.com/pymc-devs/pymc3/pull/4701" -) def test_explicit_sample(): with pm.Model() as model: a = pm.Normal("a", shape=1) - pm.HalfNormal("b") - step1 = pm.NUTS([a]) - step2 = pm.Metropolis([model["b_log__"]]) + b = pm.HalfNormal("b") + step1 = pm.NUTS([model.rvs_to_values[a]]) + step2 = pm.Metropolis([model.rvs_to_values[b]]) step = pm.CompoundStep([step1, step2]) @@ -138,7 +135,7 @@ def test_explicit_sample(): chain=3, seed=1, mp_ctx=ctx, - start={"a": np.array([1.0]), "b_log__": np.array(2.0)}, + start={"a": floatX(np.array([1.0])), "b_log__": floatX(np.array(2.0))}, step_method_pickled=None, ) proc.start() @@ -153,19 +150,16 @@ def test_explicit_sample(): proc.join() -@pytest.mark.xfail( - reason="Possibly the same issue described in https://github.com/pymc-devs/pymc3/pull/4701" -) def test_iterator(): with pm.Model() as model: a = pm.Normal("a", shape=1) - pm.HalfNormal("b") - step1 = pm.NUTS([a]) - step2 = pm.Metropolis([model["b_log__"]]) + b = pm.HalfNormal("b") + step1 = pm.NUTS([model.rvs_to_values[a]]) + step2 = pm.Metropolis([model.rvs_to_values[b]]) step = pm.CompoundStep([step1, step2]) - start = {"a": np.array([1.0]), "b_log__": np.array(2.0)} + start = {"a": floatX(np.array([1.0])), "b_log__": floatX(np.array(2.0))} sampler = ps.ParallelSampler(10, 10, 3, 2, [2, 3, 4], [start] * 3, step, 0, False) with sampler: for draw in sampler: From 211165edf6a1dbb05ae7617a2019bb95a1251164 Mon Sep 17 00:00:00 2001 From: Eelke Spaak Date: Thu, 19 Aug 2021 10:32:02 +0200 Subject: [PATCH 3/4] pickle step method for ProcessAdapter when mp_method is spawn --- pymc3/parallel_sampling.py | 4 ++++ pymc3/tests/test_parallel_sampling.py | 32 ++++++++++++++++++++++----- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/pymc3/parallel_sampling.py b/pymc3/parallel_sampling.py index 5aa49071e3..106f6cd55f 100644 --- a/pymc3/parallel_sampling.py +++ b/pymc3/parallel_sampling.py @@ -257,6 +257,10 @@ def __init__( if step_method_pickled is not None: step_method_send = step_method_pickled else: + if mp_ctx.get_start_method() == "spawn": + raise ValueError( + "please provide a pre-pickled step method when multiprocessing start method is 'spawn'" + ) step_method_send = step_method self._process = mp_ctx.Process( diff --git a/pymc3/tests/test_parallel_sampling.py b/pymc3/tests/test_parallel_sampling.py index e3074b9711..a2d4ea6a2e 100644 --- a/pymc3/tests/test_parallel_sampling.py +++ b/pymc3/tests/test_parallel_sampling.py @@ -13,9 +13,11 @@ # limitations under the License. import multiprocessing import os +import platform import aesara import aesara.tensor as at +import cloudpickle import numpy as np import pytest @@ -86,7 +88,8 @@ def test_remote_pipe_closed(): @pytest.mark.xfail(reason="Unclear") -def test_abort(): +@pytest.mark.parametrize("mp_start_method", ["spawn", "fork"]) +def test_abort(mp_start_method): with pm.Model() as model: a = pm.Normal("a", shape=1) b = pm.HalfNormal("b") @@ -95,8 +98,16 @@ def test_abort(): step = pm.CompoundStep([step1, step2]) + # on Windows we cannot fork + if platform.system() == "Windows" and mp_start_method == "fork": + return + if mp_start_method == "spawn": + step_method_pickled = cloudpickle.dumps(step, protocol=-1) + else: + step_method_pickled = None + for abort in [False, True]: - ctx = multiprocessing.get_context() + ctx = multiprocessing.get_context(mp_start_method) proc = ps.ProcessAdapter( 10, 10, @@ -105,7 +116,7 @@ def test_abort(): seed=1, mp_ctx=ctx, start={"a": floatX(np.array([1.0])), "b_log__": floatX(np.array(2.0))}, - step_method_pickled=None, + step_method_pickled=step_method_pickled, ) proc.start() while True: @@ -118,7 +129,8 @@ def test_abort(): proc.join() -def test_explicit_sample(): +@pytest.mark.parametrize("mp_start_method", ["spawn", "fork"]) +def test_explicit_sample(mp_start_method): with pm.Model() as model: a = pm.Normal("a", shape=1) b = pm.HalfNormal("b") @@ -127,7 +139,15 @@ def test_explicit_sample(): step = pm.CompoundStep([step1, step2]) - ctx = multiprocessing.get_context() + # on Windows we cannot fork + if platform.system() == "Windows" and mp_start_method == "fork": + return + if mp_start_method == "spawn": + step_method_pickled = cloudpickle.dumps(step, protocol=-1) + else: + step_method_pickled = None + + ctx = multiprocessing.get_context(mp_start_method) proc = ps.ProcessAdapter( 10, 10, @@ -136,7 +156,7 @@ def test_explicit_sample(): seed=1, mp_ctx=ctx, start={"a": floatX(np.array([1.0])), "b_log__": floatX(np.array(2.0))}, - step_method_pickled=None, + step_method_pickled=step_method_pickled, ) proc.start() while True: From 6249cf5082e7adb9c62ac78debb1cc20b8de5d9d Mon Sep 17 00:00:00 2001 From: Ricardo Date: Thu, 19 Aug 2021 13:02:52 +0200 Subject: [PATCH 4/4] Skip flaky test `test_parallel_sampling.py::test_abort` --- pymc3/tests/test_parallel_sampling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc3/tests/test_parallel_sampling.py b/pymc3/tests/test_parallel_sampling.py index a2d4ea6a2e..8bdc3ca1dc 100644 --- a/pymc3/tests/test_parallel_sampling.py +++ b/pymc3/tests/test_parallel_sampling.py @@ -87,7 +87,7 @@ def test_remote_pipe_closed(): pm.sample(step=step, mp_ctx="spawn", tune=2, draws=2, cores=2, chains=2) -@pytest.mark.xfail(reason="Unclear") +@pytest.mark.skip(reason="Unclear") @pytest.mark.parametrize("mp_start_method", ["spawn", "fork"]) def test_abort(mp_start_method): with pm.Model() as model: