Remove several functions and objects from PyMC root namespace

pymc/__init__.py

@@ -46,36 +46,19 @@ def __set_compiler_flags():
 
 __set_compiler_flags()
 
-from pymc import _version, gp, ode, sampling
-from pymc.backends import *
-from pymc.blocking import *
+from pymc import _version, gp, ode, plots, sampling, stats
 from pymc.data import *
 from pymc.distributions import *
-from pymc.exceptions import *
 from pymc.func_utils import find_constrained_prior
 from pymc.logprob import *
-from pymc.math import (
-    expand_packed_triangular,
-    invlogit,
-    invprobit,
-    logaddexp,
-    logit,
-    logsumexp,
-    probit,
-)
 from pymc.model.core import *
 from pymc.model.transform.conditioning import do, observe
 from pymc.model_graph import model_to_graphviz, model_to_networkx
-from pymc.plots import *
-from pymc.printing import *
 from pymc.pytensorf import *
 from pymc.sampling import *
 from pymc.smc import *
-from pymc.stats import *
 from pymc.step_methods import *
 from pymc.tuning import *
-from pymc.util import drop_warning_stat
 from pymc.variational import *
-from pymc.vartypes import *
 
 __version__ = _version.get_versions()["version"]

pymc/backends/base.py

@@ -45,10 +45,6 @@
 logger = logging.getLogger(__name__)
 
 
-class BackendError(Exception):
-    pass
-
-
 class IBaseTrace(ABC, Sized):
     """Minimal interface needed to record and access draws and stats for one MCMC chain."""
 

pymc/blocking.py

@@ -37,9 +37,6 @@
 
 from typing_extensions import TypeAlias
 
-__all__ = ["DictToArrayBijection"]
-
-
 T = TypeVar("T")
 PointType: TypeAlias = Dict[str, np.ndarray]
 StatsDict: TypeAlias = Dict[str, Any]

pymc/data.py

@@ -36,11 +36,11 @@
 
 import pymc as pm
 
