Replace str "output" by a dummy Op in the clients of the FunctionGraph

pytensor/compile/debugmode.py

@@ -30,6 +30,7 @@ from pytensor.configdefaults import config
 from pytensor.graph.basic import Variable, io_toposort
 from pytensor.graph.destroyhandler import DestroyHandler
 from pytensor.graph.features import AlreadyThere, BadOptimization
+from pytensor.graph.fg import Output
 from pytensor.graph.op import HasInnerGraph, Op
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined
 from pytensor.link.basic import Container, LocalLinker
@@ -628,7 +629,9 @@ def _is_used_in_graph(fgraph, var):
         True if `var` is used by another node in the graph.
 
     """
-    return not (fgraph.clients[var] == [("output", 1)] or fgraph.clients[var] == [])
+    return any(
+        client for client, _ in fgraph.clients[var] if not isinstance(client.op, Output)
+    )
 
 
 def _check_strides_match(a, b, warn_err, op):
@@ -977,7 +980,7 @@ def _check_preallocated_output(
     # disable memory checks in that mode, since they were already run.
     try:
         changed_inner_mode = False
-        if isinstance(getattr(node, "op", None), HasInnerGraph):
+        if isinstance(node.op, HasInnerGraph):
             fn = node.op.fn
             if not (fn and hasattr(fn, "maker") and hasattr(fn.maker, "mode")):
                 _logger.warning(f"Expected pytensor function not found in {node.op}.fn")
@@ -1132,18 +1135,14 @@ class _FunctionGraphEvent:
 
     def __init__(self, kind, node, idx=None, reason=None):
         self.kind = kind
-        if node == "output":
-            self.node = "output"
-            self.op = "output"
-        else:
-            self.node = node
-            self.op = node.op
+        self.node = node
+        self.op = node.op
         self.idx = idx
         self.reason = str(reason)
 
     def __str__(self):
         if self.kind == "change":
-            if self.op != "output":
+            if not isinstance(self.op, Output):
                 msg = str(len(self.node.inputs))
             else:
                 msg = ""

pytensor/compile/function/types.py

@@ -78,8 +78,6 @@ def view_tree_set(fgraph, v, treeset):
     """
     treeset.add(v)
     for cl, v_input_pos_to_cl in fgraph.clients[v]:
-        if cl == "output":
-            continue
         vmap = cl.op.view_map
         dmap = cl.op.destroy_map
         for opos, iposlist in chain(vmap.items(), dmap.items()):
@@ -1202,8 +1200,11 @@ def insert_deepcopy(fgraph, wrapped_inputs, wrapped_outputs):
     has_destroyers_attr = hasattr(fgraph, "has_destroyers")
 
     for i in range(len(fgraph.outputs)):
+        original_out = fgraph.outputs[i]
+        output_client = fgraph.get_output_client(i)
+
         views_of_output_i = set()
-        view_tree_set(fgraph, alias_root(fgraph.outputs[i]), views_of_output_i)
+        view_tree_set(fgraph, alias_root(original_out), views_of_output_i)
         copied = False
         # do not allow outputs to be aliased
         for j in range(i + 1, len(fgraph.outputs)):
@@ -1212,16 +1213,16 @@ def insert_deepcopy(fgraph, wrapped_inputs, wrapped_outputs):
             if fgraph.outputs[j] in views_of_output_i:
                 if wrapped_outputs[i].borrow and wrapped_outputs[j].borrow:
                     fgraph.change_node_input(
-                        "output", i, view_op(fgraph.outputs[i]), reason=reason
+                        *output_client, view_op(original_out), reason=reason
                     )
                 else:
                     fgraph.change_node_input(
-                        "output", i, deep_copy_op(fgraph.outputs[i]), reason=reason
+                        *output_client, deep_copy_op(original_out), reason=reason
                     )
                 copied = True
                 break
 
-        if not copied:
+        if not copied:  # no-break
             for input_j in all_graph_inputs:
                 # do not allow outputs to be aliased to an inputs (j), unless
                 # a) that j'th input has been 'destroyed' by
