Add support for update placeholders

Author: brandonwillard

aesara/compile/function/types.py

@@ -6,14 +6,16 @@
 import time
 import warnings
 from itertools import chain
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
 
 import numpy as np
+from typing_extensions import Literal
 
 import aesara
 import aesara.compile.profiling
-from aesara.compile.io import In, SymbolicInput, SymbolicOutput
-from aesara.compile.ops import deep_copy_op, view_op
+from aesara.compile.io import In, Out, SymbolicInput, SymbolicOutput
+from aesara.compile.ops import deep_copy_op, update_placeholder, view_op
+from aesara.compile.profiling import ProfileStats
 from aesara.configdefaults import config
 from aesara.graph.basic import (
     Constant,
@@ -732,10 +734,10 @@ def checkSV(sv_ori, sv_rpl):
                 message = name
             else:
                 message = str(profile.message) + " copy"
-            profile = aesara.compile.profiling.ProfileStats(message=message)
+            profile = ProfileStats(message=message)
             # profile -> object
         elif isinstance(profile, str):
-            profile = aesara.compile.profiling.ProfileStats(message=profile)
+            profile = ProfileStats(message=profile)
 
         f_cpy = maker.__class__(
             inputs=ins,
@@ -1392,13 +1394,35 @@ def check_unused_inputs(inputs, outputs, on_unused_input):
 
     @staticmethod
     def prepare_fgraph(
-        inputs,
-        outputs,
-        additional_outputs,
+        inputs: List[In],
+        outputs: List[Out],
+        additional_outputs: List[Out],
         fgraph: FunctionGraph,
         mode: "Mode",
-        profile,
+        profile: Union[Optional[ProfileStats], Literal[False]],
     ):
+        r"""Perform rewrites on a graph, insert `DeepCopyOp`\s, and remove unused updates.
+
+        .. warning::
+
+            The `additional_outputs` list and `fgraph.outputs` are updated in-place by this method.
+
+        Parameters
+        ==========
+        inputs
+            The wrapped inputs.
+        outputs
+            The wrapped outputs (i.e. wrapped with `Out`).
+        additional_outputs
+            Output graphs that essentially serve as updates to mutable `inputs`.
+        fgraph
+            The `FunctionGraph` to be prepared.
+        mode
+            The `Mode` that determines--for example--which rewrites are applied.
+        profile
+            The profile object/setting to use.
+
+        """
 
         rewriter = mode.optimizer
 
@@ -1419,6 +1443,40 @@ def prepare_fgraph(
                 rewrite_time = end_rewriter - start_rewriter
                 _logger.debug(f"Rewriting took {rewrite_time:f} seconds")
 
+                fgraph_outputs = tuple(fgraph.outputs)
+                update_mappings = tuple(fgraph.update_mapping.items())
+                outputs_to_remove = []
+                additional_outputs_to_remove = []
+
+                # Remove unused updates
+                for i, (out_idx, in_idx) in enumerate(update_mappings):
+                    update = fgraph_outputs[out_idx]
+
+                    if update.owner and update.owner.op == update_placeholder:
+
+                        # TODO: Consider removing the corresponding
+                        # `FunctionGraph` input when it has no other
+                        # references?
+                        # updated_var = fgraph_inputs[in_idx]
+                        # if not fgraph.clients[updated_var]:
+                        #     fgraph.remove_input(updated_var)
+
+                        # Remove the update entry from the wrapped inputs
+                        inputs[in_idx].update = None
+
+                        # We assume that the orders of `fgraph.update_mapping` and
+                        # `additional_outputs` correspond (and they should)
+                        additional_outputs_to_remove.append(additional_outputs[i])
+
+                        outputs_to_remove.append(fgraph.outputs[out_idx])
+
+                for add_out, out in zip(
+                    additional_outputs_to_remove, outputs_to_remove
+                ):
+                    additional_outputs.remove(add_out)
+                    fgraph_out_idx = fgraph.outputs.index(out)
+                    fgraph.remove_output(fgraph_out_idx)
+
                 # Add deep copy to respect the memory interface
                 insert_deepcopy(fgraph, inputs, outputs + additional_outputs)
         finally:

aesara/compile/ops.py

@@ -329,3 +329,45 @@ def make_op(fn):
         return FromFunctionOp(fn, itypes, otypes, infer_shape)
 
     return make_op
+
+
+class UpdatePlaceholder(Op):
+    """A placeholder for a `SharedVariable` update that hasn't been set.
+
+    These will appear in `FunctionGraph.outputs` and represent potential updates
+    (i.e. the `updates` argument to `aesara.function` and/or updates specified via
+    `SharedVariable.default_update`s) that could be specified by rewrites.
+
+    When present, these should be removed by non-rewrite steps in the
+    compilation pipeline (and before any thunks are created for them).
+
+    .. note::
+
+        One reason these can't be removed during the rewrite passes is that the
+        `FunctionGraph.outputs` list entries containing them need to be
+        entirely removed, and we don't want to add/remove
+        `FunctionGraph.outputs` during rewriting.
+
+    """
+
+    view_map = {0: [0]}
+
+    def make_node(self, x):
+        return Apply(self, [x], [x.type()])
+
+    def perform(self, node, inp, out):  # pragma: no cover
+        (x,) = inp
+        (z,) = out
+        z[0] = x
+
+    def __str__(self):  # pragma: no cover
+        return f"{self.__class__.__name__}"
+
+    def infer_shape(self, fgraph, node, input_shapes):
+        return input_shapes
+
+    def grad(self, args, g_outs):  # pragma: no cover
+        return g_outs
+
+
+update_placeholder = UpdatePlaceholder()

tests/compile/function/test_types.py

@@ -11,6 +11,7 @@
 from aesara.compile.function.types import UnusedInputError
 from aesara.compile.io import In, Out
 from aesara.compile.mode import Mode, get_default_mode
+from aesara.compile.ops import update_placeholder
 from aesara.configdefaults import config
 from aesara.graph.basic import Constant
 from aesara.graph.rewriting.basic import OpKeyGraphRewriter, PatternNodeRewriter
@@ -1265,3 +1266,22 @@ def test_empty_givens_updates():
     y = x * 2
     function([In(x)], y, givens={})
     function([In(x)], y, updates={})
+
+
+def test_update_placeholder():
+    a, x, s, m, n = scalars("axsmn")
+
+    f1 = function(
+        [
+            x,
+            In(a, value=1.0, name="a"),
+            In(m, value=0.0, update=update_placeholder(m), mutable=True),
+            In(s, value=0.0, update=s + a * x, mutable=True),
+            In(n, value=0.0, update=update_placeholder(n), mutable=True),
+        ],
+        s + a * x,
+    )
+
+    # The second update shouldn't be present
+    assert len(f1.maker.fgraph.outputs) == 2
+    assert f1.maker.fgraph.update_mapping == {1: 3}