text: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  # Copyright 2020 Huawei Technologies Co., Ltd
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import collections
  import contextlib
  import copy
  import functools
  import logging
  import mock
  import numpy as np
  import os
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  import caffe2.python.utils as putils
  import torch
  import torch.nn.functional as F
  from caffe2.proto import caffe2_pb2
  from caffe2.python import core, net_drawer, workspace
  from torch.nn.functional import interpolate as interp
  
  logger = logging.getLogger(__name__)
  
  
  # ==== torch/utils_toffee/cast.py =======================================
  
  
  def to_device(t, device_str):
      """
      This function is a replacement of .to(another_device) such that it allows the
      casting to be traced properly by explicitly calling the underlying copy ops.
      It also avoids introducing unncessary op when casting to the same device.
      """
      src = t.device
      dst = torch.device(device_str)
  
      if src == dst:
          return t
      elif src.type == "cuda" and dst.type == "cpu":
          return torch.ops._caffe2.CopyGPUToCPU(t)
      elif src.type == "cpu" and dst.type == "cuda":
          return torch.ops._caffe2.CopyCPUToGPU(t)
      else:
          raise RuntimeError("Can't cast tensor from device {} to device {}".format(src, dst))
  
  
  # ==== torch/utils_toffee/interpolate.py =======================================
  
  
  # Note: borrowed from vision/detection/fair/detectron/detectron/modeling/detector.py
  def BilinearInterpolation(tensor_in, up_scale):
      assert up_scale % 2 == 0, "Scale should be even"
  
      def upsample_filt(size):
          factor = (size + 1) // 2
          if size % 2 == 1:
              center = factor - 1
          else:
              center = factor - 0.5
  
          og = np.ogrid[:size, :size]
          return (1 - abs(og[0] - center) / factor) * (1 - abs(og[1] - center) / factor)
  
      kernel_size = int(up_scale) * 2
      bil_filt = upsample_filt(kernel_size)
  
      dim = int(tensor_in.shape[1])
      kernel = np.zeros((dim, dim, kernel_size, kernel_size), dtype=np.float32)
      kernel[range(dim), range(dim), :, :] = bil_filt
  
      tensor_out = F.conv_transpose2d(
          tensor_in,
          weight=to_device(torch.Tensor(kernel), tensor_in.device),
          bias=None,
          stride=int(up_scale),
          padding=int(up_scale / 2),
      )
  
      return tensor_out
  
  
  # NOTE: ONNX is incompatible with traced torch.nn.functional.interpolate if
  # using dynamic `scale_factor` rather than static `size`. (T43166860)
  # NOTE: Caffe2 Int8 conversion might not be able to quantize `size` properly.
  def onnx_compatibale_interpolate(
      input, size=None, scale_factor=None, mode="nearest", align_corners=None
  ):
      # NOTE: The input dimensions are interpreted in the form:
      # `mini-batch x channels x [optional depth] x [optional height] x width`.
      if size is None and scale_factor is not None:
          if input.dim() == 4:
              if isinstance(scale_factor, (int, float)):
                  height_scale, width_scale = (scale_factor, scale_factor)
              else:
                  assert isinstance(scale_factor, (tuple, list))
                  assert len(scale_factor) == 2
                  height_scale, width_scale = scale_factor
  
              assert not align_corners, "No matching C2 op for align_corners == True"
              if mode == "nearest":
                  input = input.cpu()
                  return torch.ops._caffe2.ResizeNearest(
                      input, order="NCHW", width_scale=width_scale, height_scale=height_scale
                  )
              elif mode == "bilinear":
                  logger.warning(
                      "Use F.conv_transpose2d for bilinear interpolate"
                      " because there's no such C2 op, this may cause significant"
                      " slowdown and the boundary pixels won't be as same as"
                      " using F.interpolate due to padding."
                  )
                  assert height_scale == width_scale
                  return BilinearInterpolation(input, up_scale=height_scale)
          logger.warning("Output size is not static, it might cause ONNX conversion issue")
  
      return interp(input, size, scale_factor, mode, align_corners)
  
  
  @contextlib.contextmanager
  def mock_torch_nn_functional_interpolate():
      if torch.onnx.is_in_onnx_export():
          with mock.patch(
              "torch.nn.functional.interpolate", side_effect=onnx_compatibale_interpolate
          ):
              yield
      else:
          yield
  
  
  # ==== torch/utils_caffe2/ws_utils.py ==========================================
  
  
  class ScopedWS(object):
      def __init__(self, ws_name, is_reset, is_cleanup=False):
          self.ws_name = ws_name
          self.is_reset = is_reset
          self.is_cleanup = is_cleanup
          self.org_ws = ""
  
      def __enter__(self):
          self.org_ws = workspace.CurrentWorkspace()
          if self.ws_name is not None:
              workspace.SwitchWorkspace(self.ws_name, True)
          if self.is_reset:
              workspace.ResetWorkspace()
  
          return workspace
  
      def __exit__(self, *args):
          if self.is_cleanup:
              workspace.ResetWorkspace()
          if self.ws_name is not None:
              workspace.SwitchWorkspace(self.org_ws)
  
  
  def fetch_any_blob(name):
      bb = None
      try:
          bb = workspace.FetchBlob(name)
      except TypeError:
          bb = workspace.FetchInt8Blob(name)
      except Exception as e:
          logger.error("Get blob {} error: {}".format(name, e))
  
      return bb
  
  
  # ==== torch/utils_caffe2/protobuf.py ==========================================
  
  
  def get_pb_arg(pb, arg_name):
      for x in pb.arg:
          if x.name == arg_name:
              return x
      return None
  
  
  def get_pb_arg_valf(pb, arg_name, default_val):
      arg = get_pb_arg(pb, arg_name)
      return arg.f if arg is not None else default_val
  
  
  def get_pb_arg_floats(pb, arg_name, default_val):
      arg = get_pb_arg(pb, arg_name)
      return list(map(float, arg.floats)) if arg is not None else default_val
  
  
  def get_pb_arg_ints(pb, arg_name, default_val):
      arg = get_pb_arg(pb, arg_name)
      return list(map(int, arg.ints)) if arg is not None else default_val
  
  
  def get_pb_arg_vali(pb, arg_name, default_val):
      arg = get_pb_arg(pb, arg_name)
      return arg.i if arg is not None else default_val
  
  
  def get_pb_arg_vals(pb, arg_name, default_val):
      arg = get_pb_arg(pb, arg_name)
      return arg.s if arg is not None else default_val
  
  
  def get_pb_arg_valstrings(pb, arg_name, default_val):
      arg = get_pb_arg(pb, arg_name)
      return list(arg.strings) if arg is not None else default_val
  
  
  def check_set_pb_arg(pb, arg_name, arg_attr, arg_value, allow_override=False):
      arg = get_pb_arg(pb, arg_name)
      if arg is None:
          arg = putils.MakeArgument(arg_name, arg_value)
          assert hasattr(arg, arg_attr)
          pb.arg.extend([arg])
      if allow_override and getattr(arg, arg_attr) != arg_value:
          logger.warning(
              "Override argument {}: {} -> {}".format(arg_name, getattr(arg, arg_attr), arg_value)
          )
          setattr(arg, arg_attr, arg_value)
      else:
          assert arg is not None
          assert getattr(arg, arg_attr) == arg_value, "Existing value {}, new value {}".format(
              getattr(arg, arg_attr), arg_value
          )
  
  
  def _create_const_fill_op_from_numpy(name, tensor, device_option=None):
      assert type(tensor) == np.ndarray
      kTypeNameMapper = {
          np.dtype("float32"): "GivenTensorFill",
          np.dtype("int32"): "GivenTensorIntFill",
          np.dtype("int64"): "GivenTensorInt64Fill",
          np.dtype("uint8"): "GivenTensorStringFill",
      }
  
      args_dict = {}
      if tensor.dtype == np.dtype("uint8"):
          args_dict.update({"values": [str(tensor.data)], "shape": [1]})
      else:
          args_dict.update({"values": tensor, "shape": tensor.shape})
  
      if device_option is not None:
          args_dict["device_option"] = device_option
  
      return core.CreateOperator(kTypeNameMapper[tensor.dtype], [], [name], **args_dict)
  
  
  def _create_const_fill_op_from_c2_int8_tensor(name, int8_tensor):
      assert type(int8_tensor) == workspace.Int8Tensor
      kTypeNameMapper = {
          np.dtype("int32"): "Int8GivenIntTensorFill",
          np.dtype("uint8"): "Int8GivenTensorFill",
      }
  
      tensor = int8_tensor.data
      assert tensor.dtype in [np.dtype("uint8"), np.dtype("int32")]
      values = tensor.tobytes() if tensor.dtype == np.dtype("uint8") else tensor
  
      return core.CreateOperator(
          kTypeNameMapper[tensor.dtype],
          [],
          [name],
          values=values,
          shape=tensor.shape,
          Y_scale=int8_tensor.scale,
          Y_zero_point=int8_tensor.zero_point,
      )
  
  
  def create_const_fill_op(
      name: str,
      blob: Union[np.ndarray, workspace.Int8Tensor],
      device_option: Optional[caffe2_pb2.DeviceOption] = None,
  ) -> caffe2_pb2.OperatorDef:
      """
      Given a blob object, return the Caffe2 operator that creates this blob
      as constant. Currently support NumPy tensor and Caffe2 Int8Tensor.
      """
  
      tensor_type = type(blob)
      assert tensor_type in [
          np.ndarray,
          workspace.Int8Tensor,
      ], 'Error when creating const fill op for "{}", unsupported blob type: {}'.format(
          name, type(blob)
      )
  
      if tensor_type == np.ndarray:
          return _create_const_fill_op_from_numpy(name, blob, device_option)
      elif tensor_type == workspace.Int8Tensor:
          assert device_option is None
          return _create_const_fill_op_from_c2_int8_tensor(name, blob)
  
  
  def construct_init_net_from_params(
      params: Dict[str, Any], device_options: Optional[Dict[str, caffe2_pb2.DeviceOption]] = None
  ) -> caffe2_pb2.NetDef:
      """
      Construct the init_net from params dictionary
      """
      init_net = caffe2_pb2.NetDef()
      device_options = device_options or {}
      for name, blob in params.items():
          if isinstance(blob, str):
              logger.warning(
                  (
                      "Blob {} with type {} is not supported in generating init net,"
                      " skipped.".format(name, type(blob))
                  )
              )
              continue
          init_net.op.extend(
              [create_const_fill_op(name, blob, device_option=device_options.get(name, None))]
          )
          init_net.external_output.append(name)
      return init_net
  
  
  def get_producer_map(ssa):
      """
      Return dict from versioned blob to (i, j),
          where i is index of producer op, j is the index of output of that op.
      """
      producer_map = {}
      for i in range(len(ssa)):
          outputs = ssa[i][1]
          for j, outp in enumerate(outputs):
              producer_map[outp] = (i, j)
      return producer_map
  
  
  def get_consumer_map(ssa):
      """
      Return dict from versioned blob to list of (i, j),
          where i is index of consumer op, j is the index of input of that op.
      """
      consumer_map = collections.defaultdict(list)
      for i in range(len(ssa)):
          inputs = ssa[i][0]
          for j, inp in enumerate(inputs):
              consumer_map[inp].append((i, j))
      return consumer_map
  
  
  def get_params_from_init_net(
      init_net: caffe2_pb2.NetDef,
  ) -> [Dict[str, Any], Dict[str, caffe2_pb2.DeviceOption]]:
      """
      Take the output blobs from init_net by running it.
      Outputs:
          params: dict from blob name to numpy array
          device_options: dict from blob name to the device option of its creating op
      """
      # NOTE: this assumes that the params is determined by producer op with the
      # only exception be CopyGPUToCPU which is CUDA op but returns CPU tensor.
      def _get_device_option(producer_op):
          if producer_op.type == "CopyGPUToCPU":
              return caffe2_pb2.DeviceOption()
          else:
              return producer_op.device_option
  
      with ScopedWS("__get_params_from_init_net__", is_reset=True, is_cleanup=True) as ws:
          ws.RunNetOnce(init_net)
          params = {b: fetch_any_blob(b) for b in init_net.external_output}
      ssa, versions = core.get_ssa(init_net)
      producer_map = get_producer_map(ssa)
      device_options = {
          b: _get_device_option(init_net.op[producer_map[(b, versions[b])][0]])
          for b in init_net.external_output
      }
      return params, device_options
  
  
  def _updater_raise(op, input_types, output_types):
      raise RuntimeError(
          "Failed to apply updater for op {} given input_types {} and"
          " output_types {}".format(op, input_types, output_types)
      )
  
  
  def _generic_status_identifier(
      predict_net: caffe2_pb2.NetDef,
      status_updater: Callable,
      known_status: Dict[Tuple[str, int], Any],
  ) -> Dict[Tuple[str, int], Any]:
      """
      Statically infer the status of each blob, the status can be such as device type
          (CPU/GPU), layout (NCHW/NHWC), data type (float32/int8), etc. "Blob" here
          is versioned blob (Tuple[str, int]) in the format compatible with ssa.
      Inputs:
          predict_net: the caffe2 network
          status_updater: a callable, given an op and the status of its input/output,
              it returns the updated status of input/output. `None` is used for
              representing unknown status.
          known_status: a dict containing known status, used as initialization.
      Outputs:
          A dict mapping from versioned blob to its status
      """
      ssa, versions = core.get_ssa(predict_net)
      versioned_ext_input = [(b, 0) for b in predict_net.external_input]
      versioned_ext_output = [(b, versions[b]) for b in predict_net.external_output]
      all_versioned_blobs = set().union(*[set(x[0] + x[1]) for x in ssa])
  
      allowed_vbs = all_versioned_blobs.union(versioned_ext_input).union(versioned_ext_output)
      assert all(k in allowed_vbs for k in known_status)
      assert all(v is not None for v in known_status.values())
      _known_status = copy.deepcopy(known_status)
  
      def _check_and_update(key, value):
          assert value is not None
          if key in _known_status:
              if not _known_status[key] == value:
                  raise RuntimeError(
                      "Confilict status for {}, existing status {}, new status {}".format(
                          key, _known_status[key], value
                      )
                  )
          _known_status[key] = value
  
      def _update_i(op, ssa_i):
          versioned_inputs = ssa_i[0]
          versioned_outputs = ssa_i[1]
  
          inputs_status = [_known_status.get(b, None) for b in versioned_inputs]
          outputs_status = [_known_status.get(b, None) for b in versioned_outputs]
  
          new_inputs_status, new_outputs_status = status_updater(op, inputs_status, outputs_status)
  
          for versioned_blob, status in zip(
              versioned_inputs + versioned_outputs, new_inputs_status + new_outputs_status
          ):
              if status is not None:
                  _check_and_update(versioned_blob, status)
  
      for op, ssa_i in zip(predict_net.op, ssa):
          _update_i(op, ssa_i)
      for op, ssa_i in zip(reversed(predict_net.op), reversed(ssa)):
          _update_i(op, ssa_i)
  
      # NOTE: This strictly checks all the blob from predict_net must be assgined
      # a known status. However sometimes it's impossible (eg. having deadend op),
      # we may relax this constraint if
      for k in all_versioned_blobs:
          if k not in _known_status:
              raise NotImplementedError(
                  "Can not infer the status for {}. Currently only support the case where"
                  " a single forward and backward pass can identify status for all blobs.".format(k)
              )
  
      return _known_status
  
  
  def infer_device_type(
      predict_net: caffe2_pb2.NetDef,
      known_status: Dict[Tuple[str, int], Any],
      device_name_style: str = "caffe2",
  ) -> Dict[Tuple[str, int], str]:
      """ Return the device type ("cpu" or "gpu"/"cuda") of each (versioned) blob """
  
      assert device_name_style in ["caffe2", "pytorch"]
      _CPU_STR = "cpu"
      _GPU_STR = "gpu" if device_name_style == "caffe2" else "cuda"
  
      def _copy_cpu_to_gpu_updater(op, input_types, output_types):
          if input_types[0] == _GPU_STR or output_types[0] == _CPU_STR:
              _updater_raise(op, input_types, output_types)
          return ([_CPU_STR], [_GPU_STR])
  
      def _copy_gpu_to_cpu_updater(op, input_types, output_types):
          if input_types[0] == _CPU_STR or output_types[0] == _GPU_STR:
              _updater_raise(op, input_types, output_types)
          return ([_GPU_STR], [_CPU_STR])
  
      def _other_ops_updater(op, input_types, output_types):
          non_none_types = [x for x in input_types + output_types if x is not None]
          if len(non_none_types) > 0:
              the_type = non_none_types[0]
              if not all(x == the_type for x in non_none_types):
                  _updater_raise(op, input_types, output_types)
          else:
              the_type = None
          return ([the_type for _ in op.input], [the_type for _ in op.output])
  
      def _device_updater(op, *args, **kwargs):
          return {
              "CopyCPUToGPU": _copy_cpu_to_gpu_updater,
              "CopyGPUToCPU": _copy_gpu_to_cpu_updater,
          }.get(op.type, _other_ops_updater)(op, *args, **kwargs)
  
      return _generic_status_identifier(predict_net, _device_updater, known_status)
  
  
  # ==== torch/utils_caffe2/vis.py ===============================================
  
  
  def _modify_blob_names(ops, blob_rename_f):
      ret = []
  
      def _replace_list(blob_list, replaced_list):
          del blob_list[:]
          blob_list.extend(replaced_list)
  
      for x in ops:
          cur = copy.deepcopy(x)
          _replace_list(cur.input, list(map(blob_rename_f, cur.input)))
          _replace_list(cur.output, list(map(blob_rename_f, cur.output)))
          ret.append(cur)
  
      return ret
  
  
  def _rename_blob(name, blob_sizes, blob_ranges):
      def _list_to_str(bsize):
          ret = ", ".join([str(x) for x in bsize])
          ret = "[" + ret + "]"
          return ret
  
      ret = name
      if blob_sizes is not None and name in blob_sizes:
          ret += "\n" + _list_to_str(blob_sizes[name])
      if blob_ranges is not None and name in blob_ranges:
          ret += "\n" + _list_to_str(blob_ranges[name])
  
      return ret
  
  
  # graph_name could not contain word 'graph'
  def save_graph(net, file_name, graph_name="net", op_only=True, blob_sizes=None, blob_ranges=None):
      blob_rename_f = functools.partial(_rename_blob, blob_sizes=blob_sizes, blob_ranges=blob_ranges)
      return save_graph_base(net, file_name, graph_name, op_only, blob_rename_f)
  
  
  def save_graph_base(net, file_name, graph_name="net", op_only=True, blob_rename_func=None):
      graph = None
      ops = net.op
      if blob_rename_func is not None:
          ops = _modify_blob_names(ops, blob_rename_func)
      if not op_only:
          graph = net_drawer.GetPydotGraph(ops, graph_name, rankdir="TB")
      else:
          graph = net_drawer.GetPydotGraphMinimal(
              ops, graph_name, rankdir="TB", minimal_dependency=True
          )
  
      try:
          par_dir = os.path.dirname(file_name)
          if not os.path.exists(par_dir):
              os.makedirs(par_dir)
  
          format = os.path.splitext(os.path.basename(file_name))[-1]
          if format == ".png":
              graph.write_png(file_name)
          elif format == ".pdf":
              graph.write_pdf(file_name)
          elif format == ".svg":
              graph.write_svg(file_name)
          else:
              print("Incorrect format {}".format(format))
      except Exception as e:
          print("Error when writing graph to image {}".format(e))
  
      return graph
  
  
  # ==== torch/utils_toffee/aten_to_caffe2.py ====================================
  
  
  def group_norm_replace_aten_with_caffe2(predict_net: caffe2_pb2.NetDef):
      """
      For ONNX exported model, GroupNorm will be represented as ATen op,
          this can be a drop in replacement from ATen to GroupNorm
      """
      count = 0
      for op in predict_net.op:
          if op.type == "ATen":
              op_name = get_pb_arg_vals(op, "operator", None)  # return byte in py3
              if op_name and op_name.decode() == "group_norm":
                  op.arg.remove(get_pb_arg(op, "operator"))
  
                  if get_pb_arg_vali(op, "cudnn_enabled", None):
                      op.arg.remove(get_pb_arg(op, "cudnn_enabled"))
  
                  num_groups = get_pb_arg_vali(op, "num_groups", None)
                  if num_groups is not None:
                      op.arg.remove(get_pb_arg(op, "num_groups"))
                      check_set_pb_arg(op, "group", "i", num_groups)
  
                  op.type = "GroupNorm"
                  count += 1
      if count > 1:
          logger.info("Replaced {} ATen operator to GroupNormOp".format(count))
  
  
  # ==== torch/utils_toffee/alias.py =============================================
  
  
  def alias(x, name, is_backward=False):
      if not torch.onnx.is_in_onnx_export():
          return x
      assert isinstance(x, torch.Tensor)
      x = x.cpu()
      return torch.ops._caffe2.AliasWithName(x, name, is_backward=is_backward)
  
  
  def fuse_alias_placeholder(predict_net, init_net):
      """ Remove AliasWithName placeholder and rename the input/output of it """
      # First we finish all the re-naming
      for i, op in enumerate(predict_net.op):
          if op.type == "AliasWithName":
              assert len(op.input) == 1
              assert len(op.output) == 1
              name = get_pb_arg_vals(op, "name", None).decode()
              is_backward = bool(get_pb_arg_vali(op, "is_backward", 0))
              rename_op_input(predict_net, init_net, i, 0, name, from_producer=is_backward)
              rename_op_output(predict_net, i, 0, name)
  
      # Remove AliasWithName, should be very safe since it's a non-op
      new_ops = []
      for op in predict_net.op:
          if op.type != "AliasWithName":
              new_ops.append(op)
          else:
              # safety check
              assert op.input == op.output
              assert op.input[0] == op.arg[0].s.decode()
      del predict_net.op[:]
      predict_net.op.extend(new_ops)
  
  
  # ==== torch/utils_caffe2/graph_transform.py ===================================
  
  
  class IllegalGraphTransformError(ValueError):
      """ When a graph transform function call can't be executed. """
  
  
  def _rename_versioned_blob_in_proto(
      proto: caffe2_pb2.NetDef,
      old_name: str,
      new_name: str,
      version: int,
      ssa: List[Tuple[List[Tuple[str, int]], List[Tuple[str, int]]]],
      start_versions: Dict[str, int],
      end_versions: Dict[str, int],
  ):
      """ In given proto, rename all blobs with matched version """
      # Operater list
      for op, i_th_ssa in zip(proto.op, ssa):
          versioned_inputs, versioned_outputs = i_th_ssa
          for i in range(len(op.input)):
              if versioned_inputs[i] == (old_name, version):
                  op.input[i] = new_name
          for i in range(len(op.output)):
              if versioned_outputs[i] == (old_name, version):
                  op.output[i] = new_name
      # external_input
      if start_versions.get(old_name, 0) == version:
          for i in range(len(proto.external_input)):
              if proto.external_input[i] == old_name:
                  proto.external_input[i] = new_name
      # external_output
      if end_versions.get(old_name, 0) == version:
          for i in range(len(proto.external_output)):
              if proto.external_output[i] == old_name:
                  proto.external_output[i] = new_name
  
  
  def rename_op_input(
      predict_net: caffe2_pb2.NetDef,
      init_net: caffe2_pb2.NetDef,
      op_id: int,
      input_id: int,
      new_name: str,
      from_producer: bool = False,
  ):
      """
      Rename the op_id-th operator in predict_net, change it's input_id-th input's
          name to the new_name. It also does automatic re-route and change
          external_input and init_net if necessary.
      - It requires the input is only consumed by this op.
      - This function modifies predict_net and init_net in-place.
      - When from_producer is enable, this also updates other operators that consumes
          the same input. Be cautious because may trigger unintended behavior.
      """
      assert isinstance(predict_net, caffe2_pb2.NetDef)
      assert isinstance(init_net, caffe2_pb2.NetDef)
  
      init_net_ssa, init_net_versions = core.get_ssa(init_net)
      predict_net_ssa, predict_net_versions = core.get_ssa(
          predict_net, copy.deepcopy(init_net_versions)
      )
  
      versioned_inputs, versioned_outputs = predict_net_ssa[op_id]
      old_name, version = versioned_inputs[input_id]
  
      if from_producer:
          producer_map = get_producer_map(predict_net_ssa)
          if not (old_name, version) in producer_map:
              raise NotImplementedError(
                  "Can't find producer, the input {} is probably from"
                  " init_net, this is not supported yet.".format(old_name)
              )
          producer = producer_map[(old_name, version)]
          rename_op_output(predict_net, producer[0], producer[1], new_name)
          return
  
      def contain_targets(op_ssa):
          return (old_name, version) in op_ssa[0]
  
      is_consumer = [contain_targets(op_ssa) for op_ssa in predict_net_ssa]
      if sum(is_consumer) > 1:
          raise IllegalGraphTransformError(
              (
                  "Input '{}' of operator(#{}) are consumed by other ops, please use"
                  + " rename_op_output on the producer instead. Offending op: \n{}"
              ).format(old_name, op_id, predict_net.op[op_id])
          )
  
      # update init_net
      _rename_versioned_blob_in_proto(
          init_net, old_name, new_name, version, init_net_ssa, {}, init_net_versions
      )
      # update predict_net
      _rename_versioned_blob_in_proto(
          predict_net,
          old_name,
          new_name,
          version,
          predict_net_ssa,
          init_net_versions,
          predict_net_versions,
      )
  
  
  def rename_op_output(predict_net: caffe2_pb2.NetDef, op_id: int, output_id: int, new_name: str):
      """
      Rename the op_id-th operator in predict_net, change it's output_id-th input's
          name to the new_name. It also does automatic re-route and change
          external_output and if necessary.
      - It allows multiple consumers of its output.
      - This function modifies predict_net in-place, doesn't need init_net.
      """
      assert isinstance(predict_net, caffe2_pb2.NetDef)
  
      ssa, blob_versions = core.get_ssa(predict_net)
  
      versioned_inputs, versioned_outputs = ssa[op_id]
      old_name, version = versioned_outputs[output_id]
  
      # update predict_net
      _rename_versioned_blob_in_proto(
          predict_net, old_name, new_name, version, ssa, {}, blob_versions
      )
  
  
  def get_sub_graph_external_input_output(
      predict_net: caffe2_pb2.NetDef, sub_graph_op_indices: List[int]
  ) -> Tuple[List[Tuple[str, int]], List[Tuple[str, int]]]:
      """
      Return the list of external input/output of sub-graph,
      each element is tuple of the name and corresponding version in predict_net.
  
      external input/output is defined the same way as caffe2 NetDef.
      """
      ssa, versions = core.get_ssa(predict_net)
  
      all_inputs = []
      all_outputs = []
      for op_id in sub_graph_op_indices:
          all_inputs += [inp for inp in ssa[op_id][0] if inp not in all_inputs]
          all_outputs += list(ssa[op_id][1])  # ssa output won't repeat
  
      # for versioned blobs, external inputs are just those blob in all_inputs
      # but not in all_outputs
      ext_inputs = [inp for inp in all_inputs if inp not in all_outputs]
  
      # external outputs are essentially outputs of this subgraph that are used
      # outside of this sub-graph (including predict_net.external_output)
      all_other_inputs = sum(
          (ssa[i][0] for i in range(len(ssa)) if i not in sub_graph_op_indices),
          [(outp, versions[outp]) for outp in predict_net.external_output],
      )
      ext_outputs = [outp for outp in all_outputs if outp in set(all_other_inputs)]
  
      return ext_inputs, ext_outputs
  
  
  class DiGraph:
      """ A DAG representation of caffe2 graph, each vertice is a versioned blob. """
  
      def __init__(self):
          self.vertices = set()
          self.graph = collections.defaultdict(list)
  
      def add_edge(self, u, v):
          self.graph[u].append(v)
          self.vertices.add(u)
          self.vertices.add(v)
  
      # grab from https://www.geeksforgeeks.org/find-paths-given-source-destination/
      def get_all_paths(self, s, d):
          visited = {k: False for k in self.vertices}
          path = []
          all_paths = []
  
          def _get_all_paths_util(graph, u, d, visited, path):
              visited[u] = True
              path.append(u)
              if u == d:
                  all_paths.append(copy.deepcopy(path))
              else:
                  for i in graph[u]:
                      if not visited[i]:
                          _get_all_paths_util(graph, i, d, visited, path)
              path.pop()
              visited[u] = False
  
          _get_all_paths_util(self.graph, s, d, visited, path)
          return all_paths
  
      @staticmethod
      def from_ssa(ssa):
          graph = DiGraph()
          for op_id in range(len(ssa)):
              for inp in ssa[op_id][0]:
                  for outp in ssa[op_id][1]:
                      graph.add_edge(inp, outp)
          return graph
  
  
  def _get_dependency_chain(ssa, versioned_target, versioned_source):
      """
      Return the index list of relevant operator to produce target blob from source blob,
          if there's no dependency, return empty list.
      """
  
      # finding all paths between nodes can be O(N!), thus we can only search
      # in the subgraph using the op starting from the first consumer of source blob
      # to the producer of the target blob.
      consumer_map = get_consumer_map(ssa)
      producer_map = get_producer_map(ssa)
      start_op = min(x[0] for x in consumer_map[versioned_source]) - 15
      end_op = (
          producer_map[versioned_target][0] + 15 if versioned_target in producer_map else start_op
      )
      sub_graph_ssa = ssa[start_op : end_op + 1]
      if len(sub_graph_ssa) > 30:
          logger.warning(
              "Subgraph bebetween {} and {} is large (from op#{} to op#{}), it"
              " might take non-trival time to find all paths between them.".format(
                  versioned_source, versioned_target, start_op, end_op
              )
          )
  
      dag = DiGraph.from_ssa(sub_graph_ssa)
      paths = dag.get_all_paths(versioned_source, versioned_target)  # include two ends
      ops_in_paths = [[producer_map[blob][0] for blob in path[1:]] for path in paths]
      return sorted(set().union(*[set(ops) for ops in ops_in_paths]))
  
  
  def identify_reshape_sub_graph(predict_net: caffe2_pb2.NetDef) -> List[List[int]]:
      """
      Idenfity the reshape sub-graph in a protobuf.
      The reshape sub-graph is defined as matching the following pattern:
  
      (input_blob) -> Op_1 -> ... -> Op_N -> (new_shape) -─┐
          └-------------------------------------------> Reshape -> (output_blob)
  
      Return:
          List of sub-graphs, each sub-graph is represented as a list of indices
          of the relavent ops, [Op_1, Op_2, ..., Op_N, Reshape]
      """
  
      ssa, _ = core.get_ssa(predict_net)
  
      ret = []
      for i, op in enumerate(predict_net.op):
          if op.type == "Reshape":
              assert len(op.input) == 2
              input_ssa = ssa[i][0]
              data_source = input_ssa[0]
              shape_source = input_ssa[1]
              op_indices = _get_dependency_chain(ssa, shape_source, data_source)
              ret.append(op_indices + [i])
      return ret
  
  
  def remove_reshape_for_fc(predict_net, params):
      """
      In PyTorch nn.Linear has to take 2D tensor, this often leads to reshape
          a 4D tensor to 2D by calling .view(). However this (dynamic) reshaping
          doesn't work well with ONNX and Int8 tools, and cause using extra
          ops (eg. ExpandDims) that might not be available on mobile.
      Luckily Caffe2 supports 4D tensor for FC, so we can remove those reshape
          after exporting ONNX model.
      """
      from caffe2.python import core
  
      # find all reshape sub-graph that can be removed, which is now all Reshape
      # sub-graph whose output is only consumed by FC.
      # TODO: to make it safer, we may need the actually value to better determine
      # if a Reshape before FC is removable.
      reshape_sub_graphs = identify_reshape_sub_graph(predict_net)
      sub_graphs_to_remove = []
      for reshape_sub_graph in reshape_sub_graphs:
          reshape_op_id = reshape_sub_graph[-1]
          assert predict_net.op[reshape_op_id].type == "Reshape"
          ssa, _ = core.get_ssa(predict_net)
          reshape_output = ssa[reshape_op_id][1][0]
          consumers = [i for i in range(len(ssa)) if reshape_output in ssa[i][0]]
          if all(predict_net.op[consumer].type == "FC" for consumer in consumers):
              # safety check if the sub-graph is isolated, for this reshape sub-graph,
              # it means it has one non-param external input and one external output.
              ext_inputs, ext_outputs = get_sub_graph_external_input_output(
                  predict_net, reshape_sub_graph
              )
              non_params_ext_inputs = [inp for inp in ext_inputs if inp[1] != 0]
              if len(non_params_ext_inputs) == 1 and len(ext_outputs) == 1:
                  sub_graphs_to_remove.append(reshape_sub_graph)
  
      # perform removing subgraph by:
      # 1: rename the Reshape's output to its input, then the graph can be
      #   seen as in-place itentify, meaning whose external input/output are the same.
      # 2: simply remove those ops.
      remove_op_ids = []
      params_to_remove = []
      for sub_graph in sub_graphs_to_remove:
          logger.info(
              "Remove Reshape sub-graph:\n{}".format(
                  "".join(["(#{:>4})\n{}".format(i, predict_net.op[i]) for i in sub_graph])
              )
          )
          reshape_op_id = sub_graph[-1]
          new_reshap_output = predict_net.op[reshape_op_id].input[0]
          rename_op_output(predict_net, reshape_op_id, 0, new_reshap_output)
          ext_inputs, ext_outputs = get_sub_graph_external_input_output(predict_net, sub_graph)
          non_params_ext_inputs = [inp for inp in ext_inputs if inp[1] != 0]
          params_ext_inputs = [inp for inp in ext_inputs if inp[1] == 0]
          assert len(non_params_ext_inputs) == 1 and len(ext_outputs) == 1
          assert ext_outputs[0][0] == non_params_ext_inputs[0][0]
          assert ext_outputs[0][1] == non_params_ext_inputs[0][1] + 1
          remove_op_ids.extend(sub_graph)
          params_to_remove.extend(params_ext_inputs)
  
      predict_net = copy.deepcopy(predict_net)
      new_ops = [op for i, op in enumerate(predict_net.op) if i not in remove_op_ids]
      del predict_net.op[:]
      predict_net.op.extend(new_ops)
      for versioned_params in params_to_remove:
          name = versioned_params[0]
          logger.info("Remove params: {} from init_net and predict_net.external_input".format(name))
          del params[name]
          predict_net.external_input.remove(name)
  
      return predict_net, params
  
  
  def fuse_copy_between_cpu_and_gpu(predict_net: caffe2_pb2.NetDef):
      """
      In-place fuse extra copy ops between cpu/gpu for the following case:
          a -CopyAToB-> b -CopyBToA> c1 -NextOp1-> d1
                          -CopyBToA> c2 -NextOp2-> d2
      The fused network will look like:
          a -NextOp1-> d1
            -NextOp2-> d2
      """
  
      _COPY_OPS = ["CopyCPUToGPU", "CopyGPUToCPU"]
  
      def _fuse_once(predict_net):
          ssa, blob_versions = core.get_ssa(predict_net)
          consumer_map = get_consumer_map(ssa)
          versioned_external_output = [
              (name, blob_versions[name]) for name in predict_net.external_output
          ]
  
          for op_id, op in enumerate(predict_net.op):
              if op.type in _COPY_OPS:
                  fw_copy_versioned_output = ssa[op_id][1][0]
                  consumer_ids = [x[0] for x in consumer_map[fw_copy_versioned_output]]
                  reverse_op_type = _COPY_OPS[1 - _COPY_OPS.index(op.type)]
  
                  is_fusable = (
                      len(consumer_ids) > 0
                      and fw_copy_versioned_output not in versioned_external_output
                      and all(
                          predict_net.op[_op_id].type == reverse_op_type
                          and ssa[_op_id][1][0] not in versioned_external_output
                          for _op_id in consumer_ids
                      )
                  )
  
                  if is_fusable:
                      for rv_copy_op_id in consumer_ids:
                          # making each NextOp uses "a" directly and removing Copy ops
                          rs_copy_versioned_output = ssa[rv_copy_op_id][1][0]
                          next_op_id, inp_id = consumer_map[rs_copy_versioned_output][0]
                          predict_net.op[next_op_id].input[inp_id] = op.input[0]
                      # remove CopyOps
                      new_ops = [
                          op
                          for i, op in enumerate(predict_net.op)
                          if i != op_id and i not in consumer_ids
                      ]
                      del predict_net.op[:]
                      predict_net.op.extend(new_ops)
                      return True
  
          return False
  
      # _fuse_once returns False is nothing can be fused
      while _fuse_once(predict_net):
          pass
  
  
  def remove_dead_end_ops(net_def: caffe2_pb2.NetDef):
      """ remove ops if its output is not used or not in external_output """
      ssa, versions = core.get_ssa(net_def)
      versioned_external_output = [(name, versions[name]) for name in net_def.external_output]
      consumer_map = get_consumer_map(ssa)
      removed_op_ids = set()
  
      def _is_dead_end(versioned_blob):
          return not (
              versioned_blob in versioned_external_output
              or (
                  len(consumer_map[versioned_blob]) > 0
                  and all(x[0] not in removed_op_ids for x in consumer_map[versioned_blob])
              )
          )
  
      for i, ssa_i in reversed(list(enumerate(ssa))):
          versioned_outputs = ssa_i[1]
          if all(_is_dead_end(outp) for outp in versioned_outputs):
              removed_op_ids.add(i)
  
      # simply removing those deadend ops should have no effect to external_output
      new_ops = [op for i, op in enumerate(net_def.op) if i not in removed_op_ids]
      del net_def.op[:]
      net_def.op.extend(new_ops)
  