text: # Copyright (c) Microsoft Corporation. All rights reserved.
  # Licensed under the MIT license.
  
  """Back_To_Back Optimizer.
     Collapse consecutive nodes into 1 node if possible.
  """
  
  from __future__ import unicode_literals
  
  import numpy as np
  from tf2onnx.utils import ONNX_DTYPE_NAMES  # lgtm[py/unsafe-cyclic-import]
  from .optimizer_base import GraphOptimizerBase  # lgtm[py/unsafe-cyclic-import]
  
  # pylint: disable=logging-not-lazy,unused-argument,missing-docstring,unused-variable,arguments-differ
  
  _func_map = {}
  
  
  def _register_func(op_type):
      def _internal_fun(func):
          _func_map[op_type] = func
          return func
  
      return _internal_fun
  
  
  class BackToBackOptimizer(GraphOptimizerBase):
      """Remove back-to-back nodes e.g. 'Cast'
      """
  
      def __init__(self):  # pylint: disable=useless-super-delegation
          super(BackToBackOptimizer, self).__init__()
  
      def _optimize(self, graph):
          return self._apply_optimization(graph, self._optimize_at_current_graph_level)
  
      def _optimize_at_current_graph_level(self, g):
          for optype, handler in _func_map.items():
              # candidate nodes for removal/optimization
              nodes = [n for n in g.get_nodes() if n.type in optype]
  
              # topological sort of candidates
              # simplifying assumption for back-to-back-optimizer is
              # the op_types have 1 input, 1 output, but multiple consumers
              has_dependencies = set()
              consumer_node_ids = {n.output[0]: [] for n in nodes}
              for n in nodes:
                  if n.input[0] in consumer_node_ids:
                      consumer_node_ids[n.input[0]].extend([n])
                      has_dependencies.add(n.output[0])
  
              # q = starting nodes with no dependencies
              q = list(set(consumer_node_ids.keys()) - has_dependencies)
              while q:
                  nodeid = q.pop(0)
                  node = g.get_node_by_output(nodeid, False)
                  consumer_nodes = consumer_node_ids[nodeid]
  
                  if len(consumer_nodes) > 0:
                      all_consumers = g.find_output_consumers(node.output[0])
                      if len(all_consumers) != len(consumer_nodes):
                          # if first node is used elsewhere, skip
                          continue
                      if set(node.output) & set(g.outputs):
                          # if this node is part of graph outputs, skip
                          continue
                      q2 = handler(g, node, consumer_nodes)
                      # add more nodes which can now be processed
                      q.extend(q2)
          return g
  
      @staticmethod
      @_register_func("Cast")
      def _optimize_cast(g, node, consumer_nodes):
          """remove long chains of cast ops"""
          q2 = []
          type1 = node.get_attr('to').i
          type1_name = ONNX_DTYPE_NAMES[type1] if type1 in ONNX_DTYPE_NAMES else ''
  
          # if parent node is cast node, and same type, delete this one
          pnode = node.inputs[0]
          if pnode.type == 'Cast':
              type2 = pnode.get_attr('to').i
              if type1 == type2:
                  for node2 in consumer_nodes:
                      g.replace_input(node2, node2.input[0], node.input[0], 0)
                      q2.append(node2.output[0])
                  g.remove_node(node.name)
                  return q2
  
          # otherwise, check consumer cast nodes for a target type
          # that contains more information than current type
          can_reduce = True
          for node2 in consumer_nodes:
              type2 = node2.get_attr('to').i
              type2_name = ONNX_DTYPE_NAMES[type2] if type2 in ONNX_DTYPE_NAMES else ''
  
              if 'float' in type1_name or type1_name == 'double':
                  # high information type. ok to eliminate
                  pass
              elif 'int' in type1_name:
                  # int* and uint* are mix of high and low information.
                  # for safety, keep the current node, unless type2 is bool,
                  # in which case it's ok to remove node
                  if type1 != type2 and type2_name != 'bool':
                      can_reduce = False
              elif type1_name == 'bool':
                  # bool is low information, so don't eliminate
                  if type1 != type2:
                      can_reduce = False
              elif type1_name == 'string':
                  # can always remove string
                  pass
              else:
                  # some odd type, keep node
                  can_reduce = False
              q2.append(node2.output[0])
  
          if can_reduce:
              for node2 in consumer_nodes:
                  g.replace_input(node2, node2.input[0], node.input[0], 0)
              g.remove_node(node.name)
          return q2
  
      @staticmethod
      @_register_func("Transpose")
      def _optimize_transpose(g, node, consumer_nodes):
          """remove long chains of transpose ops"""
          t1 = list(node.get_attr('perm').ints)
          q2 = []
          for node2 in consumer_nodes:
              g.replace_input(node2, node2.input[0], node.input[0], 0)
              t2 = list(node2.get_attr('perm').ints)
              new_perm = [t1[i] for i in t2]
              # check if node2 can be removed. otherwise only update
              if new_perm == list(range(len(t2))):
                  # both nodes can be deleted
                  shape = g.get_shape(node2.output[0])
                  dtype = g.get_dtype(node2.output[0])
                  node2_consumers = g.find_output_consumers(node2.output[0])
                  g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0])
                  g.remove_node(node2.name)
                  if set(node2.output) & set(g.outputs):
                      g.make_node("Identity", [node.input[0]],
                                  outputs=node2.output, shapes=[shape], dtypes=[dtype])
              else:
                  node2.set_attr('perm', [t1[i] for i in t2])
                  q2.append(node2.output[0])
          g.remove_node(node.name)
          return q2
  
      @staticmethod
      @_register_func(('Squeeze', 'Unsqueeze'))
      def _optimize_squeeze_unsqueeze(g, node, consumer_nodes):
          """remove pairs of squeeze-unsqueeze nodes"""
          if node.type != 'Squeeze' or len(consumer_nodes) != 1:
              # no need to return any value, since not removing long chain of nodes
              return []
  
          node2 = consumer_nodes[0]
          if node2.type != 'Unsqueeze':
              return []
  
          axis1 = node.get_attr('axes').ints
          axis2 = node2.get_attr('axes').ints
  
          # if squeeze followed by unsqueeze is on diff axes, skip
          if axis1 != axis2:
              return []
  
          # if unsqueeze output is graph output, skip
          if set(node2.output) & set(g.outputs):
              return []
  
          node2_consumers = g.find_output_consumers(node2.output[0])
          g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0])
          g.remove_node(node.name)
          g.remove_node(node2.name)
          return []
  
      @staticmethod
      @_register_func(('Conv', 'BatchNormalization'))
      def _optimize_conv_batchnorm_fusion(g, node, consumer_nodes):
          """fuse conv and batchnorm"""
          if node.type != 'Conv' or len(consumer_nodes) != 1:
              # can only fuse 1 conv + batchnorm
              return []
  
          node2 = consumer_nodes[0]
          if node2.type != 'BatchNormalization':
              return []
  
          # if batchnorm is a graph output, skip
          if set(node2.output) & set(g.outputs):
              return []
  
          if not node.inputs[1].is_const():
              return []
          weights = node.inputs[1].get_tensor_value(as_list=False)
          # if not 4D, NCHW skip
          if len(weights.shape) != 4:
              return []
  
          bias = 0
          # optional bias value
          if len(node.inputs) > 2:
              if not node.inputs[2].is_const():
                  return []
              bias = node.inputs[2].get_tensor_value(as_list=False)
  
          # scale, offset, mean, var be const, otherwise skip
          if False in [node2.inputs[i].is_const() for i in [1, 2, 3, 4]]:
              return []
  
          # if bn outputs used elsewhere, cannot fuse
          for i in range(1, len(node2.output)):
              if g.find_output_consumers(node2.output[i]):
                  return []
  
          weights = weights.transpose(2, 3, 1, 0)
          scale = node2.inputs[1].get_tensor_value(as_list=False)
          offset = node2.inputs[2].get_tensor_value(as_list=False)
          mean = node2.inputs[3].get_tensor_value(as_list=False)
          var = node2.inputs[4].get_tensor_value(as_list=False)
          epsilon = node2.get_attr('epsilon').f
  
          scale_new = scale / np.sqrt(var + epsilon)
          weights_new = weights * scale_new
          weights_new = weights_new.transpose(3, 2, 0, 1)
          bias_new = (bias - mean) * scale_new + offset
          bias_new_const = g.make_const(node.name + '_bias_fused_bn', bias_new)
          weights_new_const = g.make_const(node.name + '_weights_fused_bn', weights_new)
          g.replace_inputs(node, [node.input[0], weights_new_const.output[0], bias_new_const.output[0]])
  
          # fuse conv and bn, delete bn
          node2_output = node2.output[:1]
          node2_shape = g.get_shape(node2.output[0])
          node2_dtype = g.get_dtype(node2.output[0])
          g.remove_node(node2.name)
          # the setter makes a copy
          node.output = node2_output
          g.set_shape(node2_output[0], node2_shape)
          g.set_dtype(node2_output[0], node2_dtype)
          return []
  