[PyOV] python op implementation (#24487)

### Details:
 - based on: https://github.com/openvinotoolkit/openvino/pull/23612

### Tickets:
 - CVS-141051

---------

Co-authored-by: Michal Lukaszewski <michal.lukaszewski@intel.com>

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

@@ -53,6 +53,7 @@ from openvino.runtime import layout_helpers
 
 from openvino._pyopenvino import RemoteContext
 from openvino._pyopenvino import RemoteTensor
+from openvino._pyopenvino import Op
 
 # libva related:
 from openvino._pyopenvino import VAContext

src/bindings/python/src/pyopenvino/graph/discrete_type_info.cpp

@@ -14,10 +14,32 @@
 
 namespace py = pybind11;
 
+// DiscreteTypeInfo doesn't own provided memory. Wrapper allows to avoid leaks.
+class DiscreteTypeInfoWrapper : public ov::DiscreteTypeInfo {
+private:
+    const std::string name_str;
+    const std::string version_id_str;
+
+public:
+    DiscreteTypeInfoWrapper(std::string _name_str, std::string _version_id_str)
+        : DiscreteTypeInfo(nullptr, nullptr, nullptr),
+          name_str(std::move(_name_str)),
+          version_id_str(std::move(_version_id_str)) {
+        name = name_str.c_str();
+        version_id = version_id_str.c_str();
+    }
+};
+
 void regclass_graph_DiscreteTypeInfo(py::module m) {
     py::class_<ov::DiscreteTypeInfo, std::shared_ptr<ov::DiscreteTypeInfo>> discrete_type_info(m, "DiscreteTypeInfo");
     discrete_type_info.doc() = "openvino.runtime.DiscreteTypeInfo wraps ov::DiscreteTypeInfo";
 
+    discrete_type_info.def(py::init([](const std::string& name, const std::string& version_id) {
+                               return std::make_shared<DiscreteTypeInfoWrapper>(name, version_id);
+                           }),
+                           py::arg("name"),
+                           py::arg("version_id"));
+
     // operator overloading
     discrete_type_info.def(py::self < py::self);
     discrete_type_info.def(py::self <= py::self);
@@ -41,9 +63,9 @@ void regclass_graph_DiscreteTypeInfo(py::module m) {
         if (self.parent != nullptr) {
             std::string parent_version = std::string(self.parent->version_id);
             std::string parent_name = self.parent->name;
-            return "<" + class_name + ": " + name + " v" + version + " Parent(" + parent_name + " v" + parent_version +
+            return "<" + class_name + ": " + name + " " + version + " Parent(" + parent_name + " v" + parent_version +
                    ")" + ">";
         }
-        return "<" + class_name + ": " + name + " v" + version + ">";
+        return "<" + class_name + ": " + name + " " + version + ">";
     });
 }

src/bindings/python/src/pyopenvino/graph/op.cpp

@@ -0,0 +1,63 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/op/op.hpp"
+
+#include <pybind11/functional.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+
+#include <pyopenvino/graph/op.hpp>
+
+#include "openvino/core/attribute_visitor.hpp"
+#include "openvino/core/node.hpp"
+
+namespace py = pybind11;
+
+/// Trampoline class to support inheritence from TorchDecoder in Python
+class PyOp : public ov::op::Op {
+public:
+    using ov::op::Op::Op;
+
+    // Keeps a reference to the Python object to manage its lifetime
+    PyOp(const py::object& py_obj) : py_handle(py_obj) {}
+
+    void validate_and_infer_types() override {
+        PYBIND11_OVERRIDE(void, ov::op::Op, validate_and_infer_types);
+    }
+
+    bool visit_attributes(ov::AttributeVisitor& visitor) override {
+        // PYBIND11_OVERRIDE_PURE(bool, Op, visit_attributes, visitor);
+        //  Requires binding for visitor
+        //  Now works only for operations without attributes
+        return true;
+    }
+
+    std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override {
+        PYBIND11_OVERRIDE_PURE(std::shared_ptr<Node>, ov::op::Op, clone_with_new_inputs, new_args);
+    }
+
+    const type_info_t& get_type_info() const override {
+        PYBIND11_OVERRIDE(const ov::Node::type_info_t&, ov::op::Op, get_type_info);
+    }
+
+    bool evaluate(ov::TensorVector& output_values, const ov::TensorVector& input_values) const override {
+        PYBIND11_OVERRIDE(bool, ov::op::Op, evaluate, output_values, input_values);
+    }
+
+    bool has_evaluate() const override {
+        PYBIND11_OVERRIDE(bool, ov::op::Op, has_evaluate);
+    }
+
+private:
+    py::object py_handle;  // Holds the Python object to manage its lifetime
+};
+
+void regclass_graph_Op(py::module m) {
+    py::class_<ov::op::Op, std::shared_ptr<ov::op::Op>, PyOp, ov::Node>(m, "Op").def(
+        py::init([](const py::object& py_obj) {
+            return PyOp(py_obj);
+        }));
+}