@@ -1239,33 +1240,29 @@ def insert_deepcopy(fgraph, wrapped_inputs, wrapped_outputs):
                         j = fgraph.inputs.index(input_j)
                         if wrapped_outputs[i].borrow and wrapped_inputs[j].borrow:
                             fgraph.change_node_input(
-                                "output",
-                                i,
-                                view_op(fgraph.outputs[i]),
+                                *output_client,
+                                view_op(original_out),
                                 reason=reason,
                             )
                             break
                         else:
                             fgraph.change_node_input(
-                                "output",
-                                i,
-                                deep_copy_op(fgraph.outputs[i]),
+                                *output_client,
+                                deep_copy_op(original_out),
                                 reason=reason,
                             )
                             break
                     elif wrapped_outputs[i].borrow:
                         fgraph.change_node_input(
-                            "output",
-                            i,
-                            view_op(fgraph.outputs[i]),
+                            *output_client,
+                            view_op(original_out),
                             reason=reason,
                         )
                         break
                     else:
                         fgraph.change_node_input(
-                            "output",
-                            i,
-                            deep_copy_op(fgraph.outputs[i]),
+                            *output_client,
+                            deep_copy_op(original_out),
                             reason=reason,
                         )
                         break

pytensor/compile/profiling.py

@@ -16,20 +16,17 @@ import sys
 import time
 from collections import Counter, defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import numpy as np
 
 import pytensor
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Apply, Constant, Variable
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.link.utils import get_destroy_dependencies
 
 
-if TYPE_CHECKING:
-    from pytensor.graph.fg import FunctionGraph
-
-
 @contextmanager
 def extended_open(filename, mode="r"):
     if filename == "<stdout>":
@@ -1038,7 +1035,7 @@ class ProfileStats:
             executable_nodes = set()
             for var in fgraph.inputs:
                 for c, _ in fgraph.clients[var]:
-                    if c != "output":
+                    if not isinstance(c.op, Output):
                         deps = c.inputs + destroy_dependencies[c]
                         if all(compute_map[v][0] for v in deps):
                             executable_nodes.add(c)
@@ -1166,7 +1163,7 @@ class ProfileStats:
 
                         for var in node.outputs:
                             for c, _ in fgraph.clients[var]:
-                                if c != "output":
+                                if not isinstance(c.op, Output):
                                     deps = c.inputs + destroy_dependencies[c]
                                     if all(compute_map[v][0] for v in deps):
                                         new_exec_nodes.add(c)

pytensor/graph/destroyhandler.py

@@ -11,6 +11,7 @@ import pytensor
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Constant
 from pytensor.graph.features import AlreadyThere, Bookkeeper
+from pytensor.graph.fg import Output
 from pytensor.graph.utils import InconsistencyError
 from pytensor.misc.ordered_set import OrderedSet
 
@@ -401,8 +402,6 @@ class DestroyHandler(Bookkeeper):
             def recursive_destroys_finder(protected_var):
                 # protected_var is the idx'th input of app.
                 for app, idx in fgraph.clients[protected_var]:
-                    if app == "output":
-                        continue
                     destroy_maps = app.op.destroy_map.values()
                     # If True means that the apply node, destroys the protected_var.
                     if idx in [dmap for sublist in destroy_maps for dmap in sublist]:
@@ -578,7 +577,7 @@ class DestroyHandler(Bookkeeper):
         app.inputs[i] changed from old_r to new_r.
 
         """
-        if app == "output":
+        if isinstance(app.op, Output):
             # app == 'output' is special key that means FunctionGraph is redefining which nodes are being
             # considered 'outputs' of the graph.
             pass

pytensor/graph/fg.py

@@ -3,7 +3,7 @@
 import time
 from collections import defaultdict
 from collections.abc import Iterable, Sequence
-from typing import TYPE_CHECKING, Any, Literal, Union, cast
+from typing import Any, Union, cast
 
 import pytensor
 from pytensor.configdefaults import config
@@ -19,15 +19,30 @@ from pytensor.graph.basic import (
 )
 from pytensor.graph.basic import as_string as graph_as_string
 from pytensor.graph.features import AlreadyThere, Feature, ReplaceValidate
+from pytensor.graph.op import Op
 from pytensor.graph.utils import MetaObject, MissingInputError, TestValueError
 from pytensor.misc.ordered_set import OrderedSet
 
 
-if TYPE_CHECKING:
-    from pytensor.graph.op import Op
+ClientType = tuple[Apply, int]
 
-ApplyOrOutput = Apply | Literal["output"]
-ClientType = tuple[ApplyOrOutput, int]
+
+class Output(Op):
+    """A dummy `Op` that represents an output variable in a `FunctionGraph`."""
+
+    __props__ = ("idx",)
+
+    def __init__(self, idx):
+        self.idx = idx
+
+    def make_node(self, inp):
+        return Apply(self, [inp], [])
+
+    def perform(self, node, inputs, outputs):
+        raise RuntimeError("Output Ops should never be evaluated")
+
+    def __str__(self):
+        return f"output[{self.idx}]"
 
 
 class FunctionGraph(MetaObject):
@@ -157,7 +172,7 @@ class FunctionGraph(MetaObject):
         """Add a new variable as an output to this `FunctionGraph`."""
         self.outputs.append(var)
         self.import_var(var, reason=reason, import_missing=import_missing)
-        self.clients[var].append(("output", len(self.outputs) - 1))
+        self.clients[var].append((Output(len(self.outputs) - 1).make_node(var), 0))
 
     def add_input(self, var: Variable, check: bool = True) -> None:
         """Add a new variable as an input to this `FunctionGraph`.
@@ -198,10 +213,8 @@ class FunctionGraph(MetaObject):
             A ``(node, i)`` pair such that ``node.inputs[i]`` is `var`.
 
         """
-        if not isinstance(new_client[0], Apply) and new_client[0] != "output":
-            raise TypeError(
-                'The first entry of `new_client` must be an `Apply` node or the string `"output"`'
-            )
+        if not isinstance(new_client[0], Apply):
+            raise TypeError("The first entry of `new_client` must be an `Apply` node")
         self.clients[var].append(new_client)
 
     def remove_client(
@@ -278,6 +291,16 @@ class FunctionGraph(MetaObject):
                     if remove_if_empty:
                         del clients[var]
 
+    def get_output_client(self, i: int) -> ClientType:
+        """Get the dummy Output Op client to output i.
+
+        Raises lookup error if not found
+        """
+        for client in self.clients[self.outputs[i]]:
+            if isinstance(client[0].op, Output) and client[0].op.idx == i:
+                return client
+        raise LookupError
+
     def import_var(
         self, var: Variable, reason: str | None = None, import_missing: bool = False
     ) -> None:
@@ -382,7 +405,7 @@ class FunctionGraph(MetaObject):
 
     def change_node_input(
         self,
-        node: ApplyOrOutput,
+        node: Apply,
         i: int,
         new_var: Variable,
         reason: str | None = None,
@@ -401,9 +424,7 @@ class FunctionGraph(MetaObject):
         Parameters
         ----------
         node
-            The node for which an input is to be changed.  If the value is
-            the string ``"output"`` then the ``self.outputs`` will be used
-            instead of ``node.inputs``.
+            The node for which an input is to be changed.
         i
             The index in `node.inputs` that we want to change.
         new_var
@@ -417,27 +438,21 @@ class FunctionGraph(MetaObject):
             narrowed and would otherwise fail this check.
         """
         # TODO: ERROR HANDLING FOR LISTENERS (should it complete the change or revert it?)
-        if node == "output":
-            r = self.outputs[i]
-            if check and not r.type.is_super(new_var.type):
-                raise TypeError(
-                    f"The type of the replacement ({new_var.type}) must be "
-                    f"compatible with the type of the original Variable ({r.type})."
-                )
-            self.outputs[i] = new_var
-        else:
-            assert isinstance(node, Apply)
-            r = node.inputs[i]
-            if check and not r.type.is_super(new_var.type):
-                raise TypeError(
-                    f"The type of the replacement ({new_var.type}) must be "
-                    f"compatible with the type of the original Variable ({r.type})."
-                )
-            node.inputs[i] = new_var
+        r = node.inputs[i]
 
         if r is new_var:
             return
 