+from pymc.exceptions import ShapeError
 from pymc.pytensorf import convert_observed_data
 
 __all__ = [
     "get_data",
-    "GeneratorAdapter",
     "Minibatch",
     "Data",
     "ConstantData",
@@ -238,7 +238,7 @@
 
     if isinstance(value, np.ndarray) and dims is not None:
         if len(dims) != value.ndim:
-            raise pm.exceptions.ShapeError(
+            raise ShapeError(
                 "Invalid data shape. The rank of the dataset must match the " "length of `dims`.",
                 actual=value.shape,
                 expected=value.ndim,
@@ -446,7 +446,7 @@
     if isinstance(dims, str):
         dims = (dims,)
     if not (dims is None or len(dims) == x.ndim):
-        raise pm.exceptions.ShapeError(
+        raise ShapeError(
             "Length of `dims` must match the dimensions of the dataset.",
             actual=len(dims),
             expected=x.ndim,

pymc/exceptions.py

@@ -12,30 +12,11 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-__all__ = [
-    "SamplingError",
-    "IncorrectArgumentsError",
-    "TraceDirectoryError",
-    "ImputationWarning",
-    "ShapeWarning",
-    "ShapeError",
-]
-
 
 class SamplingError(RuntimeError):
     pass
 
 
-class IncorrectArgumentsError(ValueError):
-    pass
-
-
-class TraceDirectoryError(ValueError):
-    """Error from trying to load a trace from an incorrectly-structured directory,"""
-
-    pass
-
-
 class ImputationWarning(UserWarning):
     """Warning that there are missing values that will be imputed."""
 
@@ -86,3 +67,55 @@
 
 class BlockModelAccessError(RuntimeError):
     pass
+
+
+class ParallelSamplingError(Exception):
+    def __init__(self, message, chain):
+        super().__init__(message)
+        self._chain = chain
+
+
+class RemoteTraceback(Exception):
+    def __init__(self, tb):
+        self.tb = tb
+
+    def __str__(self):
+        return self.tb
+
+
+class VariationalInferenceError(Exception):
+    """Exception for VI specific cases"""
+
+
+class NotImplementedInference(VariationalInferenceError, NotImplementedError):
+    """Marking non functional parts of code"""
+
+
+class ExplicitInferenceError(VariationalInferenceError, TypeError):
+    """Exception for bad explicit inference"""
+
+
+class ParametrizationError(VariationalInferenceError, ValueError):
+    """Error raised in case of bad parametrization"""
+
+
+class GroupError(VariationalInferenceError, TypeError):
+    """Error related to VI groups"""
+
+
+class IntegrationError(RuntimeError):
+    pass
+
+
+class PositiveDefiniteError(ValueError):
+    def __init__(self, msg, idx):
+        super().__init__(msg)
+        self.idx = idx
+        self.msg = msg
+
+    def __str__(self):
+        return f"Scaling is not positive definite: {self.msg}. Check indexes {self.idx}."
+
+
+class ParameterValueError(ValueError):
+    """Exception for invalid parameters values in logprob graphs"""

pymc/logprob/utils.py

@@ -63,6 +63,7 @@
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.variable import TensorVariable
 
+from pymc.exceptions import ParameterValueError
 from pymc.logprob.abstract import MeasurableVariable, _logprob
 from pymc.util import makeiter
 
@@ -231,10 +232,6 @@ def check_potential_measurability(
     return False
 
 
-class ParameterValueError(ValueError):
-    """Exception for invalid parameters values in logprob graphs"""
-
-
 class CheckParameterValue(CheckAndRaise):
     """Implements a parameter value check in a logprob graph.
 

pymc/model/core.py

@@ -59,13 +59,13 @@
 from pymc.exceptions import (
     BlockModelAccessError,
     ImputationWarning,
+    ParameterValueError,
     SamplingError,
     ShapeError,
     ShapeWarning,
 )
 from pymc.initial_point import make_initial_point_fn
 from pymc.logprob.basic import transformed_conditional_logp
-from pymc.logprob.utils import ParameterValueError
 from pymc.model_graph import model_to_graphviz
 from pymc.pytensorf import (
     PointFunc,

pymc/plots/__init__.py

@@ -20,7 +20,6 @@
 """
 import functools
 import sys
-import warnings
 
 import arviz as az
 
@@ -29,40 +28,3 @@
     obj = getattr(az.plots, attr)
     if not attr.startswith("__"):
         setattr(sys.modules[__name__], attr, obj)
-
-
-def alias_deprecation(func, alias: str):
-    original = func.__name__
-
-    @functools.wraps(func)
-    def wrapped(*args, **kwargs):
-        raise FutureWarning(
-            f"The function `{alias}` from PyMC was an alias for `{original}` from ArviZ. "
-            "It was removed in PyMC 4.0. "
-            f"Switch to `pymc.{original}` or `arviz.{original}`."
-        )
-
-    return wrapped
-
-
-# Aliases of ArviZ functions
-autocorrplot = alias_deprecation(az.plot_autocorr, alias="autocorrplot")
-forestplot = alias_deprecation(az.plot_forest, alias="forestplot")
-kdeplot = alias_deprecation(az.plot_kde, alias="kdeplot")
-energyplot = alias_deprecation(az.plot_energy, alias="energyplot")
-densityplot = alias_deprecation(az.plot_density, alias="densityplot")
-pairplot = alias_deprecation(az.plot_pair, alias="pairplot")
-traceplot = alias_deprecation(az.plot_trace, alias="traceplot")
-compareplot = alias_deprecation(az.plot_compare, alias="compareplot")
-
-
-__all__ = tuple(az.plots.__all__) + (
-    "autocorrplot",
-    "compareplot",
-    "forestplot",
-    "kdeplot",
-    "traceplot",
-    "energyplot",
-    "densityplot",
-    "pairplot",
-)

pymc/pytensorf.py

@@ -76,16 +76,9 @@
     "hessian",
     "hessian_diag",
     "inputvars",
-    "cont_inputs",
     "floatX",
     "intX",
-    "smartfloatX",
     "jacobian",
-    "CallableTensor",
-    "join_nonshared_inputs",
-    "make_shared_replacements",
-    "generator",
-    "convert_observed_data",
     "compile_pymc",
 ]
 

pymc/sampling/forward.py

@@ -53,9 +53,11 @@
 from pytensor.tensor.sharedvar import SharedVariable
 from typing_extensions import TypeAlias
 
-import pymc as pm
-
-from pymc.backends.arviz import _DefaultTrace
+from pymc.backends.arviz import (
+    _DefaultTrace,
+    predictions_to_inference_data,
+    to_inference_data,
+)
 from pymc.backends.base import MultiTrace
 from pymc.blocking import PointType
 from pymc.model import Model, modelcontext
@@ -438,7 +440,7 @@
     ikwargs: Dict[str, Any] = dict(model=model)
     if idata_kwargs:
         ikwargs.update(idata_kwargs)
-    return pm.to_inference_data(prior=prior, **ikwargs)
+    return to_inference_data(prior=prior, **ikwargs)
 
 
 def sample_posterior_predictive(
@@ -669,8 +671,8 @@
         if extend_inferencedata:
             ikwargs.setdefault("idata_orig", idata)
             ikwargs.setdefault("inplace", True)
-        return pm.predictions_to_inference_data(ppc_trace, **ikwargs)
-    idata_pp = pm.to_inference_data(posterior_predictive=ppc_trace, **ikwargs)
+        return predictions_to_inference_data(ppc_trace, **ikwargs)
+    idata_pp = to_inference_data(posterior_predictive=ppc_trace, **ikwargs)
 
     if extend_inferencedata and idata is not None:
         idata.extend(idata_pp)

pymc/sampling/mcmc.py

@@ -52,10 +52,11 @@
     coords_and_dims_for_inferencedata,
     find_constants,
     find_observations,
+    to_inference_data,
 )
 from pymc.backends.base import IBaseTrace, MultiTrace, _choose_chains
 from pymc.blocking import DictToArrayBijection
-from pymc.exceptions import SamplingError
+from pymc.exceptions import ParallelSamplingError, SamplingError
 from pymc.initial_point import PointType, StartDict, make_initial_point_fns_per_chain
 from pymc.model import Model, modelcontext
 from pymc.sampling.parallel import Draw, _cpu_count
@@ -892,7 +893,7 @@
     if compute_convergence_checks or return_inferencedata:
         ikwargs: Dict[str, Any] = dict(model=model, save_warmup=not discard_tuned_samples)
         ikwargs.update(idata_kwargs)
-        idata = pm.to_inference_data(mtrace, **ikwargs)
+        idata = to_inference_data(mtrace, **ikwargs)
 
         if compute_convergence_checks:
             warns = run_convergence_checks(idata, model)
@@ -1198,7 +1199,7 @@
                     if callback is not None:
                         callback(trace=strace, draw=draw)
 
-        except ps.ParallelSamplingError as error:
+        except ParallelSamplingError as error:
             strace = traces[error._chain]
             for strace in traces:
                 strace.close()

pymc/sampling/parallel.py

@@ -29,25 +29,13 @@
 from fastprogress.fastprogress import progress_bar
 
 from pymc.blocking import DictToArrayBijection
-from pymc.exceptions import SamplingError
+from pymc.exceptions import ParallelSamplingError, RemoteTraceback, SamplingError
 from pymc.util import RandomSeed
 
 logger = logging.getLogger(__name__)
 
 
-class ParallelSamplingError(Exception):
-    def __init__(self, message, chain):
-        super().__init__(message)
-        self._chain = chain
-
-
 # Taken from https://hg.python.org/cpython/rev/c4f92b597074
-class RemoteTraceback(Exception):
-    def __init__(self, tb):
-        self.tb = tb
-
-    def __str__(self):
-        return self.tb
 
 
 class ExceptionWithTraceback:

pymc/stats/__init__.py

@@ -28,5 +28,3 @@
         setattr(sys.modules[__name__], attr, obj)
 
 from pymc.stats.log_likelihood import compute_log_likelihood
-
-__all__ = ("compute_log_likelihood",) + tuple(az.stats.__all__)

pymc/step_methods/hmc/base_hmc.py

@@ -23,14 +23,14 @@
 import numpy as np
 
 from pymc.blocking import DictToArrayBijection, RaveledVars, StatsType
-from pymc.exceptions import SamplingError
+from pymc.exceptions import IntegrationError, SamplingError
 from pymc.model import Point, modelcontext
 from pymc.pytensorf import floatX
 from pymc.stats.convergence import SamplerWarning, WarningType
 from pymc.step_methods import step_sizes
 from pymc.step_methods.arraystep import GradientSharedStep
 from pymc.step_methods.hmc import integration
-from pymc.step_methods.hmc.integration import IntegrationError, State
+from pymc.step_methods.hmc.integration import State
 from pymc.step_methods.hmc.quadpotential import QuadPotentialDiagAdapt, quad_potential
 from pymc.tuning import guess_scaling
 from pymc.util import get_value_vars_from_user_vars