Repository: aesara-devs/aesara

Issue:
`OpFromGraph.make_node`'s expected inputs and actual inputs are inconsistent
`OpFromGraph.make_node` expects inputs corresponding to each of its inner-graph inputs, except for the shared variables; however, `OpFromGraph.make_node` also constructs `Apply` nodes such that `Apply.inputs` _contains_ the shared variable inputs of the inner-graph.  This leads to an inherent inconsistency that prevents rewrites from reconstructing nodes via `Op.make_node(*Apply.inputs)`.

Here's an example in which this issue results in an optimization failure:
```python
import aesara
import aesara.tensor as at


from aesara.compile.builders import OpFromGraph


y = aesara.shared(1.0, name="y")

test_ofg = OpFromGraph([], [1 + y])

test_ofg.inputs
# []
test_ofg.shared_inputs
# [y]


def inner_func():
    return test_ofg()


out, out_updates = aesara.scan(inner_func, n_steps=10)

with aesara.config.change_flags(on_opt_error="raise"):
    out_fn = aesara.function([], out)
# ERROR (aesara.graph.opt): Optimization failure due to: push_out_non_seq_scan
# ERROR (aesara.graph.opt): node: for{cpu,scan_fn}(TensorConstant{10}, TensorConstant{10}, y)
# ERROR (aesara.graph.opt): TRACEBACK:
# ERROR (aesara.graph.opt): Traceback (most recent call last):
#   File ".../Aesara/aesara/graph/opt.py", line 1850, in process_node
#     replacements = lopt.transform(fgraph, node)
#   File ".../Aesara/aesara/graph/opt.py", line 1055, in transform
#     return self.fn(fgraph, node)
#   File ".../Aesara/aesara/scan/opt.py", line 282, in push_out_non_seq_scan
#     nw_outer_node = nd.op.make_node(*outside_ins)
#   File ".../Aesara/aesara/compile/builders.py", line 738, in make_node
#     raise ValueError(
# ValueError: Expected 1 inputs, got 2
```

Comments:

Patch:
diff --git a/aesara/compile/builders.py b/aesara/compile/builders.py
index b714138eb..4328e95e7 100644
--- a/aesara/compile/builders.py
+++ b/aesara/compile/builders.py
@@ -375,6 +375,11 @@ class OpFromGraph(Op, HasInnerGraph):
         self.kwargs = kwargs
         self.input_types = [inp.type for inp in inputs]
         self.output_types = [out.type for out in outputs]
+
+        self.lop_overrides = lop_overrides
+        self.grad_overrides = grad_overrides
+        self.rop_overrides = rop_overrides
+
         if lop_overrides != "default":
             if grad_overrides != "default":
                 raise ValueError(
@@ -732,19 +737,71 @@ class OpFromGraph(Op, HasInnerGraph):
         ]
         return ret_l
 
+    def __call__(self, *inputs, **kwargs):
+        # The user interface doesn't expect the shared variable inputs of the
+        # inner-graph, but, since `Op.make_node` does (and `Op.__call__`
+        # dispatches to `Op.make_node`), we need to compensate here
+        num_expected_inps = len(self.inner_inputs) - len(self.shared_inputs)
+
+        if len(inputs) == num_expected_inps:
+            actual_inputs = inputs + tuple(self.shared_inputs)
+            return super().__call__(*actual_inputs, **kwargs)
+        elif len(inputs) == len(self.inner_inputs):
+            return super().__call__(*inputs, **kwargs)
+        else:
+            raise ValueError(f"Expected at least {num_expected_inps} input(s)")
+
     def make_node(self, *inputs):
+        # The `inputs` received here should correspond to the inputs in the
+        # `Apply` nodes we produce below
+        if len(inputs) != len(self.inner_inputs):
+            raise ValueError(f"Expected {len(self.inner_inputs)} input(s)")
+
         num_expected_inps = len(self.inner_inputs) - len(self.shared_inputs)
