More strict/explicit signature in step samplers

Author: ricardoV94

pymc/step_methods/arraystep.py

@@ -174,8 +174,9 @@ class GradientSharedStep(ArrayStepShared):
     def __init__(
         self,
         vars,
+        *,
         model=None,
-        blocked=True,
+        blocked: bool = True,
         dtype=None,
         logp_dlogp_func=None,
         rng: RandomGenerator = None,

pymc/step_methods/hmc/base_hmc.py

@@ -82,11 +82,12 @@ class BaseHMC(GradientSharedStep):
     def __init__(
         self,
         vars=None,
+        *,
         scaling=None,
         step_scale=0.25,
         is_cov=False,
         model=None,
-        blocked=True,
+        blocked: bool = True,
         potential=None,
         dtype=None,
         Emax=1000,

pymc/step_methods/metropolis.py

@@ -151,6 +151,7 @@ class Metropolis(ArrayStepShared):
     def __init__(
         self,
         vars=None,
+        *,
         S=None,
         proposal_dist=None,
         scaling=1.0,
@@ -159,7 +160,7 @@ def __init__(
         model=None,
         mode=None,
         rng=None,
-        **kwargs,
+        blocked: bool = False,
     ):
         """Create an instance of a Metropolis stepper.
 
@@ -251,7 +252,7 @@ def __init__(
 
         shared = pm.make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
-        super().__init__(vars, shared, rng=rng)
+        super().__init__(vars, shared, blocked=blocked, rng=rng)
 
     def reset_tuning(self):
         """Reset the tuned sampler parameters to their initial values."""
@@ -418,7 +419,17 @@ class BinaryMetropolis(ArrayStep):
 
     _state_class = BinaryMetropolisState
 
-    def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None, rng=None):
+    def __init__(
+        self,
+        vars,
+        *,
+        scaling=1.0,
+        tune=True,
+        tune_interval=100,
+        model=None,
+        rng=None,
+        blocked: bool = True,
+    ):
         model = pm.modelcontext(model)
 
         self.scaling = scaling
@@ -432,7 +443,7 @@ def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None,
         if not all(v.dtype in pm.discrete_types for v in vars):
             raise ValueError("All variables must be Bernoulli for BinaryMetropolis")
 
-        super().__init__(vars, [model.compile_logp()], rng=rng)
+        super().__init__(vars, [model.compile_logp()], blocked=blocked, rng=rng)
 
     def astep(self, apoint: RaveledVars, *args) -> tuple[RaveledVars, StatsType]:
         logp = args[0]
@@ -530,7 +541,16 @@ class BinaryGibbsMetropolis(ArrayStep):
 
     _state_class = BinaryGibbsMetropolisState
 
-    def __init__(self, vars, order="random", transit_p=0.8, model=None, rng=None):
+    def __init__(
+        self,
+        vars,
+        *,
+        order="random",
+        transit_p=0.8,
+        model=None,
+        rng=None,
+        blocked: bool = True,
+    ):
         model = pm.modelcontext(model)
 
         # Doesn't actually tune, but it's required to emit a sampler stat
@@ -556,7 +576,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None, rng=None):
         if not all(v.dtype in pm.discrete_types for v in vars):
             raise ValueError("All variables must be binary for BinaryGibbsMetropolis")
 
-        super().__init__(vars, [model.compile_logp()], rng=rng)
+        super().__init__(vars, [model.compile_logp()], blocked=blocked, rng=rng)
 
     def reset_tuning(self):
         # There are no tuning parameters in this step method.
@@ -638,7 +658,14 @@ class CategoricalGibbsMetropolis(ArrayStep):
     _state_class = CategoricalGibbsMetropolisState
 
     def __init__(
-        self, vars, proposal="uniform", order="random", model=None, rng: RandomGenerator = None
+        self,
+        vars,
+        *,
+        proposal="uniform",
+        order="random",
+        model=None,
+        rng: RandomGenerator = None,
+        blocked: bool = True,
     ):
         model = pm.modelcontext(model)
 
@@ -693,7 +720,7 @@ def __init__(
         # that indicates whether a draw was done in a tuning phase.
         self.tune = True
 
-        super().__init__(vars, [model.compile_logp()], rng=rng)
+        super().__init__(vars, [model.compile_logp()], blocked=blocked, rng=rng)
 
     def reset_tuning(self):
         # There are no tuning parameters in this step method.
@@ -858,6 +885,7 @@ class DEMetropolis(PopulationArrayStepShared):
     def __init__(
         self,
         vars=None,
+        *,
         S=None,
         proposal_dist=None,
         lamb=None,
@@ -867,7 +895,7 @@ def __init__(
         model=None,
         mode=None,
         rng=None,
-        **kwargs,
+        blocked: bool = True,
     ):
         model = pm.modelcontext(model)
         initial_values = model.initial_point()
@@ -902,7 +930,7 @@ def __init__(
 
         shared = pm.make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
-        super().__init__(vars, shared, rng=rng)
+        super().__init__(vars, shared, blocked=blocked, rng=rng)
 
     def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
         point_map_info = q0.point_map_info
@@ -1025,6 +1053,7 @@ class DEMetropolisZ(ArrayStepShared):
     def __init__(
         self,
         vars=None,
+        *,
         S=None,
         proposal_dist=None,
         lamb=None,
@@ -1035,7 +1064,7 @@ def __init__(
         model=None,
         mode=None,
         rng=None,
-        **kwargs,
+        blocked: bool = True,
     ):
         model = pm.modelcontext(model)
         initial_values = model.initial_point()
@@ -1082,7 +1111,7 @@ def __init__(
 
         shared = pm.make_shared_replacements(initial_values, vars, model)
         self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared)
-        super().__init__(vars, shared, rng=rng)
+        super().__init__(vars, shared, blocked=blocked, rng=rng)
 
     def reset_tuning(self):
         """Reset the tuned sampler parameters and history to their initial values."""

pymc/step_methods/slicer.py

@@ -76,7 +76,15 @@ class Slice(ArrayStepShared):
     _state_class = SliceState
 
     def __init__(
-        self, vars=None, w=1.0, tune=True, model=None, iter_limit=np.inf, rng=None, **kwargs
+        self,
+        vars=None,
+        *,
+        w=1.0,
+        tune=True,
+        model=None,
+        iter_limit=np.inf,
+        rng=None,
+        blocked: bool = False,  # Could be true since tuning is independent across dims?
     ):
         model = modelcontext(model)
         self.w = np.asarray(w).copy()
@@ -97,7 +105,7 @@ def __init__(
         self.logp = compile_pymc([raveled_inp], logp)
         self.logp.trust_input = True
 
-        super().__init__(vars, shared, rng=rng)
+        super().__init__(vars, shared, blocked=blocked, rng=rng)
 
     def astep(self, apoint: RaveledVars) -> tuple[RaveledVars, StatsType]:
         # The arguments are determined by the list passed via `super().__init__(..., fs, ...)`

tests/models.py

@@ -78,7 +78,7 @@ def arbitrary_det(value):
 
 def simple_init():
     start, model, moments = simple_model()
-    step = Metropolis(model.value_vars, np.diag([1.0]), model=model)
+    step = Metropolis(model.value_vars, S=np.diag([1.0]), model=model)
     return model, start, step, moments