Merge branch 'main' into ciguaran_fix_smc_bj

.github/workflows/pre-commit.yml

@@ -11,6 +11,6 @@ jobs:
     env:
       SKIP: no-commit-to-branch
     steps:
-    - uses: actions/checkout@v2
-    - uses: actions/setup-python@v2
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
     - uses: pre-commit/[email protected]

.github/workflows/pypi.yml

@@ -12,10 +12,10 @@ jobs:
     name: build source distribution
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Build the sdist and the wheel
@@ -28,7 +28,7 @@ jobs:
             cd test-sdist
             python -m venv venv-sdist
             venv-sdist/bin/python -m pip install numpy
-            venv-sdist/bin/python -m pip install ../dist/pymc-experimental*.tar.gz
+            venv-sdist/bin/python -m pip install ../dist/pymc_experimental*.tar.gz
             echo "Checking import and version number (on release)"
             venv-sdist/bin/python -c "import pymc_experimental as pmx; assert pmx.__version__ == '${{ github.ref_name }}'[1:] if '${{ github.ref_type }}' == 'tag' else True; print(pmx.__version__)"
             cd ..
@@ -61,7 +61,7 @@ jobs:
           user: __token__
           password: ${{ secrets.TEST_PYPI_API_TOKEN }}
           repository_url: https://test.pypi.org/legacy/
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Test pip install from test.pypi

.github/workflows/test.yml

@@ -18,7 +18,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.9"]
+        python-version: ["3.10"]
         test-subset:
           - pymc_experimental/tests
       fail-fast: false
@@ -28,49 +28,23 @@ jobs:
       PYTENSOR_FLAGS: gcc__cxxflags='-march=native'
     defaults:
       run:
-        shell: bash -l {0}
+        shell: bash -leo pipefail {0}
     steps:
-      - uses: actions/checkout@v2
-      - name: Cache conda
-        uses: actions/cache@v1
-        env:
-          # Increase this value to reset cache if environment-test.yml has not changed
-          CACHE_NUMBER: 0
+      - uses: actions/checkout@v4
+      - uses: mamba-org/setup-micromamba@v1
         with:
