Check output is still integer

pymc/logprob/transforms.py

@@ -40,6 +40,7 @@
 from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import pytensor
 import pytensor.tensor as pt
 
 from pytensor import scan
@@ -670,6 +671,8 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
     if isinstance(node.op, MeasurableVariable):
         return None  # pragma: no cover
 
+    pytensor.dprint(node.outputs)
+
     rv_map_feature: Optional[PreserveRVMappings] = getattr(fgraph, "preserve_rv_mappings", None)
     if rv_map_feature is None:
         return None  # pragma: no cover
@@ -681,13 +684,16 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
         return None
 
     [measurable_input] = measurable_inputs
+    [measurable_output] = node.outputs
 
     # Do not apply rewrite to discrete variables
     if measurable_input.type.dtype.startswith("int"):
         if not (
             check_negation(node.op.scalar_op, node.inputs[0]) or isinstance(node.op.scalar_op, Add)
         ):
             return None
+        if not measurable_output.type.dtype.startswith("int"):
+            return None
 
     # Check that other inputs are not potentially measurable, in which case this rewrite
     # would be invalid

tests/logprob/test_order.py

@@ -278,6 +278,8 @@ def test_min_discrete(mu, n, test_value, axis):
         rtol=1e-06,
     )
 
+    assert 0
+
 
 def test_min_max_bernoulli():
     p = 0.7