src/bindings/python/src/pyopenvino/graph/op.hpp

@@ -0,0 +1,11 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+void regclass_graph_Op(py::module m);

src/bindings/python/src/pyopenvino/pyopenvino.cpp

@@ -19,6 +19,7 @@
 #include "pyopenvino/graph/node_factory.hpp"
 #include "pyopenvino/graph/node_input.hpp"
 #include "pyopenvino/graph/node_output.hpp"
+#include <pyopenvino/graph/op.hpp>
 #if defined(ENABLE_OV_ONNX_FRONTEND)
 #    include "pyopenvino/graph/onnx_import/onnx_import.hpp"
 #endif
@@ -223,6 +224,7 @@ PYBIND11_MODULE(_pyopenvino, m) {
     regclass_graph_Shape(m);
     regclass_graph_PartialShape(m);
     regclass_graph_Node(m);
+    regclass_graph_Op(m);
     regclass_graph_Input(m);
     regclass_graph_NodeFactory(m);
     regclass_graph_Strides(m);

src/bindings/python/tests/test_graph/test_op.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+from openvino import Op
+from openvino.runtime import CompiledModel, DiscreteTypeInfo, Model, Shape, compile_model, Tensor
+import openvino.runtime.opset14 as ops
+
+
+class CustomOp(Op):
+    class_type_info = DiscreteTypeInfo("Custom", "extension")
+
+    def __init__(self, inputs):
+        super().__init__(self)
+        self.set_arguments(inputs)
+        self.constructor_validate_and_infer_types()
+
+    def validate_and_infer_types(self):
+        self.set_output_type(0, self.get_input_element_type(0), self.get_input_partial_shape(0))
+
+    def clone_with_new_inputs(self, new_inputs):
+        return CustomOp(new_inputs)
+
+    def get_type_info(self):
+        return CustomOp.class_type_info
+
+    def evaluate(self, outputs, inputs):
+        inputs[0].copy_to(outputs[0])
+        return True
+
+    def has_evaluate(self):
+        return True
+
+
+def create_snake_model():
+    input_shape = [1, 3, 32, 32]
+    param1 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data")
+    custom_op = CustomOp([param1])
+    custom_op.set_friendly_name("custom_" + str(0))
+
+    for i in range(20):
+        custom_op = CustomOp([custom_op])
+        custom_op.set_friendly_name("custom_" + str(i + 1))
+    return Model(custom_op, [param1], "TestModel")
+
+
+class CustomAdd(Op):
+    class_type_info = DiscreteTypeInfo("CustomAdd", "extension")
+
+    def __init__(self, inputs):
+        super().__init__(self)
+        self.set_arguments(inputs)
+        self.constructor_validate_and_infer_types()
+
+    def validate_and_infer_types(self):
+        self.set_output_type(0, self.get_input_element_type(0), self.get_input_partial_shape(0))
+
+    def clone_with_new_inputs(self, new_inputs):
+        node = CustomAdd(new_inputs)
+        return node
+
+    def get_type_info(self):
+        return CustomAdd.class_type_info
+
+
+def create_add_model():
+    input_shape = [2, 1]
+
+    param1 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data1")
+    param2 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data2")
+    custom_add = CustomAdd(inputs=[param1, param2])
+    custom_add.set_friendly_name("test_add")
+    res = ops.result(custom_add, name="result")
+    return Model(res, [param1, param2], "AddModel")
+
+
+def test_custom_add_op():
+    data1 = np.array([1, 2, 3])
+    data2 = np.array([4, 5, 6])
+
+    node1 = ops.constant(data1, dtype=np.float32)
+    node2 = ops.constant(data2, dtype=np.float32)
+    inputs = [node1.output(0), node2.output(0)]
+    custom_op = CustomAdd(inputs=inputs)
+    custom_op.set_friendly_name("test_add")
+
+    assert custom_op.get_input_size() == 2
+    assert custom_op.get_output_size() == 1
+    assert custom_op.get_type_name() == "CustomAdd"
+    assert list(custom_op.get_output_shape(0)) == [3]
+    assert custom_op.friendly_name == "test_add"
+
+
+def test_custom_add_model():
+    model = create_add_model()
+
+    assert isinstance(model, Model)
+
+    ordered_ops = model.get_ordered_ops()
+    assert len(ordered_ops) == 4
+
+    op_types = [op.get_type_name() for op in ordered_ops]
+    assert op_types == ["Parameter", "Parameter", "CustomAdd", "Result"]
+
+
+def test_custom_op():
+    model = create_snake_model()
+    compiled_model = compile_model(model)
+
+    assert isinstance(compiled_model, CompiledModel)
+    request = compiled_model.create_infer_request()
+
+    input_data = np.ones([1, 3, 32, 32], dtype=np.float32)
+    expected_output = np.maximum(0.0, input_data)
+
+    input_tensor = Tensor(input_data)
+    results = request.infer({"data": input_tensor})
+    assert np.allclose(results[list(results)[0]], expected_output, 1e-4, 1e-4)

src/bindings/python/tests/test_graph/test_type_info.py

@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+from openvino import Op
+from openvino.runtime import DiscreteTypeInfo, Shape
+import openvino.runtime.opset14 as ops
+from openvino.runtime.utils.node_factory import NodeFactory
+
+
+class CustomAdd(Op):
+    class_type_info = DiscreteTypeInfo("CustomAdd", "extension")
+
+    def __init__(self, inputs):
+        super().__init__(self)
+        self.set_arguments(inputs)
+        self.constructor_validate_and_infer_types()
+
+    def validate_and_infer_types(self):
+        self.set_output_type(0, self.get_input_element_type(0), self.get_input_partial_shape(0))
+
+    def clone_with_new_inputs(self, new_inputs):
+        node = CustomAdd(new_inputs)
+        return node
+
+    def get_type_info(self):
+        return CustomAdd.class_type_info
+
+
+def test_node_factory_type_info():
+    shape = [2, 2]
+    dtype = np.int8
+    parameter_a = ops.parameter(shape, dtype=dtype, name="A")
+    parameter_b = ops.parameter(shape, dtype=dtype, name="B")
+
+    factory = NodeFactory("opset1")
+    arguments = NodeFactory._arguments_as_outputs([parameter_a, parameter_b])
+    node = factory.create("Add", arguments, {})
+
+    type_info = node.get_type_info()
+    assert isinstance(type_info, DiscreteTypeInfo)
+    assert type_info.name == "Add"
+    assert type_info.version_id == "opset1"
+
+
+def test_discrete_type_info():
+    type_info = DiscreteTypeInfo("Custom", "extension")
+    assert isinstance(type_info, DiscreteTypeInfo)
+    assert type_info.name == "Custom"
+    assert type_info.version_id == "extension"
+
+
+def test_custom_add_op_type_info():
+    data1 = np.array([1, 2, 3])
+    data2 = np.array([4, 5, 6])
+
+    node1 = ops.constant(data1, dtype=np.float32)
+    node2 = ops.constant(data2, dtype=np.float32)
+    inputs = [node1.output(0), node2.output(0)]
+    custom_op = CustomAdd(inputs=inputs)
+
+    type_info = custom_op.get_type_info()
+    assert isinstance(type_info, DiscreteTypeInfo)
+    assert type_info.name == "CustomAdd"
+    assert type_info.version_id == "extension"
+
+
+def test_add_type_info():
+    input_shape = [2, 1]
+    param1 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data1")
+    param2 = ops.parameter(Shape(input_shape), dtype=np.float32, name="data2")
+    add = ops.add(param1, param2)
+
+    type_info = add.get_type_info()
+    assert isinstance(type_info, DiscreteTypeInfo)
+    assert type_info.name == "Add"
+    assert type_info.version_id == "opset1"

tools/benchmark_tool/openvino/__init__.py

@@ -53,6 +53,7 @@ from openvino.runtime import layout_helpers
 
 from openvino._pyopenvino import RemoteContext
 from openvino._pyopenvino import RemoteTensor
+from openvino._pyopenvino import Op
 
 # libva related:
 from openvino._pyopenvino import VAContext

tools/mo/openvino/__init__.py

@@ -50,6 +50,7 @@ try:
 
     from openvino._pyopenvino import RemoteContext
     from openvino._pyopenvino import RemoteTensor
+    from openvino._pyopenvino import Op
 
     # libva related:
     from openvino._pyopenvino import VAContext

tools/openvino_dev/src/openvino/__init__.py

@@ -50,6 +50,7 @@ try:
 
     from openvino._pyopenvino import RemoteContext
     from openvino._pyopenvino import RemoteTensor
+    from openvino._pyopenvino import Op
 
     # libva related:
     from openvino._pyopenvino import VAContext

tools/ovc/openvino/__init__.py

@@ -53,6 +53,7 @@ from openvino.runtime import layout_helpers
 
 from openvino._pyopenvino import RemoteContext
 from openvino._pyopenvino import RemoteTensor
+from openvino._pyopenvino import Op
 
 # libva related:
 from openvino._pyopenvino import VAContext