Move sampling code to a submodule

[michaelosthege]

What is this PR about?
In #6257 we already wanted to do this, but had no good way to keep the from pymc import sampling_jax imports working.

Today I found a simple solution for this, so this PR moves all *sampling*.py files into a sampling submodule.

I also moved the tests accordingly.

⚠ The commits must not be squashed, otherwise the git history of jax.py would get messed-up. ⚠

You can check the diff commit-by-commit to see that I didn't actually change +721-680 lines...

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• None

Bugfixes / New features

• None

Docs / Maintenance

• All sampling related modules were moved into a new pymc.sampling submodule to facilitate further refactoring.

[codecov]

Codecov Report

Merging #6268 (f18e61c) into main (9c313cb) will increase coverage by 0.19%.
The diff coverage is 97.21%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6268      +/-   ##
==========================================
+ Coverage   94.19%   94.38%   +0.19%     
==========================================
  Files         102      108       +6     
  Lines       21480    23829    +2349     
==========================================
+ Hits        20233    22491    +2258     
- Misses       1247     1338      +91     

Impacted Files	Coverage Δ
pymc/sampling/forward.py	95.55% <ø> (ø)
pymc/sampling/parallel.py	87.32% <ø> (ø)
pymc/tests/sampling/test_mcmc.py	98.75% <94.11%> (ø)
pymc/sampling/jax.py	97.19% <97.19%> (ø)
pymc/sampling/mcmc.py	90.03% <100.00%> (ø)
pymc/sampling_jax.py	100.00% <100.00%> (+2.81%)	⬆️
pymc/smc/kernels.py	97.37% <100.00%> (ø)
pymc/smc/sampling.py	83.22% <100.00%> (ø)
pymc/tests/distributions/test_mixture.py	99.35% <100.00%> (ø)
pymc/tests/distributions/test_multivariate.py	99.44% <100.00%> (ø)
... and 10 more

[Armavica]

I only have two comments:

• There needs to be an empty __init__.py file inside tests/sampling
• There are a few places, mainly in the tests, where you import from pm.sampling.mcmc or pm.sampling.forward instead of just pm.sampling. This is not needed for the few items that are mentioned in __all__ of these two submodules, because they are reimported in pm.sampling. But it works this way too, I didn't know if you actually wanted to be more precise, and I don't have a strong opinion on this question.

[michaelosthege]

thanks, @Armavica ! I'll add the init file.
For the import routes that I went with the exact ones because it makes refactoring a little easier.
IMO we should try to eventually arrive at a point where each submodule has an __all__-advertised public API and only those functions are used by other components. Once we get there I agree that we should import from the init files.

@ricardoV94 want to add something too before I make the "apply review feedback commit"?

[ricardoV94]

@ricardoV94 want to add something too before I make the "apply review feedback commit"?

LGTM, do you want a review already or still need to merge the suggestion?

[michaelosthege]

LGTM, do you want a review already or still need to merge the suggestion?

I already added the __init__.py file, this is ready to review. You can even merge if you like it ;)