Fix __bool__ usages with PyTensor variables

Author: ricardoV94

pymc/distributions/multivariate.py

@@ -2057,14 +2057,12 @@ class KroneckerNormal(Continuous):
     rv_op = KroneckerNormalRV.rv_op
 
     @classmethod
-    def dist(cls, mu, covs=None, chols=None, evds=None, sigma=None, *args, **kwargs):
+    def dist(cls, mu, covs=None, chols=None, evds=None, sigma=0.0, *args, **kwargs):
         if len([i for i in [covs, chols, evds] if i is not None]) != 1:
             raise ValueError(
                 "Incompatible parameterization. Specify exactly one of covs, chols, or evds."
             )
 
-        sigma = sigma if sigma else 0
-
         if chols is not None:
             covs = [chol.dot(chol.T) for chol in chols]
         elif evds is not None:
@@ -2076,6 +2074,7 @@ def dist(cls, mu, covs=None, chols=None, evds=None, sigma=None, *args, **kwargs)
                 covs.append(cov_i)
 
         mu = pt.as_tensor_variable(mu)
+        sigma = pt.as_tensor_variable(sigma)
 
         return super().dist([mu, sigma, *covs], **kwargs)
 

pymc/logprob/transforms.py

@@ -470,7 +470,9 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Apply) -> list[Varia
 
     # Do not apply rewrite to discrete variables except for their addition and negation
     if measurable_input.type.dtype.startswith("int"):
-        if not (find_negated_var(measurable_output) or isinstance(node.op.scalar_op, Add)):
+        if not (
+            find_negated_var(measurable_output) is not None or isinstance(node.op.scalar_op, Add)
+        ):
             return None
         # Do not allow rewrite if output is cast to a float, because we don't have meta-info on the type of the MeasurableVariable
         if not measurable_output.type.dtype.startswith("int"):

pymc/logprob/utils.py

@@ -186,13 +186,12 @@ def expand_fn(var):
             return []
 
     if any(
-        ancestor_var
-        for ancestor_var in walk(inputs, expand=expand_fn, bfs=False)
-        if (
+        (
             ancestor_var.owner
             and isinstance(ancestor_var.owner.op, MeasurableOp)
             and not isinstance(ancestor_var.owner.op, ValuedRV)
         )
+        for ancestor_var in walk(inputs, expand=expand_fn, bfs=False)
     ):
         return True
     return False

tests/distributions/test_custom.py

@@ -604,7 +604,7 @@ def test_symbolic_dist(self):
         def dist(size):
             return Truncated.dist(Beta.dist(1, 1, size=size), lower=0.1, upper=0.9)
 
-        assert CustomDist.dist(dist=dist)
+        CustomDist.dist(dist=dist)
 
     def test_nested_custom_dist(self):
         """Test we can create CustomDist that creates another CustomDist"""

tests/distributions/test_multivariate.py

@@ -421,7 +421,7 @@ def test_matrixnormal(self, n):
 
     @pytest.mark.parametrize("n", [2, 3])
     @pytest.mark.parametrize("m", [3])
-    @pytest.mark.parametrize("sigma", [None, 1])
+    @pytest.mark.parametrize("sigma", [0, 1])
     def test_kroneckernormal(self, n, m, sigma):
         np.random.seed(5)
         N = n * m

tests/logprob/test_basic.py

@@ -382,7 +382,7 @@ def test_warn_random_found_probability_inference(func, scipy_func, test_value):
         with pytest.warns(
             UserWarning, match="RandomVariables {input} were found in the derived graph"
         ):
-            assert func(rv, 0.0)
+            func(rv, 0.0)
 
         res = func(rv, 0.0, warn_rvs=False)
         # This is the problem we are warning about, as now we can no longer identify the original rv in the graph

tests/logprob/test_order.py

@@ -291,4 +291,4 @@ def test_non_measurable_max_grad():
     joint_logp = pt.sum([term.sum() for term in logp_terms])
 
     # Test that calling gradient does not raise a NotImplementedError
-    assert pt.grad(joint_logp, x_vv)
+    pt.grad(joint_logp, x_vv)