-        if len(inputs) != num_expected_inps:
-            raise ValueError(
-                f"Expected {int(num_expected_inps)} inputs, got {len(inputs)}"
-            )
-        inputs = [
-            inp_t.filter_variable(inp) for inp, inp_t in zip(inputs, self.input_types)
+        non_shared_inputs = inputs[:num_expected_inps]
+
+        non_shared_inputs = [
+            inp_t.filter_variable(inp)
+            for inp, inp_t in zip(non_shared_inputs, self.input_types)
         ]
+
+        shared_inputs = inputs[num_expected_inps:]
+        local_shared_inputs = self.inner_inputs[num_expected_inps:]
+
+        inner_and_input_shareds = list(zip(local_shared_inputs, shared_inputs))
+
+        if not all(inp_s == inn_s for inn_s, inp_s in inner_and_input_shareds):
+            # The shared variables are not equal to the original shared
+            # variables, so we construct a new `Op` that uses the new shared
+            # variables instead
+            replace = {
+                old_inp: new_inp for old_inp, new_inp in zip(self.inner_inputs, inputs)
+            }
+            replace.update(inner_and_input_shareds)
+
+            # If the new shared variables are inconsistent with the inner-graph,
+            # such errors should arise in this step
+            new_outputs = clone_replace(
+                self.inner_outputs, replace=replace, share_inputs=True
+            )
+
+            new_op = type(self)(
+                inputs=non_shared_inputs,
+                outputs=new_outputs,
+                inline=self.is_inline,
+                lop_overrides=self.lop_overrides,
+                grad_overrides=self.grad_overrides,
+                rop_overrides=self.rop_overrides,
+                connection_pattern=self._connection_pattern,
+                name=self.name,
+            )
+        else:
+            new_op = self
+
         apply_node = Apply(
-            self,
-            list(inputs) + self.shared_inputs,
-            [type() for type in self.output_types],
+            new_op,
+            list(non_shared_inputs) + new_op.shared_inputs,
+            [type() for type in new_op.output_types],
         )
         return apply_node
 


Test patch:
diff --git a/tests/compile/test_builders.py b/tests/compile/test_builders.py
index a552c7a09..9d894aae7 100644
--- a/tests/compile/test_builders.py
+++ b/tests/compile/test_builders.py
@@ -3,6 +3,7 @@ from functools import partial
 import numpy as np
 import pytest
 
+import aesara.tensor as at
 from aesara.compile import shared
 from aesara.compile.builders import OpFromGraph
 from aesara.compile.function import function
@@ -29,6 +30,12 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
     def test_valid_input(self):
         x, y, z = matrices("xyz")
 
+        with pytest.raises(ValueError, match="Expected at least.*"):
+            OpFromGraph([x], [x])()
+
+        with pytest.raises(ValueError, match=r"Expected 1 input\(s\)"):
+            OpFromGraph([x], [x]).make_node()
+
         with pytest.raises(TypeError):
             OpFromGraph((x,), (x,))
 
@@ -451,6 +458,39 @@ class TestOpFromGraph(unittest_tools.InferShapeTester):
         grad_f = grad(f, y)
         assert grad_f.tag.test_value is not None
 
+    def test_make_node_shared(self):
+        """Make sure we can provide `OpFromGraph.make_node` new shared inputs and get a valid `OpFromGraph`."""
+
+        x = at.scalar("x")
+        y = shared(1.0, name="y")
+
+        test_ofg = OpFromGraph([x], [x + y])
+        assert test_ofg.inputs == [x]
+        assert test_ofg.shared_inputs == [y]
+
+        out = test_ofg(x)
+
+        y_clone = y.clone()
+        assert y_clone != y
+        y_clone.name = "y_clone"
+
+        out_new = test_ofg.make_node(*(out.owner.inputs[:1] + [y_clone])).outputs[0]
+
+        assert out_new.owner.op.inputs == [x]
+        assert out_new.owner.op.shared_inputs == [y_clone]
+
+        out_fn = function([x], out_new)
+
+        assert np.array_equal(out_fn(1.0), 2.0)
+
+        y_clone.set_value(2.0)
+
+        assert np.array_equal(out_fn(1.0), 3.0)
+
+        # This should also work, because the containers are the same:
+        # y.set_value(1.0)
+        # assert np.array_equal(out_fn(1.0), 2.0)
+
 
 def test_debugprint():
     x, y, z = matrices("xyz")
diff --git a/tests/scan/test_opt.py b/tests/scan/test_opt.py
index 06eb27372..af14b7930 100644
--- a/tests/scan/test_opt.py
+++ b/tests/scan/test_opt.py
@@ -4,6 +4,7 @@ import pytest
 import aesara
 import aesara.tensor.basic as at
 from aesara import function, scan, shared
+from aesara.compile.builders import OpFromGraph
 from aesara.compile.io import In
 from aesara.compile.mode import get_default_mode
 from aesara.configdefaults import config
@@ -550,6 +551,28 @@ class TestPushOutNonSeqScan:
         utt.assert_allclose(output_opt[0], output_no_opt[0])
         utt.assert_allclose(output_opt[1], output_no_opt[1])
 
+    def test_OpFromGraph_shared(self):
+        """Make sure that a simple `OpFromGraph` with a shared variable can be pushed out."""
+
+        y = shared(1.0, name="y")
+
+        test_ofg = OpFromGraph([], [1 + y])
+
+        def inner_func():
+            return test_ofg()
+
+        out, out_updates = aesara.scan(inner_func, n_steps=10)
+
+        out_fn = function([], out, updates=out_updates)
+
+        res = out_fn()
+        assert np.array_equal(res, np.repeat(2.0, 10))
+
+        y.set_value(2.0)
+
+        res = out_fn()
+        assert np.array_equal(res, np.repeat(3.0, 10))
+
 
 class TestPushOutAddScan:
     """