Merge branch 'pymc-devs:main' into model_builder_adaptations

Author: michaelraczycki

.github/workflows/test.yml

@@ -58,20 +58,20 @@ jobs:
           miniforge-variant: Mambaforge
           miniforge-version: latest
           mamba-version: "*"
-          activate-environment: pymc-test
+          activate-environment: pymc-experimental-test
           channel-priority: strict
           environment-file: conda-envs/environment-test.yml
           python-version: ${{matrix.python-version}}
           use-mamba: true
           use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
-      - name: Install-pymc
+      - name: Install pymc-experimental
         run: |
-          conda activate pymc-test
+          conda activate pymc-experimental-test
           pip install -e .
           python --version
       - name: Run tests
         run: |
-          conda activate pymc-test
+          conda activate pymc-experimental-test
           python -m pytest -vv --cov=pymc_experimental --cov-append --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
@@ -123,22 +123,22 @@ jobs:
           miniforge-variant: Mambaforge
           miniforge-version: latest
           mamba-version: "*"
-          activate-environment: pymc-test
+          activate-environment: pymc-experimental-test
           channel-priority: strict
           environment-file: conda-envs/windows-environment-test.yml
           python-version: ${{matrix.python-version}}
           use-mamba: true
           use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
-      - name: Install-pymc
+      - name: Install pymc-experimental
         run: |
-          conda activate pymc-test
+          conda activate pymc-experimental-test
           pip install -e .
           python --version
       - name: Run tests
         # This job uses a cmd shell, therefore the environment variable syntax is different!
         # The ">-" in the next line replaces newlines with spaces (see https://stackoverflow.com/a/66809682).
         run: >-
-          conda activate pymc-test &&
+          conda activate pymc-experimental-test &&
           python -m pytest -vv --cov=pymc_experimental --cov-append --cov-report=xml --cov-report term --durations=50 %TEST_SUBSET%
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2

conda-envs/environment-test.yml

@@ -1,4 +1,4 @@
-name: pymc-test
+name: pymc-experimental-test
 channels:
 - conda-forge
 - defaults
@@ -9,5 +9,5 @@ dependencies:
 - dask
 - xhistogram
 - pip:
-  - "git+https://github.com/pymc-devs/pymc.git@main"
+  - pymc>=5.1.2  # CI was failing to resolve
   - blackjax

conda-envs/windows-environment-test.yml

@@ -1,4 +1,4 @@
-name: pymc-test
+name: pymc-experimental-test
 channels:
 - conda-forge
 - defaults
@@ -9,4 +9,4 @@ dependencies:
 - dask
 - xhistogram
 - pip:
-  - "git+https://github.com/pymc-devs/pymc.git@main"
+  - pymc>=5.1.2  # CI was failing to resolve

pymc_experimental/marginal_model.py

@@ -93,8 +93,6 @@ def _delete_rv_mappings(self, rv: TensorVariable) -> None:
             self.rvs_to_initial_values.pop(rv)
         else:
             self.observed_RVs.remove(rv)
-        if rv in self.rvs_to_total_sizes:
-            self.rvs_to_total_sizes.pop(rv)
 
     def _transfer_rv_mappings(self, old_rv: TensorVariable, new_rv: TensorVariable) -> None:
         """Transfer model mappings from old_rv to new_rv"""
@@ -122,8 +120,6 @@ def _transfer_rv_mappings(self, old_rv: TensorVariable, new_rv: TensorVariable)
             index = self.observed_RVs.index(old_rv)
             self.observed_RVs.pop(index)
             self.observed_RVs.insert(index, new_rv)
-            if old_rv in self.rvs_to_total_sizes:
-                self.rvs_to_total_sizes[new_rv] = self.rvs_to_total_sizes.pop(old_rv)
 
     def _marginalize(self, user_warnings=False):
         fg = FunctionGraph(outputs=self.basic_RVs + self.marginalized_rvs, clone=False)
@@ -212,10 +208,6 @@ def clone(self):
         m.values_to_rvs = {i: vars_to_clone[rv] for i, rv in self.values_to_rvs.items()}
         m.rvs_to_values = {vars_to_clone[rv]: i for rv, i in self.rvs_to_values.items()}
         m.rvs_to_transforms = {vars_to_clone[rv]: i for rv, i in self.rvs_to_transforms.items()}
-        # Special logic due to bug in pm.Model
-        m.rvs_to_total_sizes = {
-            vars_to_clone[rv]: i for rv, i in self.rvs_to_total_sizes.items() if rv in vars_to_clone
-        }
         m.rvs_to_initial_values = {
             vars_to_clone[rv]: i for rv, i in self.rvs_to_initial_values.items()
         }

pymc_experimental/tests/distributions/test_continuous.py

@@ -21,7 +21,7 @@
 import scipy.stats.distributions as sp
 
 # test support imports from pymc
-from pymc.tests.distributions.util import (
+from pymc.testing import (
     BaseTestDistributionRandom,
     Domain,
     R,
@@ -30,8 +30,8 @@
     check_logcdf,
     check_logp,
     seeded_scipy_distribution_builder,
+    select_by_precision,
 )
-from pymc.tests.helpers import select_by_precision
 
 # the distributions to be tested
 from pymc_experimental.distributions import GenExtreme

pymc_experimental/tests/test_marginal_model.py

@@ -170,7 +170,7 @@ def test_nested_marginalized_rvs():
         sigma = pm.HalfNormal("sigma")
 
         idx = pm.Bernoulli("idx", p=0.75)
-        dep = pm.Normal("dep", mu=pt.switch(pt.eq(idx, 0), -1000, 1000), sigma=sigma)
+        dep = pm.Normal("dep", mu=pt.switch(pt.eq(idx, 0), -1000.0, 1000.0), sigma=sigma)
 
         sub_idx = pm.Bernoulli("sub_idx", p=pt.switch(pt.eq(idx, 0), 0.15, 0.95), shape=(5,))
         sub_dep = pm.Normal("sub_dep", mu=dep + sub_idx * 100, sigma=sigma, shape=(5,))

pymc_experimental/tests/test_model_builder.py

@@ -132,7 +132,7 @@ def test_predict():
     assert "y_model" in pred
     assert isinstance(pred, dict)
     assert len(prediction_data.input.values) == len(pred["y_model"])
-    assert isinstance(pred["y_model"][0], float)
+    assert isinstance(pred["y_model"][0], (np.float32, np.float64))
 
 
 def test_predict_posterior():

requirements.txt

@@ -1 +1 @@
-pymc>=5.0.0
+pymc>=5.1.2