Draft external sampler API #7880

ricardoV94 · 2025-07-30T13:34:31Z

Discussion needed !!

Draft PR motivated by #7699

The need being

Basically my question is: would PyMC be open to a PR along these lines? For public facing visibility, we'd like to put the algorithm into PyMC rather than use PyMC indirectly (i.e. extract a density from a probabilistic program in PyMC), and I thought this might be the best way.

The first idea discussed (and implemented here) was to allow pm.sample(step=ExternalSampler()) that defers sampling to an external library (including the ones we already supported before like nutpie, numpyro, ...)

import pymc as pm

with pm.Model() as m:
    x = pm.Normal("x")
    idata = pm.sample(nuts_sampler="nutpie", nuts_sampler_kwargs=kwargs)  # <- Before
    idata = pm.sample(step=pm.external.Nutpie(**kwargs))  # <- Now
    idata = pm.sample(step=pm.external.MCLMC(**kwargs))  # <- Future non NUTS methods can also be used

Pros: We are not assuming everything is a nuts_sampler, and there are objects / functions users can read to find the arguments that parametrize the samplers.

Cons: We're trying to put everything through pm.sample so it can reuse 6-8 arguments (tune, draws, chains, idata_kwargs ...) that already existed in pm.sample? Also tune/draws/chains may not make sense for some external samplers.

What I think is useful is to provide a standard API point to find external samplers (with our logic to connect pymc-library), which this PR kind of offers in pm.external.

The new ExternalSampler object is a bit awkward. It doesn't do much other other than allowing pm.sample to recognize it, so you can pass pm.sample(step=pm.external.Nutpie()), but then it has to arbitrarily split arguments between instantiation and sample, so as to make the sampler specific arguments discoverable, while reusing the few sampler-agnostic arguments.

Wouldn't it make more sense to just offer pm.external.sample_nutpie()?

IIRC this is why we went with pm.sample_smc instead of pm.sample(step=pm.SMC()) which used to exist before. It became awkward to have a function for both approaches.

📚 Documentation preview 📚: https://pymc--7880.org.readthedocs.build/en/7880/

jessegrabowski · 2025-07-30T14:37:42Z

Not against this proposal, but I haven't carefully looked at it yet. Wanted to chime in with an alternative right away. We could offer comprehensive developer docs for how to take a PyMC model and get the logp/dlogp out of it. This would encourage an ecosystem of packages like nutpie, that can work directly with model object.

If we wanted to then bridge to their thing directly via pm.whatever, it would be a thin wrapper just calling their API and maybe aligning defaults/argnames (again, like nutpie).

jessegrabowski · 2025-07-30T14:59:42Z

Looking over the actual PR, I like it much better than what we have now from a maintenance perspective. It also makes it easier to have specific parameters for each sampler, removing the horrible nuts_sampler_kwargs. We would need to have a depreciation period for the current API though?

I guess my proposal would be an "in addition to" this

ricardoV94 · 2025-07-30T16:20:48Z

We have a deprecation from now, everything should work except you get warnings.

One thing I don't love is the step name. Since these are holistic samplers. What about sampler or method?

codecov · 2025-07-30T18:04:41Z

Codecov Report

❌ Patch coverage is 44.06780% with 66 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.64%. Comparing base (58b49f2) to head (94d5ec4).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/sampling/external/nutpie.py	25.00%	36 Missing ⚠️
pymc/sampling/mcmc.py	61.53%	10 Missing ⚠️
pymc/sampling/external/base.py	43.75%	9 Missing ⚠️
pymc/sampling/external/jax.py	61.90%	8 Missing ⚠️
pymc/sampling/jax.py	0.00%	3 Missing ⚠️

