added cannot_sample_rv

Author: Rishab87

pymc/distributions/multivariate.py

@@ -619,12 +619,6 @@ def dist(cls, n, p, *args, **kwargs):
         return super().dist([n, p], *args, **kwargs)
 
     def support_point(rv, size, n, p):
-        observed = getattr(rv.tag, "observed", None)
-        if observed is None:
-            raise ValueError(
-                "Latent Multinomial variables are not supported for sampling. "
-                "Use a Categorical variable instead."
-            )
         n = pt.shape_padright(n)
         mean = n * p
         mode = pt.round(mean)

pymc/sampling/mcmc.py

@@ -51,6 +51,7 @@
 from pymc.backends.base import IBaseTrace, MultiTrace, _choose_chains
 from pymc.backends.zarr import ZarrChain, ZarrTrace
 from pymc.blocking import DictToArrayBijection
+from pymc.distributions.multivariate import Multinomial
 from pymc.exceptions import SamplingError
 from pymc.initial_point import PointType, StartDict, make_initial_point_fns_per_chain
 from pymc.model import Model, modelcontext
@@ -63,6 +64,7 @@
 )
 from pymc.step_methods import NUTS, CompoundStep
 from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
+from pymc.step_methods.cannot_sample import CannotSampleRV
 from pymc.step_methods.hmc import quadpotential
 from pymc.util import (
     ProgressBarManager,
@@ -144,6 +146,13 @@ def instantiate_steppers(
         if initial_point is None:
             initial_point = model.initial_point()
 
+        for rv in model.free_RVs:
+            if isinstance(rv.owner.op, Multinomial) and getattr(rv.tag, "observed", None) is None:
+                for step_class in list(selected_steps.keys()):
+                    if rv in selected_steps[step_class]:
+                        selected_steps[step_class].remove(rv)
+                selected_steps.setdefault(CannotSampleRV, []).append(rv)
+
         for step_class, vars in selected_steps.items():
             if vars:
                 name = getattr(step_class, "name")

pymc/step_methods/cannot_sample.py

@@ -0,0 +1,21 @@
+from pymc.step_methods.arraystep import ArrayStep
+
+class CannotSampleRV(ArrayStep):
+    """
+    A step method that raises an error when sampling a latent Multinomial variable.
+    """
+    name = "cannot_sample_rv"
+    def __init__(self, vars, **kwargs):
+        # Remove keys that ArrayStep.__init__ does not accept.
+        kwargs.pop("model", None)
+        kwargs.pop("initial_point", None)
+        kwargs.pop("compile_kwargs", None)
+        self.vars = vars
+        super().__init__(vars=vars,fs=[], **kwargs)
+
+    def astep(self, q0):
+        # This method is required by the abstract base class.
+        raise ValueError(
+            "Latent Multinomial variables are not supported"
+        )
+