[PYTHON] Expose If, Loop and TensorIterator API (openvinotoolkit#9603)

src/bindings/python/src/compatibility/pyngraph/discrete_type_info.cpp

@@ -15,7 +15,8 @@ namespace py = pybind11;
 void regclass_pyngraph_DiscreteTypeInfo(py::module m) {
     py::class_<ngraph::DiscreteTypeInfo, std::shared_ptr<ngraph::DiscreteTypeInfo>> discrete_type_info(
         m,
-        "DiscreteTypeInfo");
+        "DiscreteTypeInfo",
+        py::module_local());
     discrete_type_info.doc() = "ngraph.impl.DiscreteTypeInfo wraps ngraph::DiscreteTypeInfo";
 
     // operator overloading

src/bindings/python/src/openvino/runtime/__init__.py

@@ -53,6 +53,7 @@
 from openvino.pyopenvino import Shape
 from openvino.pyopenvino import Strides
 from openvino.pyopenvino import CoordinateDiff
+from openvino.pyopenvino import DiscreteTypeInfo
 from openvino.pyopenvino import AxisSet
 from openvino.pyopenvino import AxisVector
 from openvino.pyopenvino import Coordinate

src/bindings/python/src/openvino/runtime/op/__init__.py

@@ -21,3 +21,6 @@
 Constant.get_data = lambda self: np.array(self, copy=True)
 
 from openvino.pyopenvino.op import Parameter
+from openvino.pyopenvino.op import if_op
+from openvino.pyopenvino.op import loop
+from openvino.pyopenvino.op import tensor_iterator

src/bindings/python/src/openvino/runtime/op/util/__init__.py

@@ -15,3 +15,8 @@
 from openvino.pyopenvino.op.util import IndexReduction
 from openvino.pyopenvino.op.util import VariableInfo
 from openvino.pyopenvino.op.util import Variable
+from openvino.pyopenvino.op.util import MergedInputDescription
+from openvino.pyopenvino.op.util import InvariantInputDescription
+from openvino.pyopenvino.op.util import SliceInputDescription
+from openvino.pyopenvino.op.util import ConcatOutputDescription
+from openvino.pyopenvino.op.util import BodyOutputDescription

src/bindings/python/src/openvino/runtime/opset1/ops.py

@@ -8,7 +8,7 @@
 from functools import partial
 
 from openvino.runtime import Node, PartialShape, Shape
-from openvino.runtime.op import Constant, Parameter
+from openvino.runtime.op import Constant, Parameter, tensor_iterator
 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 (
@@ -17,14 +17,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,
@@ -2754,49 +2746,6 @@ def tanh(node: NodeInput, name: Optional[str] = None) -> Node:
     return _get_node_factory_opset1().create("Tanh", [node])
 
 
-@nameable_op
-def tensor_iterator(
-    inputs: List[Node],
-    graph_body: GraphBody,
-    slice_input_desc: List[TensorIteratorSliceInputDesc],
-    merged_input_desc: List[TensorIteratorMergedInputDesc],
-    invariant_input_desc: List[TensorIteratorInvariantInputDesc],
-    body_output_desc: List[TensorIteratorBodyOutputDesc],
-    concat_output_desc: List[TensorIteratorConcatOutputDesc],
-    name: Optional[str] = None,
-) -> Node:
-    """Perform recurrent execution of the network described in the body, iterating through the data.
-
-    @param      inputs:                The provided to TensorIterator operator.
-    @param      graph_body:            The graph representing the body we execute.
-    @param      slice_input_desc:      The descriptors describing sliced inputs, that is nodes
-                                       representing tensors we iterate through, processing single
-                                       data slice in one iteration.
-    @param      merged_input_desc:     The descriptors describing merged inputs, that is nodes
-                                       representing variables with initial value at first iteration,
-                                       which may be changing through iterations.
-    @param      invariant_input_desc:  The descriptors describing invariant inputs, that is nodes
-                                       representing variable with persistent value through all
-                                       iterations.
-    @param      body_output_desc:      The descriptors describing body outputs from specified
-                                       iteration.
-    @param      concat_output_desc:    The descriptors describing specified output values through
-                                       all the iterations concatenated into one node.
-    @param      name:                  The optional name for output node.
-    @return   Node representing TensorIterator operation.
-    """
-    attributes = {
-        "body": graph_body.serialize(),
-        "input_descriptions": {"slice_input_desc": [desc.serialize() for desc in slice_input_desc],
-                               "merged_input_desc": [desc.serialize() for desc in merged_input_desc],
-                               "invariant_input_desc": [desc.serialize() for desc in invariant_input_desc]},
-        "output_descriptions": {"body_output_desc": [desc.serialize() for desc in body_output_desc],
-                                "concat_output_desc": [desc.serialize() for desc in concat_output_desc]}
-    }
-
-    return _get_node_factory_opset1().create("TensorIterator", as_nodes(*inputs), attributes)
-
-
 @nameable_op
 def tile(data: NodeInput, repeats: NodeInput, name: Optional[str] = None) -> Node:
     """Return a node which dynamically repeats(replicates) the input data tensor.

src/bindings/python/src/openvino/runtime/opset2/ops.py

@@ -18,14 +18,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,

src/bindings/python/src/openvino/runtime/opset3/ops.py

@@ -18,14 +18,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,

src/bindings/python/src/openvino/runtime/opset4/ops.py

@@ -18,14 +18,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,

src/bindings/python/src/openvino/runtime/opset5/ops.py

@@ -8,7 +8,7 @@
 from functools import partial
 
 from openvino.runtime import Node, Shape
-from openvino.runtime.op import Constant, Parameter
+from openvino.runtime.op import Constant, Parameter, loop
 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 (
@@ -18,14 +18,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,
@@ -366,62 +358,3 @@ def rnn_sequence(
     }
 
     return _get_node_factory_opset5().create("RNNSequence", inputs, attributes)
-
-
-@nameable_op
-def loop(
-    trip_count: NodeInput,
-    execution_condition: NodeInput,
-    inputs: List[Node],
-    graph_body: GraphBody,
-    slice_input_desc: List[TensorIteratorSliceInputDesc],
-    merged_input_desc: List[TensorIteratorMergedInputDesc],
-    invariant_input_desc: List[TensorIteratorInvariantInputDesc],
-    body_output_desc: List[TensorIteratorBodyOutputDesc],
-    concat_output_desc: List[TensorIteratorConcatOutputDesc],
-    body_condition_output_idx: int,
-    current_iteration_input_idx: int = -1,
-    name: Optional[str] = None,
-) -> Node:
-    """Perform recurrent execution of the network described in the body, iterating through the data.
-
-    @param trip_count: A scalar or 1D tensor with 1 element specifying
-        maximum number of iterations.
-    @param execution_condition: A scalar or 1D tensor with 1 element
-        specifying whether to execute the first iteration or not.
-    @param      inputs:                The provided to TensorIterator operator.
-    @param      graph_body:            The graph representing the body we execute.
-    @param      slice_input_desc:      The descriptors describing sliced inputs, that is nodes
-                                       representing tensors we iterate through, processing single
-                                       data slice in one iteration.
-    @param      merged_input_desc:     The descriptors describing merged inputs, that is nodes
-                                       representing variables with initial value at first iteration,
-                                       which may be changing through iterations.
-    @param      invariant_input_desc:  The descriptors describing invariant inputs, that is nodes
-                                       representing variable with persistent value through all
-                                       iterations.
-    @param      body_output_desc:      The descriptors describing body outputs from specified
-                                       iteration.
-    @param      concat_output_desc:    The descriptors describing specified output values through
-                                       all the iterations concatenated into one node.
-    @param      body_condition_output_idx:    Determines the purpose of the corresponding result in
-                                              the graph_body. This result will determine the dynamic
-                                              exit condition. If the value of this result is False,
-                                              then iterations stop.
-    @param      current_iteration_input_idx:  Determines the purpose of the corresponding parameter
-                                              in the graph_body. This parameter will be used as
-                                              an iteration counter. Optional.
-    @return: The new node which performs Loop.
-    """
-    attributes = {
-        "body": graph_body.serialize(),
-        "input_descriptions": {"slice_input_desc": [desc.serialize() for desc in slice_input_desc],
-                               "merged_input_desc": [desc.serialize() for desc in merged_input_desc],
-                               "invariant_input_desc": [desc.serialize() for desc in invariant_input_desc]},
-        "output_descriptions": {"body_output_desc": [desc.serialize() for desc in body_output_desc],
-                                "concat_output_desc": [desc.serialize() for desc in concat_output_desc]},
-        "special_body_ports": {"body_condition_output_idx": body_condition_output_idx,
-                               "current_iteration_input_idx": current_iteration_input_idx}
-    }
-    return _get_node_factory_opset5().create("Loop", as_nodes(trip_count, execution_condition, *inputs),
-                                             attributes)

src/bindings/python/src/openvino/runtime/opset6/ops.py

@@ -18,14 +18,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,

src/bindings/python/src/openvino/runtime/opset7/ops.py

@@ -17,14 +17,6 @@
     is_positive_value,
 )
 from openvino.runtime.utils.node_factory import NodeFactory
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorSliceInputDesc,
-    TensorIteratorMergedInputDesc,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-    TensorIteratorConcatOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     NumericData,

src/bindings/python/src/openvino/runtime/opset8/ops.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 from openvino.runtime.exceptions import UserInputError
+from openvino.runtime.op import Constant, Parameter, if_op
 from openvino.runtime import Node
 from openvino.runtime.opset_utils import _get_node_factory
 from openvino.runtime.utils.decorators import nameable_op
@@ -15,11 +16,6 @@
     is_non_negative_value,
     is_positive_value,
 )
-from openvino.runtime.utils.tensor_iterator_types import (
-    GraphBody,
-    TensorIteratorInvariantInputDesc,
-    TensorIteratorBodyOutputDesc,
-)
 from openvino.runtime.utils.types import (
     NodeInput,
     TensorShape,
@@ -357,42 +353,6 @@ def random_uniform(
     return _get_node_factory_opset8().create("RandomUniform", inputs, attributes)
 
 
-@nameable_op
-def if_op(
-        condition: NodeInput,
-        inputs: List[Node],
-        bodies: Tuple[GraphBody, GraphBody],
-        input_desc: Tuple[List[TensorIteratorInvariantInputDesc], List[TensorIteratorInvariantInputDesc]],
-        output_desc: Tuple[List[TensorIteratorBodyOutputDesc], List[TensorIteratorBodyOutputDesc]],
-        name: Optional[str] = None,
-) -> Node:
-    """Execute one of the bodies depending on condtion value.
-
-    @param      condition:             A scalar or 1D tensor with 1 element specifying body will be executed.
-                                       If condition is True, then body will be executed, False - else_body.
-    @param      inputs:                The provided inputs to If operation.
-    @param      bodies:                Two graphs (then_body, else_body) which will be executed depending on
-                                       condition value.
-    @param      input_desc             Two lists (for then_body and else_body) which contain rules how If
-                                       inputs are connected with body parameters.
-    @param      output_desc:           Two lists (for then_body and else_body) which contain rules how If
-                                       outputs are connected with body results.
-    @param      name:                  The optional name for the created output node.
-
-    @return: The new node which performs If operation.
-    """
-    attributes = {
-        "then_body": bodies[0].serialize(),
-        "else_body": bodies[1].serialize(),
-        "then_inputs": {"invariant_input_desc": [desc.serialize() for desc in input_desc[0]]},
-        "else_inputs": {"invariant_input_desc": [desc.serialize() for desc in input_desc[1]]},
-        "then_outputs": {"body_output_desc": [desc.serialize() for desc in output_desc[0]]},
-        "else_outputs": {"body_output_desc": [desc.serialize() for desc in output_desc[1]]}
-    }
-    return _get_node_factory_opset8().create("If", as_nodes(condition, *inputs),
-                                             attributes)
-
-
 @nameable_op
 def slice(
         data: NodeInput,