❌ Your patch check has failed because the patch coverage (44.06%) is below the target coverage (50.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7880      +/-   ##
==========================================
- Coverage   92.94%   91.64%   -1.31%     
==========================================
  Files         116      120       +4     
  Lines       18845    18878      +33     
==========================================
- Hits        17516    17300     -216     
- Misses       1329     1578     +249

Files with missing lines	Coverage Δ
pymc/__init__.py	`100.00% <100.00%> (ø)`
pymc/sampling/external/__init__.py	`100.00% <100.00%> (ø)`
pymc/sampling/jax.py	`0.00% <0.00%> (-94.10%)`	⬇️
pymc/sampling/external/jax.py	`61.90% <61.90%> (ø)`
pymc/sampling/external/base.py	`43.75% <43.75%> (ø)`
pymc/sampling/mcmc.py	`90.26% <61.53%> (-1.12%)`	⬇️
pymc/sampling/external/nutpie.py	`25.00% <25.00%> (ø)`

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Dekermanjian · 2025-07-30T21:32:29Z

Wouldn't it make more sense to just offer pm.external.sample_nutpie()?

I personally prefer this because I think you end up with less dependencies that you have to worry about integrating nicely with each other.

For each new sampling algorithm you need to first implement the pm.external.sample_new_algorithm() and then you potentially also need to make sure that this new algorithm integrates nicely with the current pm.sample() implementation.

nataziel · 2025-07-31T03:50:02Z

pymc/sampling/external/__init__.py

+#   limitations under the License.
+from pymc.sampling.external.base import ExternalSampler
+from pymc.sampling.external.jax import Blackjax, Numpyro
+from pymc.sampling.external.nutpie import Nutpie


Suggested change

from pymc.sampling.external.nutpie import Nutpie

from pymc.sampling.external.nutpie import Nutpie

__all__ = ["Blackjax", "Numpyro", "Nutpie"]

You probably don't need to export the ABC right?

nataziel · 2025-07-31T04:56:11Z

pymc/sampling/jax.py

@@ -704,8 +666,8 @@ def sample_jax_nuts(
        dims.update(idata_kwargs.pop("dims"))

    # Use 'partial' to set default arguments before passing 'idata_kwargs'


unnecessary comment now

nataziel · 2025-07-31T05:04:55Z

pymc/sampling/external/base.py

+        compute_convergence_checks,
+        **kwargs,
+    ):
+        pass


Suggested change

pass

raise NotImplementedError

make it fail fast

nataziel · 2025-07-31T05:06:29Z

pymc/sampling/external/base.py

+    def sample(
+        self,
+        tune,
+        draws,
+        chains,
+        initvals,
+        random_seed,
+        progressbar,
+        var_names,
+        idata_kwargs,
+        compute_convergence_checks,
+        **kwargs,
+    ):


Golden opportunity to add type hints and return types early!

nataziel · 2025-07-31T05:12:16Z

pymc/sampling/external/jax.py

+
+
+class JAXSampler(ExternalSampler):
+    nuts_sampler = None  # Should be defined by subclass


This could be enforced by making an abstract method property that inheritors must implement, I've found that pattern quite useful

nataziel · 2025-07-31T05:14:04Z

pymc/sampling/external/jax.py

+from pymc.util import RandomState
+
+
+class JAXSampler(ExternalSampler):


Probably best not to assume that any JAXSampler is a NUTS sampler

nataziel · 2025-07-31T05:18:01Z

I know it's just a draft @ricardoV94, but popped some comments in where it made sense to me

jessegrabowski · 2025-07-31T06:44:11Z

We have a deprecation from now, everything should work except you get warnings.

One thing I don't love is the step name. Since these are holistic samplers. What about sampler or method?

What do you think about just calling e.g. pm.external.sample_nutpie directly, instead of passing something to pm.sample

ricardoV94 · 2025-07-31T09:54:43Z

As I wrote on the original issue I'm more inclined to that but I know some people love the pm.sample funnel

lucianopaz · 2025-08-01T11:45:35Z

I don't like using the step argument for this. The whole point of step was to be able to combine different step methods to be used on different random variables, so that people could end up writing their own kind of Gibbs sampler that was specialized for their needs. Some step methods are used for more than a single variable, like NUTS, so overloading step to make it be yet another kind of sampler that cannot be combined with other steps for our variables wouldn't be good.

Having gone through the step methods themselves not so long ago, I feel that they could be thoroughly refactored, but that would require a whole different discussion and a lot of work. So I think that the simplest option is to use pymc.external.sample_blabla or add a different argument to pymc.sample that is something like external_sampler=Blabla.

The uniform API is very nice, and I would recommend that we rely on pydantic.BaseModel to handle the API validation. I'm not sure what methods we would like to expose. From my backend perspective, I would love it if the external samplers exposed some way to collecting samples on the fly, and also something that stores the external sampler's state, so that it could be reset for continued sampling using the ZarrTrace machinery.

ricardoV94 added the request discussion label Jul 30, 2025

ricardoV94 force-pushed the external_samplers branch from 321e391 to 6a86352 Compare July 30, 2025 14:07

Draft external sampler API

94d5ec4

ricardoV94 force-pushed the external_samplers branch from 6a86352 to 94d5ec4 Compare July 30, 2025 17:45

ricardoV94 mentioned this pull request Jul 30, 2025

Implement nutpie as an external sampler #7719

Closed

nataziel reviewed Jul 31, 2025

View reviewed changes

		@@ -704,8 +666,8 @@ def sample_jax_nuts(
		dims.update(idata_kwargs.pop("dims"))

		# Use 'partial' to set default arguments before passing 'idata_kwargs'



		class JAXSampler(ExternalSampler):
		nuts_sampler = None # Should be defined by subclass

		from pymc.util import RandomState


		class JAXSampler(ExternalSampler):

Uh oh!

Draft external sampler API #7880

Are you sure you want to change the base?

Draft external sampler API #7880

Uh oh!

Conversation

ricardoV94 commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Discussion needed !!

Uh oh!

jessegrabowski commented Jul 30, 2025

Uh oh!

jessegrabowski commented Jul 30, 2025

Uh oh!

ricardoV94 commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jul 30, 2025

Codecov Report

Uh oh!

Dekermanjian commented Jul 30, 2025

Uh oh!

nataziel Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

nataziel Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

nataziel Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

nataziel Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

nataziel Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

nataziel Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

nataziel commented Jul 31, 2025

Uh oh!

jessegrabowski commented Jul 31, 2025

Uh oh!

ricardoV94 commented Jul 31, 2025

Uh oh!

lucianopaz commented Aug 1, 2025

Uh oh!

Uh oh!

ricardoV94 commented Jul 30, 2025 •

edited

Loading

ricardoV94 commented Jul 30, 2025 •

edited

Loading