Add opt_sample

Author: ricardoV94

docs/api_reference.rst

@@ -15,6 +15,7 @@ methods in the current release of PyMC experimental.
    MarginalModel
    marginalize
    model_builder.ModelBuilder
+   opt_sample
 
 Inference
 =========

pymc_experimental/__init__.py

@@ -18,6 +18,7 @@
 from pymc_experimental.inference.fit import fit
 from pymc_experimental.model.marginal.marginal_model import MarginalModel, marginalize
 from pymc_experimental.model.model_api import as_model
+from pymc_experimental.sampling.mcmc import opt_sample
 from pymc_experimental.version import __version__
 
 _log = logging.getLogger("pmx")

pymc_experimental/sampling/__init__.py

@@ -0,0 +1,65 @@
+import sys
+
+from pymc.model.core import Model, modelcontext
+from pymc.model.fgraph import fgraph_from_model, model_from_fgraph
+from pymc.sampling.mcmc import sample
+from pytensor.graph.rewriting.basic import GraphRewriter
+from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabaseQuery
+
+posterior_optimization_db = EquilibriumDB()
+posterior_optimization_db.failure_callback = None  # Raise an error if an optimization fails
+posterior_optimization_db.name = "posterior_optimization_db"
+
+
+def opt_sample(
+    *args,
+    model: Model | None = None,
+    rewriter: GraphRewriter | None = None,
+    verbose: bool = False,
+    **kwargs,
+):
+    """Sample from a model after applying optimizations.
+
+    Parameters
+    ----------
+    model : Model (optional)
+        The model to sample from. If None, use the model associated with the context.
+    rewriter : RewriteDatabaseQuery (optional)
+        The rewriter to use. If None, use the default rewriter.
+    verbose : bool, default=False
+        Print information about the optimizations applied.
+    *args, **kwargs:
+        Passed to `pm.sample`
+
+    Returns
+    -------
+    sample_output:
+        The output of `pm.sample`
+
+    Examples
+    --------
+    .. code:: python
+        import pymc as pm
+        import pymc_experimental as pmx
+
+        with pm.Model() as m:
+            p = pm.Beta("p", 1, 1)
+            y = pm.Binomial("y", n=10, p=p, observed=5)
+
+            idata = pmx.opt_sample(verbose=True)
+    """
+
+    model = modelcontext(model)
+    fgraph, _ = fgraph_from_model(model)
+
+    if rewriter is None:
+        rewriter = posterior_optimization_db.query(RewriteDatabaseQuery(include=[]))
+    _, _, rewrite_counters, *_ = rewriter.rewrite(fgraph)
+
+    if verbose:
+        for rewrite_counter in rewrite_counters:
+            for rewrite, counts in rewrite_counter.items():
+                print(f"Applied optimization: {rewrite} {counts}x", file=sys.stdout)
+
+    new_model = model_from_fgraph(fgraph)
+    return sample(*args, model=new_model, **kwargs)