+        if check and not r.type.is_super(new_var.type):
+            raise TypeError(
+                f"The type of the replacement ({new_var.type}) must be "
+                f"compatible with the type of the original Variable ({r.type})."
+            )
+        node.inputs[i] = new_var
+
+        if isinstance(node.op, Output):
+            self.outputs[node.op.idx] = new_var
+
         self.import_var(new_var, reason=reason, import_missing=import_missing)
         self.add_client(new_var, (node, i))
         self.remove_client(r, (node, i), reason=reason)
@@ -518,33 +533,6 @@ class FunctionGraph(MetaObject):
         for var, new_var in pairs:
             self.replace(var, new_var, **kwargs)
 
-    def _remove_output(self, idx: int):
-        """Remove the output at index `idx` and update the indices in the clients entries.
-
-        `FunctionGraph.clients` contains entries like ``("output", i)`` under
-        each output variable in `FunctionGraph.outputs`.  The ``i`` values
-        correspond to each output's location within the `FunctionGraph.outputs`
-        list, so, when an output is removed from the graph, all these entries
-        need to be updated.  This method performs those updates.
-
-        TODO: We could track these entries in a new instance attribute and make
-        them lists, then each could be updated in-place very easily.  This
-        seems fine, because the `FunctionGraph.clients` ``dict`` and list in
-        which they're contained are already being updated in-place.
-        """
-        old_idx_mappings = tuple((out, i) for i, out in enumerate(self.outputs))
-        self.outputs.pop(idx)
-
-        new_idx = 0
-        for out, old_idx in old_idx_mappings:
-            if old_idx == idx:
-                continue
-            out_clients = self.clients[out]
-            arrow: ClientType = ("output", old_idx)
-            arrow_idx = out_clients.index(arrow)
-            out_clients[arrow_idx] = ("output", new_idx)
-            new_idx += 1
-
     def remove_node(self, node: Apply, reason: str | None = None):
         """Remove an `Apply` node from the `FunctionGraph`.
 
@@ -571,8 +559,8 @@ class FunctionGraph(MetaObject):
             while out_clients:
                 out_client, out_idx = out_clients.pop()
 
-                if out_client == "output":
-                    self._remove_output(out_idx)
+                if isinstance(out_client.op, Output):
+                    self.remove_output(out_client.op.idx, remove_client=False)
 
                     # TODO: We could short-circuit all of the graph walking and
                     # clear everything at once when all the outputs are gone.
@@ -588,7 +576,6 @@ class FunctionGraph(MetaObject):
                     #
                     #         self.execute_callbacks("on_prune", node, reason)
                 else:
-                    assert isinstance(out_client, Apply)
                     self.remove_node(out_client, reason=reason)
 
             clients.pop(out, None)
@@ -630,32 +617,46 @@ class FunctionGraph(MetaObject):
         self.execute_callbacks("on_prune", node, reason)
 
     def remove_input(self, input_idx: int, reason: str | None = None):
-        """Remove the input at index `input_idx`."""
+        """Remove the input at index `input_idx`.
+
+        Any node that depended on such input will also be removed.
+        """
         var = self.inputs.pop(input_idx)
 
         for client, idx in list(self.clients[var]):
-            if client == "output":
-                out_var = self.outputs[idx]
-                out_node = out_var.owner
-                if out_node is None:
-                    assert out_var in self.inputs
-                    self.outputs.pop(idx)
-                    continue
-                client_node = out_node
-            else:
-                assert isinstance(client, Apply)
-                client_node = client
+            self.remove_node(client, reason=reason)
+
+    def remove_output(
+        self, output_idx: int, reason: str | None = None, remove_client: bool = True
+    ):
+        """Remove the output at index `output_idx` and update the indices in the clients entries.
+
+        `FunctionGraph.clients` contains entries like ``(output(i)(var), 0)`` under
+        each output variable in `FunctionGraph.outputs`.  The ``i`` values
+        correspond to each output's location within the `FunctionGraph.outputs`
+        list, so, when an output is removed from the graph, all these entries
+        need to be updated.  This method performs those updates.
 
