text: # SPDX-License-Identifier: Apache-2.0
  
  
  """symbolic executor
     Computes a part of the graph symbolically using SymbolicTensorElements
  """
  
  import numpy as np
  from tf2onnx import utils
  
  
  # pylint: disable=logging-not-lazy,unused-argument,missing-docstring
  
  
  class SymbolicTensorElement:
      """
      Helps with symbolic execution of the graph, in particular tensors representing shapes. Supports multiplication
      and tensor ops.
      """
      def __init__(self, terms, constant):
          # Terms is a list representing variables
          self.terms = terms
          self.constant = constant
          if self.constant == 0:
              self.terms = []
  
      def __mul__(self, other):
          if isinstance(other, SymbolicTensorElement):
              # Concat terms, multiply constant
              return SymbolicTensorElement(self.terms + other.terms, self.constant * other.constant)
          # Other term is a constant
          return SymbolicTensorElement(self.terms, self.constant * other)
  
      def __rmul__(self, other):
          return self.__mul__(other)
  
      def is_const(self):
          return len(self.terms) == 0
  
      def is_one(self):
          return len(self.terms) == 0 and self.constant == 1
  
      def is_single_var(self):
          return len(self.terms) == 1 and self.constant == 1
  
      def has_multiple_terms(self):
          return not self.is_const() and not self.is_single_var()
  
      def get_reshape_dim(self, i, offset):
          if self.is_const():
              return self.constant
          if self.get_offset(i) == offset:
              return 0
          return -1
  
      @staticmethod
      def from_const(constant):
          return SymbolicTensorElement([], constant)
  
      @staticmethod
      def from_variable(variable):
          return SymbolicTensorElement([variable], 1)
  
      @staticmethod
      def from_value(value):
          if isinstance(value, SymbolicTensorElement):
              return value
          return SymbolicTensorElement.from_const(value)
  
      @staticmethod
      def np_array(np_array):
          return np.vectorize(SymbolicTensorElement.from_value)(np_array)
  
  class SymbolicExecutionException(Exception):
      pass
  
  class SymbolicExecutor:
      def __init__(self, graph):
          self.graph = graph
          self.op_map = {
              "Unsqueeze": self.compute_squeeze_unsqueeze,
              "Squeeze": self.compute_squeeze_unsqueeze,
              "Gather": self.compute_gather,
              "Mul": self.compute_mul,
              "ReduceProd": self.compute_reduceprod,
              "Slice": self.compute_slice,
              "Cast": self.compute_cast,
              "Concat": self.compute_concat,
              "Const": self.compute_const
          }
  
      def compute_outputs(self, outputs, feed_dict):
          """Given a map of inputs to np arrays, outputs a list of np arrays of SymbolicTensorElements"""
          nodes_to_compute = self.plan_computation(outputs, feed_dict)
          if nodes_to_compute is None:
              return None
          results = feed_dict.copy()
          for node in nodes_to_compute:
              try:
                  results.update(self.compute_node(node, results))
              except Exception as e:
                  raise SymbolicExecutionException(str(e))
          # Intermediate results might be non-symbolic numpy arrays
          return [SymbolicTensorElement.np_array(results[out]) for out in outputs]
  
      def plan_computation(self, outputs, feed_dict):
          nodes = list(set(self.graph.get_node_by_output(out) for out in outputs))
          sorted_nodes = []
          while nodes:
              n = nodes.pop()
              if n.type not in self.op_map:
                  raise SymbolicExecutionException("Unsupported op %s" % n.type)
              sorted_nodes.append(n)
              for inp, inp_name in zip(n.inputs, n.input):
                  if inp_name != '' and inp_name not in feed_dict:
                      nodes.append(inp)
          return sorted_nodes[::-1]
  
      def compute_node(self, node, feed_dict):
          results = self.op_map[node.type](node, feed_dict)
          return {out: np.array(res) for out, res in zip(node.output, results)}
  
      def compute_const(self, node, feed_dict):
          return [node.get_tensor_value(as_list=False)]
  
      def compute_squeeze_unsqueeze(self, node, feed_dict):
          inp1 = feed_dict[node.input[0]]
          if self.graph.opset < 13:
              axes = node.get_attr_value("axes")
          else:
              axes = feed_dict[node.input[1]].tolist()
          shape = inp1.shape
          handler = self.compute_unsqueeze_shape if node.type == "Unsqueeze" else self.compute_squeeze_shape
          new_shape = handler(shape, axes)
          return [inp1.reshape(new_shape)]
  
      def compute_cast(self, node, feed_dict):
          inp = feed_dict[node.input[0]]
          if inp.dtype == np.object:
              return [inp]
          np_dtype = utils.ONNX_TO_NUMPY_DTYPE[node.get_attr("to").i]
          return [inp.astype(np_dtype)]
  
      def compute_mul(self, node, feed_dict):
          return [feed_dict[node.input[0]] * feed_dict[node.input[1]]]
  
      def compute_reduceprod(self, node, feed_dict):
          inp = feed_dict[node.input[0]]
          axes = node.get_attr_value("axes")
          keepdims = node.get_attr_value("keepdims", 1)
          return [np.prod(inp, axis=tuple(axes), keepdims=keepdims)]
  
      def compute_slice(self, node, feed_dict):
          inps = [feed_dict[inp] if inp != '' else None for inp in node.input]
          if self.graph.opset >= 10:
              while len(inps) < 5:
                  inps.append(None)
              data, starts, ends, axes, steps = inps
          else:
              data = inps[0]
              starts = node.get_attr_value("starts")
              ends = node.get_attr_value("ends")
              axes = node.get_attr_value("axes")
              steps = None
          rank = len(data.shape)
          ndims = len(starts)
          if axes is None:
              axes = list(range(ndims))
          if steps is None:
              steps = [1] * ndims
          slices = [slice(None, None, None) for _ in range(rank)]
          for axis, start, end, step in zip(axes, starts, ends, steps):
              slices[axis] = slice(start, end, step)
          return [data[tuple(slices)]]
  
      def compute_concat(self, node, feed_dict):
          axis = node.get_attr_value("axis")
          inps = [feed_dict[inp] for inp in node.input]
          return [np.concatenate(inps, axis=axis)]
  
      def compute_gather(self, node, feed_dict):
          data = feed_dict[node.input[0]]
          indices = feed_dict[node.input[1]]
          if indices.dtype == np.object:
              raise SymbolicExecutionException("Gather requires non-symbolic indices")
          axis = node.get_attr_value("axis", 0)
          return [np.take(data, indices, axis=axis)]
  
      def compute_unsqueeze_shape(self, shape_in, axes):
          dims_out = len(shape_in) + len(axes)
          axes = [i if i >= 0 else i + dims_out for i in axes]
          shape_in = iter(shape_in)
          shape_out = [None] * dims_out
          for ind in axes:
              shape_out[ind] = 1
          for ind, val in enumerate(shape_out):
              if val is None:
                  shape_out[ind] = next(shape_in)
          return shape_out
  
      def compute_squeeze_shape(self, shape_in, axes):
          axes = [i if i >= 0 else i + len(axes) for i in axes]
          shape_out = []
          for ind, val in enumerate(shape_in):
              if ind not in axes:
                  shape_out.append(val)
          return shape_out
  