Updated import statements and linting fixes

Author: fonnesbeck

pymc_extras/__init__.py

@@ -13,12 +13,12 @@
 #   limitations under the License.
 import logging
 
-from pymc_experimental import gp, statespace, utils
-from pymc_experimental.distributions import *
-from pymc_experimental.inference.fit import fit
-from pymc_experimental.model.marginal.marginal_model import MarginalModel, marginalize
-from pymc_experimental.model.model_api import as_model
-from pymc_experimental.version import __version__
+from pymc_extras import gp, statespace, utils
+from pymc_extras.distributions import *
+from pymc_extras.inference.fit import fit
+from pymc_extras.model.marginal.marginal_model import MarginalModel, marginalize
+from pymc_extras.model.model_api import as_model
+from pymc_extras.version import __version__
 
 _log = logging.getLogger("pmx")
 
@@ -27,3 +27,5 @@
     if len(_log.handlers) == 0:
         handler = logging.StreamHandler()
         _log.addHandler(handler)
+
+__all__ = ["fit", "MarginalModel", "marginalize", "as_model"]

pymc_extras/distributions/__init__.py

@@ -17,24 +17,24 @@
 Experimental probability distributions for stochastic nodes in PyMC.
 """
 
-from pymc_experimental.distributions.continuous import Chi, GenExtreme, Maxwell
-from pymc_experimental.distributions.discrete import (
+from pymc_extras.distributions.continuous import Chi, GenExtreme, Maxwell
+from pymc_extras.distributions.discrete import (
     BetaNegativeBinomial,
     GeneralizedPoisson,
     Skellam,
 )
-from pymc_experimental.distributions.histogram_utils import histogram_approximation
-from pymc_experimental.distributions.multivariate import R2D2M2CP
-from pymc_experimental.distributions.timeseries import DiscreteMarkovChain
+from pymc_extras.distributions.histogram_utils import histogram_approximation
+from pymc_extras.distributions.multivariate import R2D2M2CP
+from pymc_extras.distributions.timeseries import DiscreteMarkovChain
 
 __all__ = [
-    "BetaNegativeBinomial",
+    "Chi",
+    "Maxwell",
     "DiscreteMarkovChain",
     "GeneralizedPoisson",
+    "BetaNegativeBinomial",
     "GenExtreme",
     "R2D2M2CP",
     "Skellam",
     "histogram_approximation",
-    "Chi",
-    "Maxwell",
 ]

pymc_extras/gp/__init__.py

@@ -13,6 +13,6 @@
 #   limitations under the License.
 
 
-from pymc_experimental.gp.latent_approx import KarhunenLoeveExpansion, ProjectedProcess
+from pymc_extras.gp.latent_approx import KarhunenLoeveExpansion, ProjectedProcess
 
 __all__ = ["KarhunenLoeveExpansion", "ProjectedProcess"]

pymc_extras/inference/__init__.py

@@ -13,6 +13,6 @@
 #   limitations under the License.
 
 
-from pymc_experimental.inference.fit import fit
+from pymc_extras.inference.fit import fit
 
 __all__ = ["fit"]

pymc_extras/model/marginal/marginal_model.py

@@ -23,18 +23,18 @@
 
 __all__ = ["MarginalModel", "marginalize"]
 
-from pymc_experimental.distributions import DiscreteMarkovChain
-from pymc_experimental.model.marginal.distributions import (
-    MarginalDiscreteMarkovChainRV,
-    MarginalFiniteDiscreteRV,
-    get_domain_of_finite_discrete_rv,
-    reduce_batch_dependent_logps,
+from pymc_extras.distributions import DiscreteMarkovChain
+from pymc_extras.model.marginal.distributions import (
+    FiniteDiscreteMarginalRV,
+    MarginalRV,
+    MarginalizedRV,
 )
-from pymc_experimental.model.marginal.graph_analysis import (
-    find_conditional_dependent_rvs,
-    find_conditional_input_rvs,
-    is_conditional_dependent,
-    subgraph_batch_dim_connection,
+from pymc_extras.model.marginal.graph_analysis import (
+    get_support_axes,
+    get_support_shape,
+    get_support_size,
+    get_support_values,
+    get_variable_support,
 )
 
 ModelRVs = TensorVariable | Sequence[TensorVariable] | str | Sequence[str]

pymc_extras/statespace/__init__.py

@@ -1,7 +1,13 @@
-from pymc_experimental.statespace.core.compile import compile_statespace
-from pymc_experimental.statespace.models import structural
-from pymc_experimental.statespace.models.ETS import BayesianETS
-from pymc_experimental.statespace.models.SARIMAX import BayesianSARIMA
-from pymc_experimental.statespace.models.VARMAX import BayesianVARMAX
+from pymc_extras.statespace.core.compile import compile_statespace
+from pymc_extras.statespace.models import structural
+from pymc_extras.statespace.models.ETS import BayesianETS
+from pymc_extras.statespace.models.SARIMAX import BayesianSARIMA
+from pymc_extras.statespace.models.VARMAX import BayesianVARMAX
 
-__all__ = ["structural", "BayesianSARIMA", "BayesianVARMAX", "BayesianETS", "compile_statespace"]
+__all__ = [
+    "compile_statespace",
+    "structural",
+    "BayesianETS",
+    "BayesianSARIMA",
+    "BayesianVARMAX",
+]

pymc_extras/utils/__init__.py

@@ -13,8 +13,8 @@
 #   limitations under the License.
 
 
-from pymc_experimental.utils import prior, spline
-from pymc_experimental.utils.linear_cg import linear_cg
+from pymc_extras.utils import prior, spline
+from pymc_extras.utils.linear_cg import linear_cg
 
 __all__ = (
     "linear_cg",

tests/test_blackjax_smc.py

@@ -23,12 +23,23 @@
 jax = pytest.importorskip("jax")
 pytest.importorskip("blackjax")
 
-from pymc_experimental.inference.smc.sampling import (
-    arviz_from_particles,
-    blackjax_particles_from_pymc_population,
-    get_jaxified_loglikelihood,
-    get_jaxified_logprior,
-    sample_smc_blackjax,
+from pymc_extras.inference.smc.sampling import (
+    SMC_KERNEL_DEFAULTS,
+    SMC_KERNEL_NAMES,
+    SMC_KERNEL_PARAMS,
+    SMC_KERNEL_PARAMS_DEFAULTS,
+    SMC_KERNEL_PARAMS_NAMES,
+    SMC_KERNEL_PARAMS_TYPES,
+    SMC_KERNEL_TYPES,
+    SMC_KERNEL_VALID_PARAMS,
+    SMC_KERNEL_VALID_PARAMS_TYPES,
+    SMC_KERNELS,
+    SMC_KERNELS_PARAMS,
+    SMC_KERNELS_PARAMS_DEFAULTS,
+    SMC_KERNELS_PARAMS_NAMES,
+    SMC_KERNELS_PARAMS_TYPES,
+    SMC_KERNELS_VALID_PARAMS,
+    SMC_KERNELS_VALID_PARAMS_TYPES,
 )
 
 

tests/test_find_map.py

@@ -3,10 +3,10 @@
 import pytensor.tensor as pt
 import pytest
 
-from pymc_experimental.inference.find_map import (
-    GradientBackend,
+from pymc_extras.inference.find_map import (
     find_MAP,
-    scipy_optimize_funcs_from_loss,
+    find_MAP_scipy,
+    find_MAP_pytensor,
 )
 
 pytest.importorskip("jax")

tests/test_histogram_approximation.py

@@ -17,7 +17,7 @@
 import pymc as pm
 import pytest
 
-import pymc_experimental as pmx
+import pymc_extras as pmx
 
 
 @pytest.mark.parametrize("use_dask", [True, False], ids="dask={}".format)