-          path: ~/conda_pkgs_dir
-          key: ${{ runner.os }}-py${{matrix.python-version}}-conda-${{ env.CACHE_NUMBER }}-${{
-            hashFiles('conda-envs/environment-test.yml') }}
-      - name: Cache multiple paths
-        uses: actions/cache@v2
-        env:
-          # Increase this value to reset cache if requirements.txt has not changed
-          CACHE_NUMBER: 0
-        with:
-          path: |
-            ~/.cache/pip
-            $RUNNER_TOOL_CACHE/Python/*
-            ~\AppData\Local\pip\Cache
-          key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{
-            hashFiles('requirements.txt') }}
-      - uses: conda-incubator/setup-miniconda@v2
-        with:
-          miniforge-variant: Mambaforge
-          miniforge-version: latest
-          mamba-version: "*"
-          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!
+          create-args: >-
+            python=${{matrix.python-version}}
+          environment-name: pymc-experimental-test
+          init-shell: bash
+          cache-environment: true
       - name: Install pymc-experimental
         run: |
-          conda activate pymc-experimental-test
           pip install -e .
           python --version
       - name: Run tests
         run: |
-          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
@@ -82,7 +56,7 @@ jobs:
     strategy:
       matrix:
         os: [windows-latest]
-        python-version: ["3.11"]
+        python-version: ["3.12"]
         test-subset:
           - pymc_experimental/tests
       fail-fast: false
@@ -92,51 +66,25 @@ jobs:
       PYTENSOR_FLAGS: gcc__cxxflags='-march=core2'
     defaults:
       run:
-        shell: cmd
+        shell: cmd /C call {0}
     steps:
-      - uses: actions/checkout@v2
-      - name: Cache conda
-        uses: actions/cache@v1
-        env:
-          # Increase this value to reset cache if conda-envs/windows-environment-test.yml has not changed
-          CACHE_NUMBER: 0
-        with:
-          path: ~/conda_pkgs_dir
-          key: ${{ runner.os }}-py${{matrix.python-version}}-conda-${{ env.CACHE_NUMBER }}-${{
-            hashFiles('conda-envs/windows-environment-test.yml') }}
-      - name: Cache multiple paths
-        uses: actions/cache@v2
-        env:
-          # Increase this value to reset cache if requirements.txt has not changed
-          CACHE_NUMBER: 0
-        with:
-          path: |
-            ~/.cache/pip
-            $RUNNER_TOOL_CACHE/Python/*
-            ~\AppData\Local\pip\Cache
-          key: ${{ runner.os }}-build-${{ matrix.python-version }}-${{
-            hashFiles('requirements.txt') }}
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: actions/checkout@v4
+      - uses: mamba-org/setup-micromamba@v1
         with:
-          miniforge-variant: Mambaforge
-          miniforge-version: latest
-          mamba-version: "*"
-          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!
+          create-args: >-
+            python=${{matrix.python-version}}
+          environment-name: pymc-experimental-test
+          init-shell: cmd.exe
+          cache-environment: true
       - name: Install pymc-experimental
         run: |
-          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-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

.readthedocs.yaml

@@ -5,7 +5,7 @@ version: 2
 build:
   os: ubuntu-20.04
   tools:
-    python: "3.9"
+    python: "3.10"
 
 python:
    install:

conda-envs/environment-test.yml

@@ -4,13 +4,12 @@ channels:
 - defaults
 dependencies:
 - pip
-
 - pytest-cov>=2.5
 - pytest>=3.0
 - dask
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.11.0  # CI was failing to resolve
+  - pymc>=5.13.0  # CI was failing to resolve
   - blackjax
   - scikit-learn

conda-envs/windows-environment-test.yml

@@ -10,6 +10,6 @@ dependencies:
 - xhistogram
 - statsmodels
 - pip:
-  - pymc>=5.11.0  # CI was failing to resolve
+  - pymc>=5.13.0  # CI was failing to resolve
   - blackjax
   - scikit-learn

pymc_experimental/distributions/discrete.py

@@ -20,6 +20,15 @@
 from pytensor.tensor.random.op import RandomVariable
 
 
+def log1mexp(x):
+    cond = x < np.log(0.5)
+    return np.piecewise(
+        x,
+        [cond, ~cond],
+        [lambda x: np.log1p(-np.exp(x)), lambda x: np.log(-np.expm1(x))],
+    )
+
+
 class GeneralizedPoissonRV(RandomVariable):
     name = "generalized_poisson"
     ndim_supp = 0
@@ -74,7 +83,7 @@ def _inverse_rng_fn(cls, rng, theta, lam, dist_size, idxs_mask):
                 log1p_lam_m_C = np.where(
                     pos_lam,
                     np.log1p(np.exp(abs_log_lam - log_c)),
-                    pm.math.log1mexp_numpy(abs_log_lam - log_c, negative_input=True),
+                    log1mexp(abs_log_lam - log_c),
                 )
                 log_p = log_c + log1p_lam_m_C * (x_ - 1) + log_p - np.log(x_) - lam
                 log_s = np.logaddexp(log_s, log_p)

pymc_experimental/distributions/multivariate/r2d2m2cp.py

@@ -418,9 +418,7 @@ def R2D2M2CP(
     *broadcast_dims, dim = dims
     input_sigma = pt.as_tensor(input_sigma)
     output_sigma = pt.as_tensor(output_sigma)
-    with pm.Model(name) as model:
-        if not all(isinstance(model.dim_lengths[d], pt.TensorConstant) for d in dims):
-            raise ValueError(f"{dims!r} should be constant length immutable dims")
+    with pm.Model(name):
         if r2_std is not None:
             r2 = pm.Beta("r2", mu=r2, sigma=r2_std, dims=broadcast_dims)
         phi = _phi(

pymc_experimental/distributions/timeseries.py

@@ -202,13 +202,13 @@ def transition(*args):
         discrete_mc_ = pt.moveaxis(pt.concatenate([init_dist_, markov_chain], axis=0), 0, -1)
 
         discrete_mc_op = DiscreteMarkovChainRV(
-            inputs=[P_, steps_, init_dist_],
+            inputs=[P_, steps_, init_dist_, state_rng],
             outputs=[state_next_rng, discrete_mc_],
             ndim_supp=1,
             n_lags=n_lags,
         )
 
-        discrete_mc = discrete_mc_op(P, steps, init_dist)
+        discrete_mc = discrete_mc_op(P, steps, init_dist, state_rng)
         return discrete_mc
 
 

pymc_experimental/linearmodel.py

@@ -69,8 +69,8 @@ def build_model(self, X: pd.DataFrame, y: pd.Series):
 
         # Data array size can change but number of dimensions must stay the same.
         with pm.Model() as self.model:
-            x = pm.MutableData("x", np.zeros((1,)), dims="observation")
-            y_data = pm.MutableData("y_data", np.zeros((1,)), dims="observation")
+            x = pm.Data("x", np.zeros((1,)), dims="observation")
+            y_data = pm.Data("y_data", np.zeros((1,)), dims="observation")
 
             # priors
             intercept = pm.Normal(

pymc_experimental/model/marginal_model.py

@@ -6,15 +6,15 @@
 import pytensor.tensor as pt
 from arviz import dict_to_dataset
 from pymc import SymbolicRandomVariable
-from pymc.backends.arviz import coords_and_dims_for_inferencedata
+from pymc.backends.arviz import coords_and_dims_for_inferencedata, dataset_to_point_list
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.transforms import Chain
 from pymc.logprob.abstract import _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
 from pymc.pytensorf import compile_pymc, constant_fold, inputvars
-from pymc.util import _get_seeds_per_chain, dataset_to_point_list, treedict
+from pymc.util import _get_seeds_per_chain, treedict
 from pytensor import Mode, scan
 from pytensor.compile import SharedVariable
 from pytensor.compile.builders import OpFromGraph
@@ -410,7 +410,7 @@ def transform_input(inputs):
                 marginalized_rv.type, dependent_logps
             )
 
-            rv_shape = constant_fold(tuple(marginalized_rv.shape))
+            rv_shape = constant_fold(tuple(marginalized_rv.shape), raise_not_constant=False)
             rv_domain = get_domain_of_finite_discrete_rv(marginalized_rv)
             rv_domain_tensor = pt.moveaxis(
                 pt.full(
@@ -579,6 +579,15 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs):
     return True
 
 
+from pytensor.graph.basic import graph_inputs
+
+
+def collect_shared_vars(outputs, blockers):
+    return [
+        inp for inp in graph_inputs(outputs, blockers=blockers) if isinstance(inp, SharedVariable)
+    ]
+
+
 def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs):
     # TODO: This should eventually be integrated in a more general routine that can
     #  identify other types of supported marginalization, of which finite discrete
@@ -621,27 +630,21 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     rvs_to_marginalize = [rv_to_marginalize, *dependent_rvs]
 
     outputs = rvs_to_marginalize
-    # Clone replace inner RV rng inputs so that we can be sure of the update order
-    # replace_inputs = {rng: rng.type() for rng in updates_rvs_to_marginalize.keys()}
-    # Clone replace outter RV inputs, so that their shared RNGs don't make it into
-    # the inner graph of the marginalized RVs
-    # FIXME: This shouldn't be needed!
-    replace_inputs = {}
-    replace_inputs.update({input_rv: input_rv.type() for input_rv in input_rvs})
-    cloned_outputs = clone_replace(outputs, replace=replace_inputs)
+    # We are strict about shared variables in SymbolicRandomVariables
+    inputs = input_rvs + collect_shared_vars(rvs_to_marginalize, blockers=input_rvs)
 
     if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
         marginalize_constructor = DiscreteMarginalMarkovChainRV
     else:
         marginalize_constructor = FiniteDiscreteMarginalRV
 
     marginalization_op = marginalize_constructor(
-        inputs=list(replace_inputs.values()),
-        outputs=cloned_outputs,
+        inputs=inputs,
+        outputs=outputs,
         ndim_supp=ndim_supp,
     )
 
-    marginalized_rvs = marginalization_op(*replace_inputs.keys())
+    marginalized_rvs = marginalization_op(*inputs)
     fgraph.replace_all(tuple(zip(rvs_to_marginalize, marginalized_rvs)))
     return rvs_to_marginalize, marginalized_rvs