-            self.remove_node(client_node, reason=reason)
+        """
+        outputs = self.outputs
 
-    def remove_output(self, output_idx: int, reason: str | None = None):
-        """Remove the output at index `input_idx`."""
+        # We have to update all the output indexes to the right of the removed index
+        for old_idx, out in enumerate(outputs[output_idx + 1 :], output_idx + 1):
+            old_client = self.get_output_client(old_idx)
+            out_clients = self.clients[out]
+            out_clients[out_clients.index(old_client, 0)] = (
+                Output(old_idx - 1).make_node(out),
+                0,
+            )
 
-        var = self.outputs[output_idx]
-        self._remove_output(output_idx)
-        self.remove_client(
-            var, ("output", output_idx), reason=reason, remove_if_empty=True
-        )
+        # Remove the Output Op client from the clients list
+        # This is false when called from `remove_node` which removes the clients ahead of time
+        if remove_client:
+            output_client = self.get_output_client(output_idx)
+            self.remove_client(
+                outputs[output_idx], output_client, reason=reason, remove_if_empty=True
+            )
+        outputs.pop(output_idx)
 
     def attach_feature(self, feature: Feature) -> None:
         """Add a ``graph.features.Feature`` to this function graph and trigger its ``on_attach`` callback."""
@@ -832,19 +833,17 @@ class FunctionGraph(MetaObject):
             ):
                 raise Exception(f"Undeclared input: {variable}")
             for cl_node, i in clients[variable]:
-                if cl_node == "output":
-                    if self.outputs[i] is not variable:
+                if isinstance(cl_node.op, Output):
+                    out_idx = cl_node.op.idx
+                    if self.outputs[out_idx] is not variable:
                         raise Exception(
-                            f"Inconsistent clients list: {variable}, {self.outputs[i]}"
+                            f"Inconsistent clients list: {variable}, {self.outputs[out_idx]}"
                         )
-                    continue
-
-                assert isinstance(cl_node, Apply)
-
-                if cl_node not in nodes:
+                elif cl_node not in nodes:
                     raise Exception(
                         f"Client not in FunctionGraph: {variable}, {(cl_node, i)}"
                     )
+
                 if cl_node.inputs[i] is not variable:
                     raise Exception(
                         f"Inconsistent clients list: {variable}, {cl_node.inputs[i]}"

pytensor/graph/rewriting/basic.py

@@ -30,7 +30,7 @@ from pytensor.graph.basic import (
     vars_between,
 )
 from pytensor.graph.features import AlreadyThere, Feature, NodeFinder
-from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.op import Op
 from pytensor.graph.utils import AssocList, InconsistencyError
 from pytensor.misc.ordered_set import OrderedSet
@@ -738,7 +738,7 @@ class MergeOptimizer(GraphRewriter):
                         if any(
                             i in flatten(c.op.destroy_map.values())
                             for c, i in clients
-                            if c != "output" and c.op.destroy_map
+                            if c.op.destroy_map
                         ):
                             continue
 
@@ -1612,8 +1612,6 @@ class PatternNodeRewriter(NodeRewriter):
 
         if get_nodes and self.get_nodes is not None:
             for real_node in self.get_nodes(fgraph, node):
-                if real_node == "output":
-                    continue
                 ret = self.transform(fgraph, real_node, get_nodes=False)
                 if ret is not False and ret is not None:
                     return dict(zip(real_node.outputs, ret))
@@ -2399,7 +2397,7 @@ class EquilibriumGraphRewriter(NodeProcessingGraphRewriter):
             if self.tracks_on_change_inputs:
 
                 def chin_(node, i, r, new_r, reason):
-                    if node is not current_node and not isinstance(node, str):
+                    if node is not current_node and not isinstance(node.op, Output):
                         q.append(node)
 
                 chin = chin_

pytensor/graph/rewriting/utils.py

 import copy
 from collections.abc import Generator, Sequence
-from typing import TYPE_CHECKING, Optional, cast
+from typing import TYPE_CHECKING, Optional
 
 import pytensor
 from pytensor.graph.basic import (
@@ -10,7 +10,7 @@ from pytensor.graph.basic import (
     graph_inputs,
     vars_between,
 )
-from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 
 
@@ -230,11 +230,9 @@ def get_clients_at_depth(
     for var in node.outputs:
         if depth > 0:
             for out_node, _ in fgraph.clients[var]:
-                if out_node == "output":
+                if isinstance(out_node.op, Output):
                     continue
-                yield from get_clients_at_depth(
-                    fgraph, cast(Apply, out_node), depth - 1
-                )
+                yield from get_clients_at_depth(fgraph, out_node, depth - 1)
         else:
             assert var.owner is not None
             yield var.owner

pytensor/link/c/basic.py

@@ -354,9 +354,7 @@ def get_c_declare(fgraph, r, name, sub):
     # it means they need `r`'s dtype to be declared, so
     # we have to pass `check_input=True` to `c_declare`.
     if any(
-        getattr(c.op, "check_input", config.check_input)
-        for (c, _) in fgraph.clients[r]
-        if not isinstance(c, str)
+        getattr(c.op, "check_input", config.check_input) for (c, _) in fgraph.clients[r]
     ) or (r.owner and getattr(r.owner.op, "check_input", config.check_input)):
         c_declare = r.type.c_declare(name, sub, True)
     else:

pytensor/link/vm.py

@@ -954,8 +954,7 @@ class VMLinker(LocalLinker):
             if k.owner and self.fgraph.clients[k]:
                 ls = []
                 for cl in self.fgraph.clients[k]:
-                    if cl[0] != "output":
-                        ls += cl[0].outputs
+                    ls += cl[0].outputs
                 dependencies[k] += ls
         return dependencies
 

pytensor/printing.py

@@ -437,10 +437,15 @@ N.B.:
 
             for out in inner_outputs:
                 if (
-                    isinstance(getattr(out.owner, "op", None), HasInnerGraph)
-                    or hasattr(getattr(out.owner, "op", None), "scalar_op")
-                    and isinstance(out.owner.op.scalar_op, HasInnerGraph)
-                ) and out not in inner_graph_vars:
+                    out.owner is not None
+                    and (
+                        isinstance(out.owner.op, HasInnerGraph)
+                        or isinstance(
+                            getattr(out.owner.op, "scalar_op", None), HasInnerGraph
+                        )
+                    )
+                    and out not in inner_graph_vars
+                ):
                     inner_graph_vars.append(out)
 
                 _debugprint(

pytensor/scan/rewriting.py

@@ -27,7 +27,7 @@ from pytensor.graph.basic import (
 )
 from pytensor.graph.destroyhandler import DestroyHandler
 from pytensor.graph.features import ReplaceValidate
-from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import (
@@ -1303,7 +1303,7 @@ def scan_save_mem(fgraph, node):
         for cl, _ in fgraph.clients[out]:
             # 2.1 outputs of the function
             # => output needs all its intermediate values
-            if isinstance(cl, str):
+            if isinstance(cl.op, Output):
                 # if the node is actually an output, then
                 # we need to store the entire thing
                 global_nsteps = None
@@ -1412,7 +1412,7 @@ def scan_save_mem(fgraph, node):
     for i, out in enumerate(node.outputs[:c_outs]):
         # look at all its clients
         for cl, _ in fgraph.clients[out]:
-            if isinstance(cl, str):
+            if isinstance(cl.op, Output):
                 store_steps[i] = 0
                 break
             elif not isinstance(cl.op, Subtensor):
@@ -2309,7 +2309,7 @@ def scan_push_out_dot1(fgraph, node):
             and isinstance(out.owner.op.scalar_op, ps.Add)
             and inp in out.owner.inputs
             and len(fgraph.clients[outer_out]) == 1
-            and not isinstance(fgraph.clients[outer_out][0][0], str)
+            and not isinstance(fgraph.clients[outer_out][0][0], Output)
             and isinstance(fgraph.clients[outer_out][0][0].op, Subtensor)
             and fgraph.clients[outer_out][0][0].op.idx_list == (-1,)
         ):

pytensor/tensor/basic.py

@@ -24,7 +24,7 @@ from pytensor import scalar as ps
 from pytensor.gradient import DisconnectedType, grad_undefined
 from pytensor.graph import RewriteDatabaseQuery
 from pytensor.graph.basic import Apply, Constant, Variable, equal_computations
-from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.op import Op
 from pytensor.graph.replace import _vectorize_node
 from pytensor.graph.rewriting.db import EquilibriumDB
@@ -1681,7 +1681,7 @@ class Alloc(COp):
             return False
 
         for client, idx in clients:
-            if client == "output":
+            if isinstance(client.op, Output):
                 # If the output is a constant, it will have to be deepcopied
                 # each time the function is called.  So we do not fold.
                 return False

pytensor/tensor/random/rewriting/basic.py

@@ -31,15 +31,12 @@ def is_rv_used_in_graph(base_rv, node, fgraph):
     TODO: We should apply all the shape rewrites before these rewrites, since
     that would properly remove the unnecessary dependencies on `base_rv` (when
     possible).
-
     """
