Remove several functions and objects from PyMC root namespace

Author: ricardoV94

.github/workflows/tests.yml

@@ -60,6 +60,7 @@ jobs:
             tests/test_util.py
             tests/test_pytensorf.py
             tests/test_math.py
+            tests/test_root_namespace.py
             tests/backends/test_base.py
             tests/backends/test_ndarray.py
             tests/step_methods/hmc/test_hmc.py

docs/source/api.rst

@@ -16,6 +16,7 @@ API
    api/data
    api/ode
    api/logprob
+   api/stats
    api/tuning
    api/math
    api/pytensorf
@@ -43,11 +44,11 @@ Plots, stats and diagnostics are delegated to the
 library, a general purpose library for
 "exploratory analysis of Bayesian models".
 
-* Functions from the ``arviz.plots`` module are available through ``pymc.<function>`` or ``pymc.plots.<function>``,
-  but for their API documentation please refer to the :ref:`ArviZ documentation <arviz:plot_api>`.
+* Functions from the ``arviz_plots`` package are available through ``pymc.plots.<function>``,
+  but for their API documentation please refer to the :doc:`arviz_plots documentation <arviz_plots:index>`.
 
-* Functions from the ``arviz.stats`` module are available through ``pymc.<function>`` or ``pymc.stats.<function>``,
-  but for their API documentation please refer to the :ref:`ArviZ documentation <arviz:stats_api>`.
+* Functions from the ``arviz_stats`` package are available through ``pymc.stats.<function>``,
+  but for their API documentation please refer to the :doc:`arviz_stats documentation <arviz_stats:index>`.
 
 ArviZ is a dependency of PyMC and so, in addition to the locations described above,
 importing ArviZ and using ``arviz.<function>`` will also work without any extra installation.

docs/source/api/math.rst

@@ -2,115 +2,10 @@
 Math
 ====
 
-This submodule contains various mathematical functions. Most of them are imported directly
-from pytensor.tensor (see there for more details). Doing any kind of math with PyMC random
-variables, or defining custom likelihoods or priors requires you to use these PyTensor
-expressions rather than NumPy or Python code.
-
-.. currentmodule:: pymc
-
-Functions exposed in pymc namespace
------------------------------------
-.. autosummary::
-   :toctree: generated/
-
-   expand_packed_triangular
-   logit
-   invlogit
-   probit
-   invprobit
-   logaddexp
-   logsumexp
-
-
-Functions exposed in pymc.math
-------------------------------
+This submodule contains various mathematical functions. Most are re-exported directly from
+:mod:`pytensor.tensor` and :mod:`pytensor.tensor.linalg` (see there for full signatures and
+details). Doing any kind of math with PyMC random variables, or defining custom likelihoods
+or priors, requires you to use these PyTensor expressions rather than NumPy or Python code.
 
 .. automodule:: pymc.math
-.. autosummary::
-   :toctree: generated/
-
-   abs
-   prod
-   dot
-   eq
-   neq
-   ge
-   gt
-   le
-   lt
-   exp
-   log
-   sgn
-   sqr
-   sqrt
-   sum
-   ceil
-   floor
-   sin
-   sinh
-   arcsin
-   arcsinh
-   cos
-   cosh
-   arccos
-   arccosh
-   tan
-   tanh
-   arctan
-   arctanh
-   cumprod
-   cumsum
-   matmul
-   and_
-   broadcast_to
-   clip
-   concatenate
-   flatten
-   or_
-   stack
-   switch
-   where
-   flatten_list
-   constant
-   max
-   maximum
-   mean
-   min
-   minimum
-   round
-   erf
-   erfc
-   erfcinv
-   erfinv
-   log1pexp
-   log1mexp
-   logaddexp
-   logsumexp
-   logdiffexp
-   logit
-   invlogit
-   probit
-   invprobit
-   sigmoid
-   softmax
-   log_softmax
-   logbern
-   full
-   full_like
-   ones
-   ones_like
-   zeros
-   zeros_like
-   kronecker
-   cartesian
-   kron_dot
-   kron_solve_lower
-   kron_solve_upper
-   kron_diag
-   flat_outer
-   expand_packed_triangular
-   batched_diag
-   block_diagonal
-   matrix_inverse
-   logdet
+   :members:

docs/source/api/misc.rst

@@ -6,8 +6,12 @@ Other utils
 .. autosummary::
    :toctree: generated/
 
-   compute_log_likelihood
-   compute_log_prior
    find_constrained_prior
