class_docstr: Graph optimizer that fuses consecutive Elemwise operations.
class_docstr_tok: ["Graph","optimizer","that","fuses","consecutive","Elemwise","operations","."]
code: def apply(self, fgraph):
      nb_replacement = 0
  
      if fgraph.profile:
          validate_before = fgraph.profile.validate_time
          callbacks_before = fgraph.execute_callbacks_times.copy()
          callback_before = fgraph.execute_callbacks_time
  
      max_operands = elemwise_max_operands_fct(None)
  
      def find_next_fuseable_subgraph(
          fg: FunctionGraph,
      ) -> Generator[tuple[list[Variable], list[Variable]], None, None]:
          """Find all subgraphs in a FunctionGraph that can be fused together
  
          Yields
          -------
          List of inputs and outputs that determine subgraphs which can be fused.
          This generator assumes that such subgraph is replaced by a single
          Elemwise Composite before being accessed again in the next iteration.
          """
  
          FUSEABLE_MAPPING = defaultdict[Variable, list[Apply]]
          UNFUSEABLE_MAPPING = defaultdict[Variable, set[ApplyOrOutput]]
  
          def initialize_fuseable_mappings(
              *, fg: FunctionGraph
          ) -> tuple[FUSEABLE_MAPPING, UNFUSEABLE_MAPPING]:
              @cache
              def elemwise_scalar_op_has_c_code(node: Apply) -> bool:
                  # TODO: This should not play a role in non-c backends!
                  if node.op.scalar_op.supports_c_code(
                      node.inputs, node.outputs
                  ):
                      return True
                  else:
                      warn(
                          "Optimization Warning: "
                          f"The Op {node.op.scalar_op} does not provide a C implementation."
                          " As well as being potentially slow, this also disables "
                          "loop fusion."
                      )
                      return False
  
              # Fuseable nodes have to be accessed in a deterministic manner
              # to ensure the rewrite remains deterministic.
              # This is not a problem from unfuseable ones, as they can never
              # become part of the graph.
              fuseable_clients: FUSEABLE_MAPPING = defaultdict(list)
              unfuseable_clients: UNFUSEABLE_MAPPING = defaultdict(set)
              for out, clients in fg.clients.items():
                  # Old FunctionGraph nodes remain in the clients dictionary
                  # even after they are removed by rewrites
                  if not clients:
                      continue
  
                  out_maybe_fuseable = (
                      out.owner
                      and isinstance(out.owner.op, Elemwise)
                      # and not isinstance(out.owner.op.scalar_op, ps.Composite)
                      and len(out.owner.outputs) == 1
                      and elemwise_scalar_op_has_c_code(out.owner)
                  )
                  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
                          and out.type.broadcastable
                          == client.outputs[0].type.broadcastable
                          and elemwise_scalar_op_has_c_code(client)
                      ):
                          if client not in fuseable_clients[out]:
                              fuseable_clients[out].append(client)
                      else:
                          unfuseable_clients[out].add(client)
  
              return fuseable_clients, unfuseable_clients
  
          def find_fuseable_subgraph(
              *,
              fg: FunctionGraph,
              visited_nodes: set[Apply],
              fuseable_clients: FUSEABLE_MAPPING,
              unfuseable_clients: UNFUSEABLE_MAPPING,
          ) -> tuple[list[Variable], list[Variable]]:
              KT = TypeVar("KT")
              VT = TypeVar("VT", list, set)
  
              def shallow_clone_defaultdict(
                  d: defaultdict[KT, VT],
              ) -> defaultdict[KT, VT]:
                  new_dict: defaultdict[KT, VT] = defaultdict(
                      d.default_factory
                  )
                  new_dict.update({k: v.copy() for k, v in d.items()})
                  return new_dict
  
              def variables_depend_on(
                  variables, depend_on, stop_search_at=None
              ) -> bool:
                  return any(
                      a in depend_on
                      for a in ancestors(
                          variables, blockers=stop_search_at
                      )
                  )
  
              toposort = fg.toposort()
              for starting_node in toposort:
                  if starting_node in visited_nodes:
                      continue
  
                  starting_out = starting_node.outputs[0]
                  if not fuseable_clients.get(starting_out):
                      visited_nodes.add(starting_node)
                      continue
  
                  subgraph_inputs: list[Variable] = []
                  subgraph_outputs: list[Variable] = []
                  unfuseable_clients_subgraph: set[Variable] = set()
  
                  # Shallow cloning of maps so that they can be manipulated in place
                  fuseable_clients_temp = shallow_clone_defaultdict(
                      fuseable_clients
                  )
                  unfuseable_clients_clone = shallow_clone_defaultdict(
                      unfuseable_clients
                  )
  
                  fuseable_nodes_to_visit = deque([starting_node])
  
                  # We now try to expand as much as possible towards the potentially
                  # fuseable clients and ancestors to detect the largest possible
                  # subgraph that can be Composed together into a single `Op`. The
                  # largest issue to watch out is for cyclical dependencies, where
                  # some inputs or clients may depend on other nodes of the same
                  # subgraph via a path that cannot be included in the Composite
                  # (unfuseable)
                  while fuseable_nodes_to_visit:
                      next_node = fuseable_nodes_to_visit.popleft()
                      visited_nodes.add(next_node)
                      next_out = next_node.outputs[0]
  
                      # If the output variable of next_node has no fuseable clients
                      # or has unfuseable clients, then next_node must become an output
                      # if it is to be fused.
                      must_become_output = (
                          next_out not in fuseable_clients_temp
                          or next_out in unfuseable_clients_clone
                      )
  
                      # We have backtracked to this node, and it may no longer be a viable output,
                      # so we remove it and check again as if we had never seen this node
                      if (
                          must_become_output
                          and next_out in subgraph_outputs
                      ):
                          subgraph_outputs.remove(next_out)
  
                      required_unfuseable_inputs = [
                          inp
                          for inp in next_node.inputs
                          if next_node
                          in unfuseable_clients_clone.get(inp, ())
                      ]
                      new_required_unfuseable_inputs = [
                          inp
                          for inp in required_unfuseable_inputs
                          if inp not in subgraph_inputs
                      ]
  
                      must_backtrack = False
                      if (
                          new_required_unfuseable_inputs
                          and subgraph_outputs
                      ):
                          # We need to check that any new inputs required by this node
                          # do not depend on other outputs of the current subgraph,
                          # via an unfuseable path.
                          if variables_depend_on(
                              [next_out],
                              depend_on=unfuseable_clients_subgraph,
                              stop_search_at=subgraph_outputs,
                          ):
                              must_backtrack = True
  
                      if not must_backtrack:
                          implied_unfuseable_clients = {
                              c
                              for client in unfuseable_clients_clone.get(
                                  next_out, ()
                              )
                              if not isinstance(client, str)  # "output"
                              for c in client.outputs
                          }
  
                          new_implied_unfuseable_clients = (
                              implied_unfuseable_clients
                              - unfuseable_clients_subgraph
                          )
  
                          if (
                              new_implied_unfuseable_clients
                              and subgraph_inputs
                          ):
                              # We need to check that any inputs of the current subgraph
                              # do not depend on other clients of this node,
                              # via an unfuseable path.
                              if variables_depend_on(
                                  subgraph_inputs,
                                  depend_on=new_implied_unfuseable_clients,
                              ):
                                  must_backtrack = True
  
                      if must_backtrack:
                          for inp in next_node.inputs:
                              if (
                                  inp.owner in visited_nodes
                                  # next_node could have the same input repeated
                                  and next_node
                                  in fuseable_clients_temp[inp]
                              ):
                                  fuseable_clients_temp[inp].remove(
                                      next_node
                                  )
                                  unfuseable_clients_clone[inp].add(
                                      next_node
                                  )
                                  # This input must become an output of the subgraph,
                                  # because it can't be merged with next_node.
                                  # We will revisit it to make sure this is safe.
                                  fuseable_nodes_to_visit.appendleft(
                                      inp.owner
                                  )
  
                          for client in fuseable_clients_temp[next_out]:
                              if client in visited_nodes:
                                  fuseable_clients_temp[
                                      next_out
                                  ].remove(client)
                                  unfuseable_clients_clone[
                                      next_out
                                  ].add(client)
                                  # next_out must become an input of the subgraph.
                                  # We will revisit any of its clients currently
                                  # in the subgraph to make sure this is safe.
                                  fuseable_nodes_to_visit.appendleft(
                                      client
                                  )
  
                          # Revisit node at a later time
                          visited_nodes.remove(next_node)
                          continue
  
                      # Adding next_node to subgraph does not result in any
                      # immediate dependency problems. Update subgraph
                      # mappings as if it next_node was part of it.
                      # Useless inputs will be removed by the useless Composite rewrite
                      for inp in new_required_unfuseable_inputs:
                          if inp not in subgraph_inputs:
                              subgraph_inputs.append(inp)
  
                      if must_become_output:
                          subgraph_outputs.append(next_out)
                          unfuseable_clients_subgraph.update(
                              new_implied_unfuseable_clients
                          )
  
                      # Expand through unvisited fuseable ancestors
                      for inp in sorted(
                          (
                              inp
                              for inp in next_node.inputs
                              if (
                                  inp not in required_unfuseable_inputs
                                  and inp.owner not in visited_nodes
                              )
                          ),
                          key=lambda inp: toposort.index(inp.owner),
                          reverse=True,
                      ):
                          fuseable_nodes_to_visit.appendleft(inp.owner)
  
                      # Expand through unvisited fuseable clients
                      for next_node in sorted(
                          (
                              node
                              for node in fuseable_clients_temp.get(
                                  next_out, ()
                              )
                              if node not in visited_nodes
                          ),
                          key=lambda node: toposort.index(node),
                      ):
                          fuseable_nodes_to_visit.append(next_node)
  
                  # Don't return if final subgraph is just the original Elemwise
                  if len(subgraph_outputs) == 1 and set(
                      subgraph_outputs[0].owner.inputs
                  ) == set(subgraph_inputs):
                      # Update global fuseable mappings
                      # No input was actually fuseable
                      for inp in starting_node.inputs:
                          if starting_node in fuseable_clients.get(
                              inp, ()
                          ):
                              fuseable_clients[inp].remove(
                                  starting_node
                              )
                              unfuseable_clients[inp].add(starting_node)
                      # No client was actually fuseable
                      unfuseable_clients[starting_out].update(
                          fuseable_clients.pop(starting_out, ())
                      )
                      continue
  
                  return subgraph_inputs, subgraph_outputs
              raise ValueError
  
          def update_fuseable_mappings_after_fg_replace(
              *,
              fg: FunctionGraph,
              visited_nodes: set[Apply],
              fuseable_clients: FUSEABLE_MAPPING,
              unfuseable_clients: UNFUSEABLE_MAPPING,
              starting_nodes: set[Apply],
          ) -> None:
              # Find new composite node and dropped intermediate nodes
              # by comparing the current fg.apply nodes with the cached
              # original nodes
              next_nodes = fg.apply_nodes
              (new_composite_node,) = next_nodes - starting_nodes
              dropped_nodes = starting_nodes - next_nodes
  
              # Remove intermediate Composite nodes from mappings
              for dropped_node in dropped_nodes:
                  (dropped_out,) = dropped_node.outputs
                  fuseable_clients.pop(dropped_out, None)
                  unfuseable_clients.pop(dropped_out, None)
                  visited_nodes.remove(dropped_node)
  
              # Update fuseable information for subgraph inputs
              for inp in subgraph_inputs:
                  if inp in fuseable_clients:
                      new_fuseable_clients = [
                          client
                          for client in fuseable_clients[inp]
                          if client not in dropped_nodes
                      ]
                      if new_fuseable_clients:
                          fuseable_clients[inp] = new_fuseable_clients
                      else:
                          fuseable_clients.pop(inp)
                  unfuseable_clients[inp] = (
                      unfuseable_clients[inp] - dropped_nodes
                  ) | {new_composite_node}
  
              # Update fuseable information for subgraph outputs
              for out in new_composite_node.outputs:
                  unfuseable_clients[out] = {
                      client for client, _ in fg.clients[out]
                  }
  
              visited_nodes.add(new_composite_node)
              return
  
          # We start by creating two maps, 1) from each node to each potentially
          # fuseable client (both nodes must be single output Elemwise with same
          # broadcast type) and 2) from each node to each certainly unfuseable
          # client (those that don't fit into 1))
          (
              fuseable_clients,
              unfuseable_clients,
          ) = initialize_fuseable_mappings(fg=fg)
          visited_nodes: set[Apply] = set()
          while True:
              starting_nodes = fg.apply_nodes.copy()
              try:
                  (
                      subgraph_inputs,
                      subgraph_outputs,
                  ) = find_fuseable_subgraph(
                      fg=fg,
                      visited_nodes=visited_nodes,
                      fuseable_clients=fuseable_clients,
                      unfuseable_clients=unfuseable_clients,
                  )
              except ValueError:
                  return
              else:
                  # The caller is now expected to update fg in place,
                  # by replacing the subgraph with a Composite Op
                  yield subgraph_inputs, subgraph_outputs
  
                  # This is where we avoid repeated work by using a stateful
                  # generator. For large models (as in `TestFusion.test_big_fusion`)
                  # this can provide huge speedups
                  update_fuseable_mappings_after_fg_replace(
                      fg=fg,
                      visited_nodes=visited_nodes,
                      fuseable_clients=fuseable_clients,
                      unfuseable_clients=unfuseable_clients,
                      starting_nodes=starting_nodes,
                  )
  
      for inputs, outputs in find_next_fuseable_subgraph(fgraph):
          if (len(inputs) + len(outputs)) > max_operands:
              warn(
                  "Loop fusion failed because the resulting node would exceed "
                  "the kernel argument limit."
              )
              break
  
          scalar_inputs, scalar_outputs = self.elemwise_to_scalar(
              inputs, outputs
          )
          composite_outputs = Elemwise(
              ps.Composite(scalar_inputs, scalar_outputs)
          )(*inputs)
          if not isinstance(composite_outputs, list):
              composite_outputs = [composite_outputs]
          for old_out, composite_out in zip(outputs, composite_outputs):
              if old_out.name:
                  composite_out.name = old_out.name
  
          fgraph.replace_all_validate(
              list(zip(outputs, composite_outputs)),
              reason=self.__class__.__name__,
          )
          nb_replacement += 1
  
      if fgraph.profile:
          validate_time = fgraph.profile.validate_time - validate_before
          callback_time = (
              fgraph.execute_callbacks_time - callback_before
          )
          callbacks_time = {}
          for k, v in fgraph.execute_callbacks_times.items():
              if k in callbacks_before:
                  callbacks_time[k] = v - callbacks_before[k]
              else:
                  callbacks_time[k] = v
      else:
          validate_time = None
          callback_time = None
          callbacks_time = {}
  
      return (
          self,
          1,  # nb_iter
          nb_replacement,
          0,  # nb_inconsintency_replace
          validate_time,
          callback_time,
          callbacks_time,
          -1,  # toposort_time
      )
  
