Finish refactoring BlockedStep.competence implementations

Author: brandonwillard

pymc3/sampling.py

@@ -201,13 +201,14 @@ def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None
             has_gradient = var.dtype not in discrete_types
             if has_gradient:
                 try:
-                    tg.grad(model.logpt, var.tag.value_var)
-                except (AttributeError, NotImplementedError, tg.NullTypeGradError):
+                    tg.grad(model.logpt, var)
+                except (NotImplementedError, tg.NullTypeGradError):
                     has_gradient = False
             # select the best method
+            rv_var = model.values_to_rvs[var]
             selected = max(
                 methods,
-                key=lambda method, var=var, has_gradient=has_gradient: method._competence(
+                key=lambda method, var=rv_var, has_gradient=has_gradient: method._competence(
                     var, has_gradient
                 ),
             )

pymc3/step_methods/metropolis.py

@@ -14,12 +14,12 @@
 from typing import Any, Dict, List, Tuple
 
 import aesara
-import aesara.tensor as at
 import numpy as np
 import numpy.random as nr
 import scipy.linalg
 
-from aesara.tensor.random.basic import CategoricalRV
+from aesara.graph.fg import MissingInputError
+from aesara.tensor.random.basic import BernoulliRV, CategoricalRV
 
 import pymc3 as pm
 
@@ -362,13 +362,22 @@ def competence(var):
         and Categorical variables with k=1.
         """
         distribution = getattr(var.owner, "op", None)
-        if isinstance(distribution, pm.Bernoulli) or (var.dtype in pm.bool_types):
-            return Competence.IDEAL
+
+        if isinstance(distribution, BernoulliRV):
+            return Competence.COMPATIBLE
 
         if isinstance(distribution, CategoricalRV):
-            k = at.get_scalar_constant_value(distribution.owner.inputs[2])
-            if k == 2:
-                return Competence.IDEAL
+            # TODO: We could compute the initial value of `k`
+            # if we had a model object.
+            # k_graph = var.owner.inputs[3].shape[-1]
+            # (k_graph,), _ = rvs_to_value_vars((k_graph,), apply_transforms=True)
+            # k = model.fn(k_graph)(initial_point)
+            try:
+                k = var.owner.inputs[3].shape[-1].eval()
+                if k == 2:
+                    return Competence.COMPATIBLE
+            except MissingInputError:
+                pass
         return Competence.INCOMPATIBLE
 
 
@@ -449,13 +458,22 @@ def competence(var):
         and Categorical variables with k=2.
         """
         distribution = getattr(var.owner, "op", None)
-        if isinstance(distribution, pm.Bernoulli) or (var.dtype in pm.bool_types):
+
+        if isinstance(distribution, BernoulliRV):
             return Competence.IDEAL
 
         if isinstance(distribution, CategoricalRV):
-            k = at.get_scalar_constant_value(distribution.owner.inputs[2])
-            if k == 2:
-                return Competence.IDEAL
+            # TODO: We could compute the initial value of `k`
+            # if we had a model object.
+            # k_graph = var.owner.inputs[3].shape[-1]
+            # (k_graph,), _ = rvs_to_value_vars((k_graph,), apply_transforms=True)
+            # k = model.fn(k_graph)(initial_point)
+            try:
+                k = var.owner.inputs[3].shape[-1].eval()
+                if k == 2:
+                    return Competence.IDEAL
+            except MissingInputError:
+                pass
         return Competence.INCOMPATIBLE
 
 
@@ -585,13 +603,23 @@ def competence(var):
         Categorical variables.
         """
         distribution = getattr(var.owner, "op", None)
+
         if isinstance(distribution, CategoricalRV):
-            k = at.get_scalar_constant_value(distribution.owner.inputs[2])
-            if k == 2:
-                return Competence.IDEAL
+            # TODO: We could compute the initial value of `k`
+            # if we had a model object.
+            # k_graph = var.owner.inputs[3].shape[-1]
+            # (k_graph,), _ = rvs_to_value_vars((k_graph,), apply_transforms=True)
+            # k = model.fn(k_graph)(initial_point)
+            try:
+                k = var.owner.inputs[3].shape[-1].eval()
+                if k > 2:
+                    return Competence.IDEAL
+            except MissingInputError:
+                pass
+
             return Competence.COMPATIBLE
 
-        if isinstance(distribution, pm.Bernoulli) or (var.dtype in pm.bool_types):
+        if isinstance(distribution, BernoulliRV):
             return Competence.COMPATIBLE
 
         return Competence.INCOMPATIBLE

pymc3/step_methods/slicer.py

@@ -119,7 +119,7 @@ def astep(self, q0, logp):
     @staticmethod
     def competence(var, has_grad):
         if var.dtype in continuous_types:
-            if not has_grad and (var.shape is None or var.shape.ndim == 1):
+            if not has_grad and var.ndim == 0:
                 return Competence.PREFERRED
             return Competence.COMPATIBLE
         return Competence.INCOMPATIBLE