Create arrow to observation nodes subject to arbitrary dtype casting in pm.model_to_graphviz (#6011)

* Adding graphviz area for observed nodes with dtype casting

* Add unit test for obs dtype casting in ModelGraph

* check data container for owner attribute

* Adding comments for make_compute_graph pattern matching

* Repositioned 'While True' loop to better account for arbitary number of dtype casts

* Fix grammar in comment on dtype casting

Author: larryshamalama

pymc/model_graph.py

@@ -20,6 +20,8 @@
 from aesara.compile.sharedvalue import SharedVariable
 from aesara.graph import Apply
 from aesara.graph.basic import ancestors, walk
+from aesara.scalar.basic import Cast
+from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.var import TensorConstant, TensorVariable
 
@@ -98,13 +100,28 @@ def make_compute_graph(
             input_map[var_name] = input_map[var_name].union(parent_name)
 
             if hasattr(var.tag, "observations"):
-                try:
-                    obs_name = var.tag.observations.name
+                obs_node = var.tag.observations
+
+                # loop created so that the elif block can go through this again
+                # and remove any intermediate ops, notably dtype casting, to observations
+                while True:
+
+                    obs_name = obs_node.name
                     if obs_name and obs_name != var_name:
                         input_map[var_name] = input_map[var_name].difference({obs_name})
                         input_map[obs_name] = input_map[obs_name].union({var_name})
-                except AttributeError:
-                    pass
+                        break
+                    elif (
+                        # for cases where observations are cast to a certain dtype
+                        # see issue 5795: https://github.com/pymc-devs/pymc/issues/5795
+                        obs_node.owner
+                        and isinstance(obs_node.owner.op, Elemwise)
+                        and isinstance(obs_node.owner.op.scalar_op, Cast)
+                    ):
+                        # we can retrieve the observation node by going up the graph
+                        obs_node = obs_node.owner.inputs[0]
+                    else:
+                        break
 
         return input_map
 

pymc/tests/test_model_graph.py

@@ -16,6 +16,7 @@
 import pytest
 
 from aesara.compile.sharedvalue import SharedVariable
+from aesara.tensor.var import TensorConstant
 
 import pymc as pm
 
@@ -154,6 +155,25 @@ def model_unnamed_observed_node():
     return model, compute_graph, plates
 
 
+def model_observation_dtype_casting():
+    """
+    Model at the source of the following issue: https://github.com/pymc-devs/pymc/issues/5795
+    """
+    with pm.Model() as model:
+        data = pm.ConstantData("data", [0, 0, 1, 1], dtype=int)
+        p = pm.Beta("p", 1, 1)
+        bern = pm.Bernoulli("response", p, observed=data)
+
+    compute_graph = {
+        "p": set(),
+        "response": {"p"},
+        "data": {"response"},
+    }
+    plates = {"": {"p"}, "4": {"data", "response"}}
+
+    return model, compute_graph, plates
+
+
 class BaseModelGraphTest(SeededTest):
     model_func = None
 
@@ -166,7 +186,7 @@ def test_inputs(self):
         for child, parents_in_plot in self.compute_graph.items():
             var = self.model[child]
             parents_in_graph = self.model_graph.get_parent_names(var)
-            if isinstance(var, SharedVariable):
+            if isinstance(var, (SharedVariable, TensorConstant)):
                 # observed data also doesn't have parents in the compute graph!
                 # But for the visualization we like them to become decendants of the
                 # RVs that these observations belong to.
@@ -236,6 +256,10 @@ class TestUnnamedObservedNodes(BaseModelGraphTest):
     model_func = model_unnamed_observed_node
 
 
+class TestObservationDtypeCasting(BaseModelGraphTest):
+    model_func = model_observation_dtype_casting
+
+
 class TestVariableSelection:
     @pytest.mark.parametrize(
         "var_names, vars_to_plot, compute_graph",