-
-    def _node_check(n, i):
-        if n == "output":
-            n = fgraph.outputs[i].owner
-        return n == node or isinstance(n.op, Shape | Shape_i)
-
-    return not all(_node_check(n, i) for n, i in fgraph.clients.get(base_rv, ()))
+    return any(
+        n
+        for n, i in fgraph.clients.get(base_rv, ())
+        if not (n is node or isinstance(n.op, Shape | Shape_i))
+    )
 
 
 @node_rewriter([RandomVariable], inplace=True)

pytensor/tensor/rewriting/elemwise.py

@@ -14,7 +14,7 @@ from pytensor.configdefaults import config
 from pytensor.graph import FunctionGraph
 from pytensor.graph.basic import Apply, Constant, Variable, ancestors, io_toposort
 from pytensor.graph.features import ReplaceValidate
-from pytensor.graph.fg import ApplyOrOutput
+from pytensor.graph.fg import Output
 from pytensor.graph.rewriting.basic import (
     EquilibriumGraphRewriter,
     GraphRewriter,
@@ -688,7 +688,7 @@ class FusionOptimizer(GraphRewriter):
             """
 
             FUSEABLE_MAPPING = defaultdict[Variable, list[Apply]]
-            UNFUSEABLE_MAPPING = defaultdict[Variable, set[ApplyOrOutput]]
+            UNFUSEABLE_MAPPING = defaultdict[Variable, set[Apply]]
 
             def initialize_fuseable_mappings(
                 *, fg: FunctionGraph
@@ -727,7 +727,6 @@ class FusionOptimizer(GraphRewriter):
                     for client, _ in clients:
                         if (
                             out_maybe_fuseable
-                            and not isinstance(client, str)  # "output"
                             and isinstance(client.op, Elemwise)
                             # and not isinstance(client.op.scalar_op, ps.Composite)
                             and len(client.outputs) == 1
@@ -841,7 +840,7 @@ class FusionOptimizer(GraphRewriter):
                             implied_unfuseable_clients = {
                                 c
                                 for client in unfuseable_clients_clone.get(next_out, ())
-                                if not isinstance(client, str)  # "output"
+                                if not isinstance(client.op, Output)
                                 for c in client.outputs
                             }
 

pytensor/tensor/rewriting/linalg.py

@@ -299,8 +299,6 @@ def local_det_chol(fgraph, node):
     """
     (x,) = node.inputs
     for cl, xpos in fgraph.clients[x]:
-        if cl == "output":
-            continue
         if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, Cholesky):
             L = cl.outputs[0]
             return [prod(diagonal(L, axis1=-2, axis2=-1) ** 2, axis=-1)]

