text: # Copyright (C) 2018-2021 Intel Corporation
  # SPDX-License-Identifier: Apache-2.0
  
  """Factory functions for all openvino ops."""
  from typing import Callable, Iterable, List, Optional, Set, Union
  
  import numpy as np
  from functools import partial
  
  from openvino.runtime import Node, Shape
  from openvino.runtime.op import Constant, Parameter
  from openvino.runtime.opset_utils import _get_node_factory
  from openvino.runtime.utils.decorators import binary_op, nameable_op, unary_op
  from openvino.runtime.utils.input_validation import (
      assert_list_of_ints,
      check_valid_attributes,
      is_non_negative_value,
      is_positive_value,
  )
  from openvino.runtime.utils.node_factory import NodeFactory
  from openvino.runtime.utils.types import (
      NodeInput,
      NumericData,
      NumericType,
      ScalarData,
      TensorShape,
      as_node,
      as_nodes,
      get_dtype,
      get_element_type,
      get_element_type_str,
      make_constant_node,
  )
  
  _get_node_factory_opset3 = partial(_get_node_factory, "opset3")
  
  # -------------------------------------------- ops ------------------------------------------------
  
  
  @nameable_op
  def assign(new_value: NodeInput, variable_id: str, name: Optional[str] = None) -> Node:
      """Return a node which produces the Assign operation.
  
      @param new_value:    Node producing a value to be assigned to a variable.
      @param variable_id:  Id of a variable to be updated.
      @param name:         Optional name for output node.
      @return Assign node
      """
      return _get_node_factory_opset3().create(
          "Assign",
          [as_node(new_value)],
          {"variable_id": variable_id}
      )
  
  
  @nameable_op
  def broadcast(
      data: NodeInput,
      target_shape: NodeInput,
      axes_mapping: Optional[NodeInput] = None,
      broadcast_spec: str = "NUMPY",
      name: Optional[str] = None,
  ) -> Node:
      """Create a node which broadcasts the input node's values along specified axes to a desired shape.
  
      @param data: The node with input tensor data.
      @param target_shape: The node with a new shape we want to broadcast tensor to.
      @param axes_mapping: The node with a axis positions (0-based) in the result
                             that are being broadcast.
      @param broadcast_spec: The type of broadcasting that specifies mapping of input tensor axes
                             to output shape axes. Range of values: NUMPY, EXPLICIT, BIDIRECTIONAL.
      @param name: Optional new name for output node.
      @return New node with broadcast shape.
      """
      inputs = as_nodes(data, target_shape)
      if broadcast_spec.upper() == "EXPLICIT":
          inputs.append(as_node(axes_mapping))
      return _get_node_factory_opset3().create(
          "Broadcast", inputs, {"mode": broadcast_spec.upper()}
      )
  
  
  @nameable_op
  def bucketize(
      data: Node,
      buckets: NodeInput,
      output_type: str = "i64",
      with_right_bound: bool = True,
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which produces the Bucketize operation.
  
      @param data:              Input data to bucketize
      @param buckets:           1-D of sorted unique boundaries for buckets
      @param output_type:       Output tensor type, "i64" or "i32", defaults to i64
      @param with_right_bound:  indicates whether bucket includes the right or left
                                edge of interval. default true = includes right edge
      @param name:              Optional name for output node.
      @return Bucketize node
      """
      return _get_node_factory_opset3().create(
          "Bucketize",
          [data, as_node(buckets)],
          {"output_type": output_type, "with_right_bound": with_right_bound},
      )
  
  
  @nameable_op
  def cum_sum(
      arg: NodeInput,
      axis: NodeInput,
      exclusive: bool = False,
      reverse: bool = False,
      name: Optional[str] = None,
  ) -> Node:
      """Construct a cumulative summation operation.
  
      @param arg: The tensor to be summed.
      @param axis: zero dimension tensor specifying axis position along which sum will be performed.
      @param exclusive: if set to true, the top element is not included
      @param reverse: if set to true, will perform the sums in reverse direction
      @return New node performing the operation
      """
      return _get_node_factory_opset3().create(
          "CumSum", as_nodes(arg, axis), {"exclusive": exclusive, "reverse": reverse}
      )
  
  
  @nameable_op
  def embedding_bag_offsets_sum(
      emb_table: Node,
      indices: NodeInput,
      offsets: NodeInput,
      default_index: Optional[NodeInput] = None,
      per_sample_weights: Optional[NodeInput] = None,
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which performs sums of bags of embeddings without the intermediate embeddings.
  
      @param emb_table: Tensor containing the embedding lookup table.
      @param indices: Tensor with indices.
      @param offsets: Tensor containing the starting index positions of each bag in indices.
      @param per_sample_weights: Tensor with weights for each sample.
      @param default_index: Scalar containing default index in embedding table to fill empty bags.
      @param name: Optional name for output node.
      @return The new node which performs EmbeddingBagOffsetsSum
      """
      inputs = [emb_table, as_node(indices), as_node(offsets)]
      if per_sample_weights is not None:
          inputs.append(default_index)
          inputs.append(per_sample_weights)
      elif default_index is not None:
          inputs.append(default_index)
  
      return _get_node_factory_opset3().create("EmbeddingBagOffsetsSum", inputs, {})
  
  
  @nameable_op
  def embedding_bag_packed_sum(
      emb_table: NodeInput,
      indices: NodeInput,
      per_sample_weights: Optional[NodeInput] = None,
      name: Optional[str] = None,
  ) -> Node:
      """Return an EmbeddingBagPackedSum node.
  
      EmbeddingSegmentsSum constructs an output tensor by replacing every index in a given
      input tensor with a row (from the weights matrix) at that index
  
      @param emb_table: Tensor containing the embedding lookup table.
      @param indices: Tensor with indices.
      @param per_sample_weights: Weights to be multiplied with embedding table.
      @param name: Optional name for output node.
      @return EmbeddingBagPackedSum node
      """
      inputs = [as_node(emb_table), as_node(indices)]
      if per_sample_weights is not None:
          inputs.append(as_node(per_sample_weights))
  
      return _get_node_factory_opset3().create("EmbeddingBagPackedSum", inputs, {})
  
  
  @nameable_op
  def embedding_segments_sum(
      emb_table: Node,
      indices: NodeInput,
      segment_ids: NodeInput,
      num_segments: Optional[NodeInput] = None,
      default_index: Optional[NodeInput] = None,
      per_sample_weights: Optional[NodeInput] = None,
      name: Optional[str] = None,
  ) -> Node:
      """Return an EmbeddingSegmentsSum node.
  
      EmbeddingSegmentsSum constructs an output tensor by replacing every index in a given
      input tensor with a row (from the weights matrix) at that index
  
      @param emb_table: Tensor containing the embedding lookup table.
      @param indices: Tensor with indices.
      @param segment_ids: Tensor with indices into the output Tensor
      @param num_segments: Tensor with number of segments.
      @param default_index: Scalar containing default index in embedding table to fill empty bags.
      @param per_sample_weights: Weights to be multiplied with embedding table.
      @param name: Optional name for output node.
      @return EmbeddingSegmentsSum node
      """
      inputs = [as_node(emb_table), as_node(indices), as_node(segment_ids)]
      if per_sample_weights is not None:
          inputs.append(as_node(num_segments))
          inputs.append(as_node(default_index))
          inputs.append(as_node(per_sample_weights))
      elif default_index is not None:
          inputs.append(as_node(num_segments))
          inputs.append(as_node(default_index))
      elif num_segments is not None:
          inputs.append(as_node(num_segments))
  
      return _get_node_factory_opset3().create("EmbeddingSegmentsSum", inputs, {})
  
  
  @nameable_op
  def extract_image_patches(
      image: NodeInput,
      sizes: TensorShape,
      strides: List[int],
      rates: TensorShape,
      auto_pad: str,
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which produces the ExtractImagePatches operation.
  
      @param image:     4-D Input data to extract image patches.
      @param sizes:     Patch size in the format of [size_rows, size_cols].
      @param strides:   Patch movement stride in the format of [stride_rows, stride_cols]
      @param rates:     Element seleciton rate for creating a patch.
      @param auto_pad:  Padding type.
      @param name:      Optional name for output node.
      @return ExtractImagePatches node
      """
      return _get_node_factory_opset3().create(
          "ExtractImagePatches",
          [as_node(image)],
          {"sizes": sizes, "strides": strides, "rates": rates, "auto_pad": auto_pad},
      )
  
  
  @nameable_op
  def gru_cell(
      X: NodeInput,
      initial_hidden_state: NodeInput,
      W: NodeInput,
      R: NodeInput,
      B: NodeInput,
      hidden_size: int,
      activations: List[str] = None,
      activations_alpha: List[float] = None,
      activations_beta: List[float] = None,
      clip: float = 0.0,
      linear_before_reset: bool = False,
      name: Optional[str] = None,
  ) -> Node:
      """Perform GRUCell operation on the tensor from input node.
  
      GRUCell represents a single GRU Cell that computes the output
      using the formula described in the paper: https://arxiv.org/abs/1406.1078
  
      Note this class represents only single *cell* and not whole *layer*.
  
      @param X:                       The input tensor with shape: [batch_size, input_size].
      @param initial_hidden_state:    The hidden state tensor at current time step with shape:
                                      [batch_size, hidden_size].
      @param W:                       The weights for matrix multiplication, gate order: zrh.
                                      Shape: [3*hidden_size, input_size].
      @param R:                       The recurrence weights for matrix multiplication.
                                      Shape: [3*hidden_size, hidden_size].
      @param B:                       The sum of biases (weight and recurrence).
                                      For linear_before_reset set True the shape is [4*hidden_size].
                                      Otherwise the shape is [3*hidden_size].
      @param hidden_size:             The number of hidden units for recurrent cell.
                                      Specifies hidden state size.
      @param activations:             The vector of activation functions used inside recurrent cell.
      @param activation_alpha:        The vector of alpha parameters for activation functions in
                                      order respective to activation list.
      @param activation_beta:         The vector of beta parameters for activation functions in order
                                      respective to activation list.
      @param clip:                    The value defining clipping range [-clip, clip] on input of
                                      activation functions.
      @param linear_before_reset:     Flag denotes if the layer behaves according to the modification
                                      of GRUCell described in the formula in the ONNX documentation.
      @param name:                    Optional output node name.
      @return   The new node performing a GRUCell operation on tensor from input node.
      """
      if activations is None:
          activations = ["sigmoid", "tanh"]
      if activations_alpha is None:
          activations_alpha = []
      if activations_beta is None:
          activations_beta = []
  
      input_nodes = as_nodes(X, initial_hidden_state, W, R, B)
      attributes = {
          "hidden_size": hidden_size,
          "activations": activations,
          "activations_alpha": activations_alpha,
          "activations_beta": activations_beta,
          "linear_before_reset": linear_before_reset,
          "clip": clip,
      }
      return _get_node_factory_opset3().create("GRUCell", input_nodes, attributes)
  
  
  @nameable_op
  def non_max_suppression(
      boxes: NodeInput,
      scores: NodeInput,
      max_output_boxes_per_class: Optional[NodeInput] = None,
      iou_threshold: Optional[NodeInput] = None,
      score_threshold: Optional[NodeInput] = None,
      box_encoding: str = "corner",
      sort_result_descending: bool = True,
      output_type: str = "i64",
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which performs NonMaxSuppression.
  
      @param boxes: Tensor with box coordinates.
      @param scores: Tensor with box scores.
      @param max_output_boxes_per_class: Tensor Specifying maximum number of boxes
                                          to be selected per class.
      @param iou_threshold: Tensor specifying intersection over union threshold
      @param score_threshold: Tensor specifying minimum score to consider box for the processing.
      @param box_encoding: Format of boxes data encoding.
      @param sort_result_descending: Flag that specifies whenever it is necessary to sort selected
                                     boxes across batches or not.
      @param output_type: Output element type.
      @return The new node which performs NonMaxSuppression
      """
      if max_output_boxes_per_class is None:
          max_output_boxes_per_class = make_constant_node(0, np.int64)
      if iou_threshold is None:
          iou_threshold = make_constant_node(0, np.float32)
      if score_threshold is None:
          score_threshold = make_constant_node(0, np.float32)
  
      inputs = as_nodes(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold)
      attributes = {
          "box_encoding": box_encoding,
          "sort_result_descending": sort_result_descending,
          "output_type": output_type,
      }
  
      return _get_node_factory_opset3().create("NonMaxSuppression", inputs, attributes)
  
  
  @nameable_op
  def non_zero(data: NodeInput, output_type: str = "i64", name: Optional[str] = None,) -> Node:
      """Return the indices of the elements that are non-zero.
  
      @param data: Input data.
      @param output_type: Output tensor type.
  
      @return The new node which performs NonZero
      """
      return _get_node_factory_opset3().create(
          "NonZero",
          [as_node(data)],
          {"output_type": output_type}
      )
  
  
  @nameable_op
  def read_value(init_value: NodeInput, variable_id: str, name: Optional[str] = None) -> Node:
      """Return a node which produces the Assign operation.
  
      @param init_value:   Node producing a value to be returned instead of an unassigned variable.
      @param variable_id:  Id of a variable to be read.
      @param name:         Optional name for output node.
      @return ReadValue node
      """
      return _get_node_factory_opset3().create(
          "ReadValue",
          [as_node(init_value)],
          {"variable_id": variable_id}
      )
  
  
  @nameable_op
  def rnn_cell(
      X: NodeInput,
      initial_hidden_state: NodeInput,
      W: NodeInput,
      R: NodeInput,
      B: NodeInput,
      hidden_size: int,
      activations: List[str],
      activations_alpha: List[float],
      activations_beta: List[float],
      clip: float = 0.0,
      name: Optional[str] = None,
  ) -> Node:
      """Perform RNNCell operation on tensor from input node.
  
      It follows notation and equations defined as in ONNX standard:
      https://github.com/onnx/onnx/blob/master/docs/Operators.md#RNN
  
      Note this class represents only single *cell* and not whole RNN *layer*.
  
      @param X:                       The input tensor with shape: [batch_size, input_size].
      @param initial_hidden_state:    The hidden state tensor at current time step with shape:
                                      [batch_size, hidden_size].
      @param W:                       The weight tensor with shape: [hidden_size, input_size].
      @param R:                       The recurrence weight tensor with shape: [hidden_size,
                                      hidden_size].
      @param B:                       The sum of biases (weight and recurrence) with shape: [hidden_size].
      @param hidden_size:             The number of hidden units for recurrent cell.
                                      Specifies hidden state size.
      @param activations:             The vector of activation functions used inside recurrent cell.
      @param activation_alpha:        The vector of alpha parameters for activation functions in
                                      order respective to activation list.
      @param activation_beta:         The vector of beta parameters for activation functions in order
                                      respective to activation list.
      @param clip:                    The value defining clipping range [-clip, clip] on input of
                                      activation functions.
      @param name:                    Optional output node name.
      @return   The new node performing a RNNCell operation on tensor from input node.
      """
      if activations is None:
          activations = ["tanh"]
      if activations_alpha is None:
          activations_alpha = []
      if activations_beta is None:
          activations_beta = []
  
      input_nodes = as_nodes(X, initial_hidden_state, W, R, B)
      attributes = {
          "hidden_size": hidden_size,
          "activations": activations,
          "activations_alpha": activations_alpha,
          "activations_beta": activations_beta,
          "clip": clip,
      }
      return _get_node_factory_opset3().create("RNNCell", input_nodes, attributes)
  
  
  @nameable_op
  def roi_align(
      data: NodeInput,
      rois: NodeInput,
      batch_indices: NodeInput,
      pooled_h: int,
      pooled_w: int,
      sampling_ratio: int,
      spatial_scale: float,
      mode: str,
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which performs ROIAlign.
  
      @param data: Input data.
      @param rois: RoIs (Regions of Interest) to pool over.
      @param batch_indices: Tensor with each element denoting the index of
                            the corresponding image in the batch.
      @param pooled_h: Height of the ROI output feature map.
      @param pooled_w: Width of the ROI output feature map.
      @param sampling_ratio: Number of bins over height and width to use to calculate
                             each output feature map element.
      @param spatial_scale: Multiplicative spatial scale factor to translate ROI coordinates.
      @param mode: Method to perform pooling to produce output feature map elements.
  
      @return The new node which performs ROIAlign
      """
      inputs = as_nodes(data, rois, batch_indices)
      attributes = {
          "pooled_h": pooled_h,
          "pooled_w": pooled_w,
          "sampling_ratio": sampling_ratio,
          "spatial_scale": spatial_scale,
          "mode": mode,
      }
      return _get_node_factory_opset3().create("ROIAlign", inputs, attributes)
  
  
  @nameable_op
  def scatter_elements_update(
      data: NodeInput,
      indices: NodeInput,
      updates: NodeInput,
      axis: NodeInput,
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which produces a ScatterElementsUpdate operation.
  
      @param data:    The input tensor to be updated.
      @param indices: The tensor with indexes which will be updated.
      @param updates: The tensor with update values.
      @param axis:    The axis for scatter.
      @return ScatterElementsUpdate node
  
      ScatterElementsUpdate creates a copy of the first input tensor with updated elements
      specified with second and third input tensors.
  
      For each entry in `updates`, the target index in `data` is obtained by combining
      the corresponding entry in `indices` with the index of the entry itself: the
      index-value for dimension equal to `axis` is obtained from the value of the
      corresponding entry in `indices` and the index-value for dimension not equal
      to `axis` is obtained from the index of the entry itself.
  
      """
      return _get_node_factory_opset3().create(
          "ScatterElementsUpdate", as_nodes(data, indices, updates, axis)
      )
  
  
  @nameable_op
  def scatter_update(
      data: Node, indices: NodeInput, updates: NodeInput, axis: NodeInput, name: Optional[str] = None
  ) -> Node:
      """Return a node which produces a ScatterUpdate operation.
  
      ScatterUpdate sets new values to slices from data addressed by indices.
  
      @param data:    The input tensor to be updated.
      @param indices: The tensor with indexes which will be updated.
      @param updates: The tensor with update values.
      @param axis:    The axis at which elements will be updated.
      @return ScatterUpdate node
      """
      return _get_node_factory_opset3().create(
          "ScatterUpdate",
          as_nodes(data, indices, updates, axis)
      )
  
  
  @nameable_op
  def shape_of(data: NodeInput, output_type: str = "i64", name: Optional[str] = None) -> Node:
      """Return a node which produces a tensor containing the shape of its input data.
  
      @param data: The tensor containing the input data.
      @param output_type: Output element type.
      @return ShapeOf node
      """
      return _get_node_factory_opset3().create(
          "ShapeOf",
          [as_node(data)],
          {"output_type": output_type}
      )
  
  
  @nameable_op
  def shuffle_channels(data: Node, axis: int, group: int, name: Optional[str] = None) -> Node:
      """Perform permutation on data in the channel dimension of the input tensor.
  
      @param data: The node with input tensor.
      @param axis: Channel dimension index in the data tensor.
                   A negative value means that the index should be calculated
                   from the back of the input data shape.
      @param group: The channel dimension specified by the axis parameter
                   should be split into this number of groups.
      @param name: Optional output node name.
      @return The new node performing a permutation on data in the channel dimension
               of the input tensor.
  
      The operation is the equivalent with the following transformation of the input tensor
      `data` of shape [N, C, H, W]:
  
      `data_reshaped` = reshape(`data`, [N, group, C / group, H * W])
  
      `data_trnasposed` = transpose(`data_reshaped`, [0, 2, 1, 3])
  
      `output` = reshape(`data_trnasposed`, [N, C, H, W])
  
      For example:
  
      @code{.py}
          Inputs: tensor of shape [1, 6, 2, 2]
  
                  data = [[[[ 0.,  1.], [ 2.,  3.]],
                           [[ 4.,  5.], [ 6.,  7.]],
                           [[ 8.,  9.], [10., 11.]],
                           [[12., 13.], [14., 15.]],
                           [[16., 17.], [18., 19.]],
                           [[20., 21.], [22., 23.]]]]
  
                  axis = 1
                  groups = 3
  
          Output: tensor of shape [1, 6, 2, 2]
  
                  output = [[[[ 0.,  1.], [ 2.,  3.]],
                             [[ 8.,  9.], [10., 11.]],
                             [[16., 17.], [18., 19.]],
                             [[ 4.,  5.], [ 6.,  7.]],
                             [[12., 13.], [14., 15.]],
                             [[20., 21.], [22., 23.]]]]
      @endcode
      """
      return _get_node_factory_opset3().create(
          "ShuffleChannels", [as_node(data)], {"axis": axis, "group": group}
      )
  
  
  @nameable_op
  def topk(
      data: NodeInput,
      k: NodeInput,
      axis: int,
      mode: str,
      sort: str,
      index_element_type: str = "i32",
      name: Optional[str] = None,
  ) -> Node:
      """Return a node which performs TopK.
  
      @param data: Input data.
      @param k: K.
      @param axis: TopK Axis.
      @param mode: Compute TopK largest ('max') or smallest ('min')
      @param sort: Order of output elements (sort by: 'none', 'index' or 'value')
      @param index_element_type: Type of output tensor with indices.
      @return The new node which performs TopK (both indices and values)
      """
      return _get_node_factory_opset3().create(
          "TopK",
          as_nodes(data, k),
          {"axis": axis, "mode": mode, "sort": sort, "index_element_type": index_element_type},
      )
  