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","."]
code: def merge(self, nodes):
      if nodes[0].op.info.as_while:
          as_while = True
          condition = nodes[0].op.inner_outputs[-1]
      else:
          as_while = False
  
      # We keep the inner_ins and inner_outs of each original node separated.
      # To be able to recombine them in the right order after the clone,
      # we also need to split them by types (seq, mitmot, ...).
      # On the other hand, outer_ins, outer_outs and info are held together.
      inner_ins = [[] for nd in nodes]
      outer_ins = []
      inner_outs = [[] for nd in nodes]
      outer_outs = []
  
      for idx, nd in enumerate(nodes):
          inner_ins[idx].append(nd.op.inner_seqs(nd.op.inner_inputs))
          outer_ins += nd.op.outer_seqs(nd.inputs)
  
      mit_mot_out_slices = ()
  
      mit_mot_in_slices = ()
      for idx, nd in enumerate(nodes):
          inner_ins[idx].append(nd.op.inner_mitmot(nd.op.inner_inputs))
          inner_outs[idx].append(
              nd.op.inner_mitmot_outs(nd.op.inner_outputs)
          )
          mit_mot_in_slices += nd.op.info.mit_mot_in_slices
          mit_mot_out_slices += nd.op.info.mit_mot_out_slices[
              : nd.op.info.n_mit_mot
          ]
          outer_ins += nd.op.outer_mitmot(nd.inputs)
          outer_outs += nd.op.outer_mitmot_outs(nd.outputs)
  
      mit_sot_in_slices = ()
      for idx, nd in enumerate(nodes):
          inner_ins[idx].append(nd.op.inner_mitsot(nd.op.inner_inputs))
          inner_outs[idx].append(
              nd.op.inner_mitsot_outs(nd.op.inner_outputs)
          )
          mit_sot_in_slices += nd.op.info.mit_sot_in_slices
          outer_ins += nd.op.outer_mitsot(nd.inputs)
          outer_outs += nd.op.outer_mitsot_outs(nd.outputs)
  
      sit_sot_in_slices = ()
      for idx, nd in enumerate(nodes):
          inner_ins[idx].append(nd.op.inner_sitsot(nd.op.inner_inputs))
          sit_sot_in_slices += tuple(
              (-1,) for x in range(nd.op.info.n_sit_sot)
          )
          inner_outs[idx].append(
              nd.op.inner_sitsot_outs(nd.op.inner_outputs)
          )
          outer_ins += nd.op.outer_sitsot(nd.inputs)
          outer_outs += nd.op.outer_sitsot_outs(nd.outputs)
  
      for idx, nd in enumerate(nodes):
          # Shared
          inner_ins[idx].append(nd.op.inner_shared(nd.op.inner_inputs))
          outer_ins += nd.op.outer_shared(nd.inputs)
  
      for idx, nd in enumerate(nodes):
          # NitSot
          inner_outs[idx].append(
              nd.op.inner_nitsot_outs(nd.op.inner_outputs)
          )
          outer_ins += nd.op.outer_nitsot(nd.inputs)
          outer_outs += nd.op.outer_nitsot_outs(nd.outputs)
  
      for idx, nd in enumerate(nodes):
          # Shared
          outer_outs += nd.op.outer_shared_outs(nd.outputs)
          inner_outs[idx].append(
              nd.op.inner_shared_outs(nd.op.inner_outputs)
          )
  
      n_non_seqs = 0
      for idx, nd in enumerate(nodes):
          # Non Seqs
          node_inner_non_seqs = nd.op.inner_non_seqs(nd.op.inner_inputs)
          n_non_seqs += len(node_inner_non_seqs)
          inner_ins[idx].append(node_inner_non_seqs)
          outer_ins += nd.op.outer_non_seqs(nd.inputs)
  
      # Add back the number of steps
      outer_ins = [nodes[0].inputs[0]] + outer_ins
  
      if as_while:
          # add the condition, which was the one of nodes[0]
          inner_outs[0].append([condition])
  
      # Clone the inner graph of each node independently
      for idx, nd in enumerate(nodes):
          # concatenate all inner_ins and inner_outs of nd
          flat_inner_ins = sum(inner_ins[idx], [])
          flat_inner_outs = sum(inner_outs[idx], [])
          # clone
          flat_inner_ins, flat_inner_outs = reconstruct_graph(
              flat_inner_ins, flat_inner_outs
          )
          # split the new inner variables again in seq, mitmot, etc.
          new_inner_ins = []
          count = 0
          for nl in inner_ins[idx]:
              seq_len = len(nl)
              new_inner_ins.append(
                  flat_inner_ins[count : (count + seq_len)]
              )
              count += seq_len
  
          new_inner_outs = []
          count = 0
          for nl in inner_outs[idx]:
              seq_len = len(nl)
              new_inner_outs.append(
                  flat_inner_outs[count : (count + seq_len)]
              )
              count += seq_len
  
          inner_ins[idx] = new_inner_ins
          inner_outs[idx] = new_inner_outs
  
      # Flatten inner_ins and inner_outs so that all seqs are first,
      # then mitmot, etc.
      new_inner_ins = []
      new_inner_outs = []
      nb_ins_groups = len(inner_ins[0])
      nb_outs_groups = len(inner_outs[0])
      for idx, nd in enumerate(nodes):
          # All inner_ins should have the same length
          assert len(inner_ins[idx]) == nb_ins_groups
  
          # All inner_outs should have the same length, except if as_while,
          # in which case the first one should have one more element
          if as_while and idx > 0:
              assert len(inner_outs[idx]) == nb_outs_groups - 1
          else:
              assert len(inner_outs[idx]) == nb_outs_groups
  
      for gr_idx in range(nb_ins_groups):
          for idx, nd in enumerate(nodes):
              new_inner_ins += inner_ins[idx][gr_idx]
  
      for gr_idx in range(nb_outs_groups):
          for idx, nd in enumerate(nodes):
              if (
                  as_while
                  and idx > 0
                  and gr_idx == (nb_outs_groups - 1)
              ):
                  # There is no condition on that node, skip it
                  pass
              else:
                  new_inner_outs += inner_outs[idx][gr_idx]
  
      info = ScanInfo(
          n_seqs=sum(nd.op.info.n_seqs for nd in nodes),
          mit_mot_in_slices=mit_mot_in_slices,
          mit_mot_out_slices=mit_mot_out_slices,
          mit_sot_in_slices=mit_sot_in_slices,
          sit_sot_in_slices=sit_sot_in_slices,
          n_nit_sot=sum(nd.op.info.n_nit_sot for nd in nodes),
          n_shared_outs=sum(nd.op.info.n_shared_outs for nd in nodes),
          n_non_seqs=n_non_seqs,
          as_while=as_while,
      )
  
      old_op = nodes[0].op
      new_op = Scan(
          new_inner_ins,
          new_inner_outs,
          info,
          mode=old_op.mode,
          profile=old_op.profile,
          truncate_gradient=old_op.truncate_gradient,
          allow_gc=old_op.allow_gc,
          name="&".join([nd.op.name for nd in nodes]),
      )
      new_outs = new_op(*outer_ins)
  
      if not isinstance(new_outs, (list, tuple)):
          new_outs = [new_outs]
  
      return list(zip(outer_outs, new_outs))
  
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 2 for calling the ScanMerge obj's underlying member method code with example usage: obj.merge(nodes) and returns: list