class_docstr: Graph optimizer that merges different scan ops.
  
  This optimization attempts to fuse distinct `Scan` `Op`s into a single `Scan` `Op`
  that performs all the computation. The main advantage of merging `Scan` `Op`\s
  together comes from the possibility of both original `Op`\s having some
  computation in common. In such a setting, this computation ends up being done
  twice. The fused `Scan` `Op`, however, would only need to do it once and could
  therefore be more computationally efficient. Also, since every `Scan` node
  involves a certain overhead, at runtime, reducing the number of `Scan` nodes in
  the graph can improve performance.
class_docstr_tok: ["Graph","optimizer","that","merges","different","scan","ops",".","This","optimization","attempts","to","fuse","distinct","`","Scan","`","`","Op","`","s","into","a","single","`","Scan","`","`","Op","`","that","performs","all","the","computation",".","The","main","advantage","of","merging","`","Scan","`","`","Op","`","\\s","together","comes","from","the","possibility","of","both","original","`","Op","`","\\s","having","some","computation","in","common",".","In","such","a","setting",",","this","computation","ends","up","being","done","twice",".","The","fused","`","Scan","`","`","Op","`",",","however",",","would","only","need","to","do","it","once","and","could","therefore","be","more","computationally","efficient",".","Also",",","since","every","`","Scan","`","node","involves","a","certain","overhead",",","at","runtime",",","reducing","the","number","of","`","Scan","`","nodes","in","the","graph","can","improve","performance","."]
docstr: This function checks if node `node` belongs to `set_nodes`, in the
  sense that it can be merged together with every other node in
  `set_nodes`. In order for two nodes to be mergeable, they have to go
  over the same number of steps, have the same condition (if any),
  have the same value for truncate_gradient, and have the same mode.
  Questionable, we should also consider profile ?
docstr_tok: ["This","function","checks","if","node","`","node","`","belongs","to","`","set_nodes","`",",","in","the","sense","that","it","can","be","merged","together","with","every","other","node","in","`","set_nodes","`",".","In","order","for","two","nodes","to","be","mergeable",",","they","have","to","go","over","the","same","number","of","steps",",","have","the","same","condition","(","if","any",")",",","have","the","same","value","for","truncate_gradient",",","and","have","the","same","mode",".","Questionable",",","we","should","also","consider","profile","?"]
code: def belongs_to_set(self, node, set_nodes):
      """
      This function checks if node `node` belongs to `set_nodes`, in the
      sense that it can be merged together with every other node in
      `set_nodes`. In order for two nodes to be mergeable, they have to go
      over the same number of steps, have the same condition (if any),
      have the same value for truncate_gradient, and have the same mode.
      Questionable, we should also consider profile ?
  
      """
      op = node.op
      rep_node = set_nodes[0]
      rep_op = rep_node.op
      if (
          op.info.as_while != rep_op.info.as_while
          or op.truncate_gradient != rep_op.truncate_gradient
          or op.mode != rep_op.mode
      ):
          return False
  
      nsteps = node.inputs[0]
      try:
          nsteps = int(get_underlying_scalar_constant_value(nsteps))
      except NotScalarConstantError:
          pass
  
      rep_nsteps = rep_node.inputs[0]
      try:
          rep_nsteps = int(
              get_underlying_scalar_constant_value(rep_nsteps)
          )
      except NotScalarConstantError:
          pass
  
      if nsteps != rep_nsteps:
          return False
  
      # Check to see if it is an input of a different node
      for nd in set_nodes:
          if apply_depends_on(node, nd) or apply_depends_on(nd, node):
              return False
  
      if not op.info.as_while:
          return True
  
      # We need to check the while conditions are identical
      conds = [op.inner_outputs[-1]]
      rep_conds = [rep_op.inner_outputs[-1]]
      if not equal_computations(
          conds, rep_conds, op.inner_inputs, rep_op.inner_inputs
      ):
          return False
  
      # If they depend on inner inputs we need to check for equivalence on the respective outer inputs
      nominal_inputs = [
          a for a in ancestors(conds) if isinstance(a, NominalVariable)
      ]
      if not nominal_inputs:
          return True
      rep_nominal_inputs = [
          a
          for a in ancestors(rep_conds)
          if isinstance(a, NominalVariable)
      ]
  
      conds = []
      rep_conds = []
      mapping = op.get_oinp_iinp_iout_oout_mappings()[
          "outer_inp_from_inner_inp"
      ]
      rep_mapping = rep_op.get_oinp_iinp_iout_oout_mappings()[
          "outer_inp_from_inner_inp"
      ]
      inner_inputs = op.inner_inputs
      rep_inner_inputs = rep_op.inner_inputs
      for nominal_input, rep_nominal_input in zip(
          nominal_inputs, rep_nominal_inputs
      ):
          conds.append(
              node.inputs[mapping[inner_inputs.index(nominal_input)]]
          )
          rep_conds.append(
              rep_node.inputs[
                  rep_mapping[rep_inner_inputs.index(rep_nominal_input)]
              ]
          )
  
      return equal_computations(conds, rep_conds)
  
imports: import copy
  import dataclasses
  from itertools import chain
  from sys import maxsize
  from typing import Optional, cast
  import numpy
  import pytensor
  from pytensor import scalar
  from pytensor import tensor
  from pytensor.compile import optdb
  from pytensor.compile.function.types import deep_copy_op
  from pytensor.configdefaults import config
  from pytensor.graph.basic import Apply, Constant, NominalVariable, Variable, ancestors, apply_depends_on, equal_computations, graph_inputs, io_toposort
  from pytensor.graph.destroyhandler import DestroyHandler
  from pytensor.graph.features import ReplaceValidate
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import compute_test_value
  from pytensor.graph.replace import clone_replace
  from pytensor.graph.rewriting.basic import GraphRewriter, copy_stack_trace, in2out, node_rewriter
  from pytensor.graph.rewriting.db import EquilibriumDB, SequenceDB
  from pytensor.graph.rewriting.utils import get_clients_at_depth
  from pytensor.graph.type import HasShape
  from pytensor.graph.utils import InconsistencyError
  from pytensor.raise_op import Assert
  from pytensor.scalar import ScalarConstant
  from pytensor.scan.op import Scan, ScanInfo
  from pytensor.scan.utils import ScanArgs, compress_outs, expand_empty, reconstruct_graph, safe_new, scan_can_remove_outs
  from pytensor.tensor.basic import Alloc, AllocEmpty, get_underlying_scalar_constant_value
  from pytensor.tensor.elemwise import DimShuffle, Elemwise
  from pytensor.tensor.exceptions import NotScalarConstantError
  from pytensor.tensor.math import Dot, dot, maximum, minimum
  from pytensor.tensor.rewriting.basic import constant_folding, local_useless_switch
  from pytensor.tensor.rewriting.elemwise import local_upcast_elemwise_constant_inputs
  from pytensor.tensor.rewriting.math import local_abs_merge, local_mul_switch_sink
  from pytensor.tensor.shape import shape
  from pytensor.tensor.subtensor import IncSubtensor, Subtensor, get_canonical_form_slice, get_idx_list, get_slice_elements, set_subtensor
  from pytensor.tensor.variable import TensorConstant, get_unique_constant_value
label_desc: Use image node_id 3 for calling the ScanMerge obj's underlying member method code with example usage: obj.belongs_to_set(node, set_nodes) and returns: equal_computations, False, False, True, False, True, False