imports: import sys
  from collections import defaultdict, deque
  from collections.abc import Generator
  from functools import cache
  from typing import TypeVar
  from warnings import warn
  import pytensor
  import pytensor.scalar.basic
  from pytensor import clone_replace, compile
  from pytensor.compile.mode import get_target_language
  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.rewriting.basic import EquilibriumGraphRewriter, GraphRewriter, copy_stack_trace, in2out, node_rewriter
  from pytensor.graph.rewriting.db import SequenceDB
  from pytensor.graph.utils import InconsistencyError, MethodNotDefined
  from pytensor.scalar.loop import ScalarLoop
  from pytensor.scalar.math import Grad2F1Loop, _grad_2f1_loop
  from pytensor.tensor.basic import MakeVector, alloc, cast, get_underlying_scalar_constant_value
  from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
  from pytensor.tensor.exceptions import NotScalarConstantError
  from pytensor.tensor.math import exp
  from pytensor.tensor.rewriting.basic import alloc_like, register_canonicalize, register_specialize
  from pytensor.tensor.shape import shape_padleft
  from pytensor.tensor.variable import TensorConstant, get_unique_constant_value
label_desc: Use image node_id 3 for calling the FusionOptimizer obj's underlying member method code with example usage: obj.apply(fgraph) and returns: self,  int,  nb_replacement,  int,  validate_time,  callback_time,  callbacks_time,  unknown, fuseable_clients,  unfuseable_clients, new_dict, any, subgraph_inputs,  subgraph_outputs, True, False