pytensor/tensor/rewriting/math.py

@@ -1126,14 +1126,11 @@ class AlgebraicCanonizer(NodeRewriter):
         # this canonized graph...  if so, we do nothing and wait for
         # them to be transformed.
         for c, c_idx in out_clients:
-            if c == "output":
-                continue
             while (
-                isinstance(getattr(c, "op", None), DimShuffle)
-                and len(fgraph.clients[c.outputs[0]]) <= 1
+                isinstance(c.op, DimShuffle) and len(fgraph.clients[c.outputs[0]]) <= 1
             ):
                 c = fgraph.clients[c.outputs[0]][0][0]
-            if getattr(c, "op", "") in [self.main, self.inverse, self.reciprocal]:
+            if c.op in [self.main, self.inverse, self.reciprocal]:
                 return False
 
         # Here we make the canonical version of the graph around this node

pytensor/tensor/rewriting/shape.py

@@ -401,7 +401,7 @@ class ShapeFeature(Feature):
                 merged_shape.append(other_shape[i])
             elif (
                 ps.owner
-                and isinstance(getattr(ps.owner, "op", None), Shape_i)
+                and isinstance(ps.owner.op, Shape_i)
                 and ps.owner.op.i == i
                 and ps.owner.inputs[0] in (r, other_r)
             ):
