Move sampling code to a submodule

.github/workflows/tests.yml

@@ -59,16 +59,16 @@ jobs:
             pymc/tests/distributions/test_censored.py
             pymc/tests/distributions/test_simulator.py
             pymc/tests/distributions/test_truncated.py
-            pymc/tests/test_sampling_forward.py
+            pymc/tests/sampling/test_forward.py
             pymc/tests/stats/test_convergence.py
 
           - |
             pymc/tests/tuning/test_scaling.py
             pymc/tests/tuning/test_starting.py
-            pymc/tests/test_sampling.py
             pymc/tests/distributions/test_dist_math.py
             pymc/tests/distributions/test_transform.py
-            pymc/tests/test_parallel_sampling.py
+            pymc/tests/sampling/test_mcmc.py
+            pymc/tests/sampling/test_parallel.py
             pymc/tests/test_printing.py
 
           - |
@@ -150,8 +150,8 @@ jobs:
         test-subset:
           - pymc/tests/variational/test_approximations.py pymc/tests/variational/test_callbacks.py pymc/tests/variational/test_inference.py pymc/tests/variational/test_opvi.py pymc/tests/test_initial_point.py
           - pymc/tests/test_model.py pymc/tests/step_methods/test_compound.py pymc/tests/step_methods/hmc/test_hmc.py
-          - pymc/tests/gp/test_cov.py pymc/tests/gp/test_gp.py pymc/tests/gp/test_mean.py pymc/tests/gp/test_util.py pymc/tests/ode/test_ode.py pymc/tests/ode/test_utils.py pymc/tests/smc/test_smc.py pymc/tests/test_parallel_sampling.py
-          - pymc/tests/test_sampling.py pymc/tests/step_methods/test_metropolis.py pymc/tests/step_methods/test_slicer.py pymc/tests/step_methods/hmc/test_nuts.py
+          - pymc/tests/gp/test_cov.py pymc/tests/gp/test_gp.py pymc/tests/gp/test_mean.py pymc/tests/gp/test_util.py pymc/tests/ode/test_ode.py pymc/tests/ode/test_utils.py pymc/tests/smc/test_smc.py pymc/tests/sampling/test_parallel.py
+          - pymc/tests/sampling/test_mcmc.py pymc/tests/step_methods/test_metropolis.py pymc/tests/step_methods/test_slicer.py pymc/tests/step_methods/hmc/test_nuts.py
 
       fail-fast: false
     runs-on: ${{ matrix.os }}
@@ -221,12 +221,12 @@ jobs:
         python-version: ["3.9"]
         test-subset:
           - |
-            pymc/tests/test_parallel_sampling.py
+            pymc/tests/sampling/test_parallel.py
             pymc/tests/test_data.py
             pymc/tests/test_model.py
 
           - |
-            pymc/tests/test_sampling.py
+            pymc/tests/sampling/test_mcmc.py
 
           - |
             pymc/tests/backends/test_arviz.py
@@ -294,7 +294,7 @@ jobs:
         floatx: [float64]
         python-version: ["3.9"]
         test-subset:
-          - pymc/tests/test_sampling_jax.py
+          - pymc/tests/sampling/test_jax.py
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
@@ -363,7 +363,7 @@ jobs:
         floatx: [float32]
         python-version: ["3.10"]
         test-subset:
-        - pymc/tests/test_sampling.py pymc/tests/ode/test_ode.py pymc/tests/ode/test_utils.py
+        - pymc/tests/sampling/test_mcmc.py pymc/tests/ode/test_ode.py pymc/tests/ode/test_utils.py
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:

docs/source/api/samplers.rst

@@ -13,8 +13,8 @@ This submodule contains functions for MCMC and forward sampling.
    sample_prior_predictive
    sample_posterior_predictive
    sample_posterior_predictive_w
-   sampling_jax.sample_blackjax_nuts
-   sampling_jax.sample_numpyro_nuts
+   sampling.jax.sample_blackjax_nuts
+   sampling.jax.sample_numpyro_nuts
    iter_sample
    init_nuts
    draw

docs/source/contributing/build_docs.md

@@ -9,7 +9,7 @@ To build the docs, run these commands at PyMC repository root:
 
 ```bash
 pip install -r requirements-dev.txt  # Make sure the dev requirements are installed
-pip install numpyro  # Make sure `sampling_jax` docs can be built
+pip install numpyro  # Make sure `sampling/jax` docs can be built
 pip install -e .  # Install local pymc version as installable package
 make clean  # clean built docs from previous runs and intermediate outputs
 make html   # Build docs

pymc/__init__.py

@@ -68,7 +68,6 @@ def __set_compiler_flags():
 from pymc.plots import *
 from pymc.printing import *
 from pymc.sampling import *
-from pymc.sampling_forward import *
 from pymc.smc import *
 from pymc.stats import *
 from pymc.step_methods import *

pymc/sampling/__init__.py

@@ -0,0 +1,16 @@
+#   Copyright 2022 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from pymc.sampling.forward import *
+from pymc.sampling.mcmc import *