Check all ancestors for MinibatchOp

Author: zaxtax

pymc/model/core.py

@@ -1241,15 +1241,17 @@ def register_rv(
             self.add_named_variable(rv_var, dims)
             self.set_initval(rv_var, initval)
         else:
-            if (
-                isinstance(observed, TensorVariable)
-                and observed.owner is not None
-                and isinstance(observed.owner.op, MinibatchOp)
-                and total_size is None
-            ):
-                warnings.warn(
-                    f"total_size not provided for observed variable `{name}` that uses pm.Minibatch"
-                )
+            if total_size is None:
+                for node in ancestors([observed]):
+                    if (
+                        isinstance(node, TensorVariable)
+                        and node.owner is not None
+                        and isinstance(node.owner.op, MinibatchOp)
+                    ):
+                        warnings.warn(
+                            f"total_size not provided for observed variable `{name}` that uses pm.Minibatch"
+                        )
+                        break
             if not is_valid_observed(observed):
                 raise TypeError(
                     "Variables that depend on other nodes cannot be used for observed data."