@@ -602,7 +602,7 @@ class ShapeFeature(Feature):
         # r is *scheduled*.
         # At that point, node is no longer a client of r, but of new_r
         for shpnode, idx in fgraph.clients[r] + [(node, i)]:
-            if isinstance(getattr(shpnode, "op", None), Shape_i):
+            if isinstance(shpnode.op, Shape_i):
                 idx = shpnode.op.i
                 repl = self.shape_of[new_r][idx]
                 if repl.owner is shpnode:
@@ -1057,7 +1057,10 @@ def local_Shape_of_SpecifyShape(fgraph, node):
 
     specified_shape = node.inputs[0]
 
-    if not isinstance(getattr(specified_shape.owner, "op", None), SpecifyShape):
+    if not (
+        specified_shape.owner is not None
+        and isinstance(specified_shape.owner.op, SpecifyShape)
+    ):
         return False
 
     x, *shape = specified_shape.owner.inputs

tests/graph/test_fg.py

@@ -6,7 +6,7 @@ import pytest
 
 from pytensor.configdefaults import config
 from pytensor.graph.basic import NominalVariable
-from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.utils import MissingInputError
 from pytensor.printing import debugprint
 from tests.graph.utils import (
@@ -78,8 +78,13 @@ class TestFunctionGraph:
         assert fg.variables == {var1, var2, var3, var4}
         assert fg.get_clients(var1) == [(var3.owner, 0)]
         assert fg.get_clients(var2) == [(var4.owner, 1)]
-        assert fg.get_clients(var3) == [("output", 0), (var4.owner, 0)]
-        assert fg.get_clients(var4) == [("output", 1)]
+        var3_clients = fg.get_clients(var3)
+        assert len(var3_clients) == 2
+        assert var3_clients[0][0].op == Output(0)
+        assert var3_clients[1] == (var4.owner, 0)
+        var4_clients = fg.get_clients(var4)
+        assert len(var4_clients) == 1
+        assert var4_clients[0][0].op == Output(1)
 
         varC = MyConstant("varC")
         var5 = op1(var1, varC)
@@ -208,8 +213,11 @@ class TestFunctionGraph:
         fg = FunctionGraph([var1, var2], [var3, var5], clone=False)
 
         var6 = MyVariable2("var6")
+        [out_client] = [
+            cl for cl, _ in fg.clients[fg.outputs[0]] if isinstance(cl.op, Output)
+        ]
         with pytest.raises(TypeError):
-            fg.change_node_input("output", 1, var6)
+            fg.change_node_input(out_client, 0, var6)
 
         with pytest.raises(TypeError):
             fg.change_node_input(var5.owner, 1, var6)
@@ -358,12 +366,13 @@ class TestFunctionGraph:
 
         # TODO: What if the index value is greater than 1?  It will throw an
         # `IndexError`, but that doesn't sound like anything we'd want.
+        out_node = Output(idx=1).make_node(var4)
         with pytest.raises(Exception, match="Inconsistent clients list.*"):
-            fg.add_client(var4, ("output", 1))
+            fg.add_client(var4, (out_node, 0))
 
             fg.check_integrity()
 
-        fg.remove_client(var4, ("output", 1))
+        fg.remove_client(var4, (out_node, 0))
 
         with pytest.raises(TypeError, match="The first entry of.*"):
             fg.add_client(var4, (None, 0))