Wrap Ops during etuplization

Evaluating etuplized objects fails for some `RandomVariable` ops whose
`__call__` function does not defer to `make_node`.

In this commit we wrap `RandomVariable` ops during etuplization with a
class that always defers `__call__` to `make_node`. We also add a
dispatch rule for `etuplize` so it also wraps `RandomVariable` with the
same class.

Author: rlouf

aesara/graph/rewriting/unify.py

@@ -11,6 +11,7 @@
 """
 
 from collections.abc import Mapping
+from dataclasses import dataclass
 from numbers import Number
 from typing import Dict, Optional, Tuple, Union
 
@@ -25,6 +26,7 @@
 from aesara.graph.basic import Constant, Variable
 from aesara.graph.op import Op
 from aesara.graph.type import Type
+from aesara.tensor.random.op import RandomVariable
 
 
 def eval_if_etuple(x):
@@ -72,9 +74,50 @@ def __repr__(self):
         return f"{type(self).__name__}({repr(self.constraint)}, {self.token})"
 
 
+@dataclass
+class MakeRandomVariableNodeOp:
+    r"""Wrapper around `RandomVariable` `Op`\s.
+
+    Some `Op.__call__` signatures (e.g. `RandomVariable`) do not match their
+    `Op.make_node` signatures, causing `ExpressionTuple.eval_if_etuple` to
+    fail. To circumvent this limitation, we wrap `Op`\s with this class, which
+    maps `Op.__call__` directly to `Op.make_node`.
+
+    """
+
+    op: Op
+
+    def __call__(self, *args):
+        return self.op.make_node(*args)
+
+
+def car_MakeRandomVariableNodeOp(x):
+    return type(x)
+
+
+_car.add((MakeRandomVariableNodeOp,), car_MakeRandomVariableNodeOp)
+
+
+def cdr_MakeRandomVariableNodeOp(x):
+    x_e = etuple(_car(x), x.op, evaled_obj=x)
+    return x_e[1:]
+
+
+_cdr.add((MakeRandomVariableNodeOp,), cdr_MakeRandomVariableNodeOp)
+
+
+@etuplize.register(RandomVariable)
+def etuplize_random(*args, **kwargs):
+    """Wrap RandomVariable Ops with a MakeNodeOp object."""
+    return etuple(MakeRandomVariableNodeOp, etuplize.funcs[(object,)](*args, **kwargs))
+
+
 def car_Variable(x):
     if x.owner:
-        return x.owner.op
+        if issubclass(type(x.owner.op), RandomVariable):
+            return MakeRandomVariableNodeOp(x.owner.op)
+        else:
+            return x.owner.op
     else:
         raise ConsError("Not a cons pair.")
 

tests/graph/rewriting/test_unify.py

@@ -144,6 +144,17 @@ def perform(self, node, inputs, outputs):
     assert res[0].owner.op == op1_np
     assert res[1].owner.op == op1_np
 
+    # Etuplize RandomVariable ops and the Variables they create
+    x_rv = at.random.normal(0, 1, name="x")
+    x_et = etuplize(x_rv)
+    res = x_et.evaled_obj
+    assert res.owner.op == at.random.normal
+
+    y_et = etuple(etuplize(at.random.normal), *x_et[1:])
+    res = y_et.evaled_obj
+    assert x_et == y_et
+    assert res.op == at.random.normal
+
 
 def test_unify_Variable():
     x_at = at.vector("x")