+
+.. currentmodule:: pymc.blocking
+
+.. autosummary::
+   :toctree: generated/
+
    DictToArrayBijection
    model_table

docs/source/api/stats.rst

@@ -0,0 +1,15 @@
+Stats
+*****
+
+.. currentmodule:: pymc.stats
+
+.. autosummary::
+   :toctree: generated/
+
+   compute_log_prior
+   compute_log_likelihood
+
+PyMC re-exports functions from the ``arviz_stats`` library under the ``pymc.stats``
+namespace, allowing functions like ``summary``, ``ess``, ``rhat``, ``loo`` etc. to be
+accessed as ``pymc.stats.<function>``. For the API documentation of those functions,
+see the :doc:`arviz_stats documentation <arviz_stats:index>`.

docs/source/api/vi.rst

@@ -11,14 +11,22 @@ Variational Inference
    ASVGD
    SVGD
    FullRankADVI
+   fit
+
+.. currentmodule:: pymc.variational
+
+.. autosummary::
+   :toctree: generated/
+
    ImplicitGradient
    Inference
    KLqp
-   fit
 
 Approximations
 --------------
 
+.. currentmodule:: pymc
+
 .. autosummary::
    :toctree: generated/
 
@@ -33,9 +41,15 @@ OPVI
 .. autosummary::
    :toctree: generated/
 
-   Approximation
    Group
 
+.. currentmodule:: pymc.variational
+
+.. autosummary::
+   :toctree: generated/
+
+   Approximation
+
 Operators
 ---------
 
@@ -48,11 +62,16 @@ Operators
 
 Special
 -------
-.. currentmodule:: pymc
+.. currentmodule:: pymc.variational
 .. autosummary::
    :toctree: generated/
 
    Stein
+
+.. currentmodule:: pymc
+.. autosummary::
+   :toctree: generated/
+
    adadelta
    adagrad
    adagrad_window

docs/source/conf.py

@@ -299,6 +299,8 @@
 # intersphinx configuration to ease linking arviz docs
 intersphinx_mapping = {
     "arviz": ("https://python.arviz.org/en/latest/", None),
+    "arviz_plots": ("https://python.arviz.org/projects/plots/en/latest/", None),
+    "arviz_stats": ("https://python.arviz.org/projects/stats/en/latest/", None),
     "pytensor": ("https://pytensor.readthedocs.io/en/latest/", None),
     "home": ("https://www.pymc.io", None),
     "pmx": ("https://www.pymc.io/projects/experimental/en/latest", None),

docs/source/contributing/developer_guide.md

@@ -525,7 +525,7 @@ Moreover, if we are limiting to the problem that we can solved most confidently
 ## Inference
 
 ### MCMC
-The ability for model instance to generate conditional logp and dlogp function enable one of the unique feature of PyMC - {class}`~pymc.CompoundStep` method.
+The ability for model instance to generate conditional logp and dlogp function enable one of the unique feature of PyMC - {class}`~pymc.step_methods.CompoundStep` method.
 On a conceptual level it is a Metropolis-within-Gibbs sampler.
 Users can specify different sampler for different RVs.
 Alternatively, it is implemented as yet another interceptor:

docs/source/learn/core_notebooks/model_comparison.ipynb

@@ -565,7 +565,7 @@
     "\n",
     "## Reference\n",
     "\n",
-    "[Gelman, A., Hwang, J., & Vehtari, A. (2014). Understanding predictive information criteria for Bayesian models. Statistics and Computing, 24(6), 997–1016.](https://doi.org/10.1007/s11222-013-9416-2)\n",
+    "[Gelman, A., Hwang, J., & Vehtari, A. (2014). Understanding predictive information criteria for Bayesian models. Statistics and Computing, 24(6), 997\u20131016.](https://doi.org/10.1007/s11222-013-9416-2)\n",
     "\n",
     "[Vehtari, A, Gelman, A, Gabry, J. (2016). Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC. Statistics and Computing](http://link.springer.com/article/10.1007/s11222-016-9696-4)"
    ]

docs/source/learn/usage_overview.rst

@@ -33,7 +33,7 @@ Models are defined using a context manager (``with`` statement). The model is sp
         beta = pm.Normal('beta', 0, sigma=1)
 
         # Linear combination of parameters
-        theta = pm.invlogit(alpha + beta * dose)
+        theta = pm.math.invlogit(alpha + beta * dose)
 
         # Model likelihood
         deaths = pm.Binomial('deaths', n=n, p=theta, observed=y)