Merge branch 'pymc-devs:main' into implement-minimiser-for-INLA

Author: Michal-Novomestsky

conda-envs/environment-test.yml

@@ -12,7 +12,6 @@ dependencies:
   - numba
   - pytest
   - pytest-cov
-  - libgcc<15
   - pydantic>=2.0.0
   - preliz
   - pip

pymc_extras/linearmodel.py

@@ -2,10 +2,12 @@
 import pandas as pd
 import pymc as pm
 
+from sklearn.base import BaseEstimator
+
 from pymc_extras.model_builder import ModelBuilder
 
 
-class LinearModel(ModelBuilder):
+class LinearModel(ModelBuilder, BaseEstimator):
     def __init__(
         self, model_config: dict | None = None, sampler_config: dict | None = None, nsamples=100
     ):

tests/statespace/core/test_statespace_JAX.py

@@ -7,6 +7,7 @@
 import pytest
 
 from pymc.model.transform.optimization import freeze_dims_and_data
+from pymc.testing import mock_sample_setup_and_teardown
 
 from pymc_extras.statespace.utils.constants import (
     FILTER_OUTPUT_NAMES,
@@ -30,6 +31,8 @@
 nile = load_nile_test_data()
 ALL_SAMPLE_OUTPUTS = MATRIX_NAMES + FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES
 
+mock_pymc_sample = pytest.fixture(scope="session")(mock_sample_setup_and_teardown)
+
 
 @pytest.fixture(scope="session")
 def pymc_mod(ss_mod):
@@ -66,7 +69,7 @@ def exog_pymc_mod(exog_ss_mod, rng):
 
 
 @pytest.fixture(scope="session")
-def idata(pymc_mod, rng):
+def idata(pymc_mod, rng, mock_pymc_sample):
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         with pymc_mod:
@@ -88,7 +91,7 @@ def idata(pymc_mod, rng):
 
 
 @pytest.fixture(scope="session")
-def idata_exog(exog_pymc_mod, rng):
+def idata_exog(exog_pymc_mod, rng, mock_pymc_sample):
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")