[PyOV] OpExtension

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

@@ -32,6 +32,7 @@
 from openvino._pyopenvino import Symbol
 from openvino._pyopenvino import Dimension
 from openvino._pyopenvino import Input
+from openvino._pyopenvino import OpExtension
 from openvino._pyopenvino import Output
 from openvino._pyopenvino import Node
 from openvino._pyopenvino import Strides
@@ -49,14 +50,14 @@
 from openvino._pyopenvino import layout_helpers
 from openvino._pyopenvino import RemoteContext
 from openvino._pyopenvino import RemoteTensor
-from openvino._pyopenvino import Op
 
 # Import public classes from _ov_api
 from openvino._ov_api import Model
 from openvino._ov_api import Core
 from openvino._ov_api import CompiledModel
 from openvino._ov_api import InferRequest
 from openvino._ov_api import AsyncInferQueue
+from openvino._ov_api import Op
 
 # Import all public modules
 from openvino import runtime as runtime

src/bindings/python/src/openvino/_ov_api.py

@@ -3,16 +3,16 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from types import TracebackType
-from typing import Any, Iterable, Union, Optional, Dict, Type
+from typing import Any, Iterable, Union, Optional, Dict, Type, List
 from pathlib import Path
 
 
 from openvino._pyopenvino import Model as ModelBase
 from openvino._pyopenvino import Core as CoreBase
 from openvino._pyopenvino import CompiledModel as CompiledModelBase
 from openvino._pyopenvino import AsyncInferQueue as AsyncInferQueueBase
-from openvino._pyopenvino import Tensor
-from openvino._pyopenvino import Node
+from openvino._pyopenvino import Op as OpBase
+from openvino._pyopenvino import Node, Output, Tensor
 
 from openvino.utils.data_helpers import (
     OVDict,
@@ -22,6 +22,15 @@
 )
 
 
+class Op(OpBase):
+    def __init__(self, py_obj: "Op", inputs: Optional[List[Union[Node, Output]]] = None) -> None:
+        super().__init__(py_obj)
+        self._update_type_info()
+        if inputs is not None:
+            self.set_arguments(inputs)
+            self.constructor_validate_and_infer_types()
+
+
 class ModelMeta(type):
     def __dir__(cls) -> list:
         return list(set(cls.__dict__.keys()) | set(dir(ModelBase)))

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

@@ -22,7 +22,7 @@
 # extensions
 from openvino._pyopenvino import DecoderTransformationExtension
 from openvino._pyopenvino import ConversionExtension
-from openvino._pyopenvino import OpExtension
+from openvino._pyopenvino.frontend import OpExtension
 from openvino._pyopenvino import ProgressReporterExtension
 from openvino._pyopenvino import TelemetryExtension
 

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

@@ -14,6 +14,7 @@
 
 #include "common.hpp"
 #include "pyopenvino/core/remote_context.hpp"
+#include "pyopenvino/graph/op_extension.hpp"
 #include "pyopenvino/utils/utils.hpp"
 
 namespace py = pybind11;
@@ -703,6 +704,19 @@ void regclass_Core(py::module m) {
             :type extensions: list[openvino.runtime.Extension]
         )");
 
+    cls.def(
+        "add_extension",
+        [](ov::Core& self, py::object dtype) {
+            self.add_extension(PyOpExtension(dtype));
+        },
+        py::arg("custom_op"),
+        R"(
+            Registers custom Op to a Core object.
+
+            :param custom_op: type of custom Op
+            :type custom_op: openvino.Op
+        )");
+
     cls.def("get_available_devices",
             &ov::Core::get_available_devices,
             py::call_guard<py::gil_scoped_release>(),

src/bindings/python/src/pyopenvino/core/extension.cpp

@@ -9,7 +9,7 @@
 #include <pybind11/stl.h>
 #include <pybind11/stl_bind.h>
 
-#include "openvino/frontend/manager.hpp"
+#include "openvino/core/extension.hpp"
 #include "pyopenvino/core/common.hpp"
 
 namespace py = pybind11;

src/bindings/python/src/pyopenvino/frontend/extension.cpp

@@ -142,9 +142,9 @@ void regclass_frontend_ProgressReporterExtension(py::module m) {
 }
 
 void regclass_frontend_OpExtension(py::module m) {
-    py::class_<OpExtension<void>, std::shared_ptr<OpExtension<void>>, ConversionExtension> ext(m,
-                                                                                               "OpExtension",
-                                                                                               py::dynamic_attr());
+    py::module frontend = m.def_submodule("frontend");
+    py::class_<ov::frontend::OpExtension<void>, std::shared_ptr<ov::frontend::OpExtension<void>>, ConversionExtension>
+        ext(frontend, "OpExtension", py::dynamic_attr());
 
     ext.def(py::init([](const std::string& fw_type_name,
                         const std::map<std::string, std::string>& attr_names_map,

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

@@ -8,28 +8,11 @@
 #include <pybind11/stl.h>
 #include <pybind11/stl_bind.h>
 
-#include "openvino/core/type.hpp"
 #include "pyopenvino/core/common.hpp"
 #include "pyopenvino/utils/utils.hpp"
 
 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";

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

@@ -6,6 +6,26 @@
 
 #include <pybind11/pybind11.h>
 
+#include "openvino/core/type.hpp"
+#include <string>
 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);

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

@@ -439,6 +439,7 @@ void regclass_graph_Node(py::module m) {
     node.def_property_readonly("type_info", &ov::Node::get_type_info);
     node.def_property("friendly_name", &ov::Node::get_friendly_name, &ov::Node::set_friendly_name);
 
+    node.def("visit_attributes", &ov::Node::visit_attributes);
     node.def("get_attributes", [](const std::shared_ptr<ov::Node>& self) {
         util::DictAttributeSerializer dict_serializer(self);
         return dict_serializer.get_attributes();

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

@@ -16,52 +16,70 @@
 
 namespace py = pybind11;
 
-/// Trampoline class to support inheritance 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 PyOp::validate_and_infer_types() {
+    PYBIND11_OVERRIDE(void, ov::op::Op, validate_and_infer_types);
+}
 
-    void validate_and_infer_types() override {
-        PYBIND11_OVERRIDE(void, ov::op::Op, validate_and_infer_types);
+bool PyOp::visit_attributes(ov::AttributeVisitor& value) {
+    py::gil_scoped_acquire gil;  // Acquire the GIL while in this scope.
+    // Try to look up the overridden method on the Python side.
+    py::function overrided_py_method = pybind11::get_override(this, "visit_attributes");
+    if (overrided_py_method) {                                       // method is found
+        return static_cast<py::bool_>(overrided_py_method(&value));  // Call the Python function.
     }
+    return true;
+}
 
-    bool visit_attributes(ov::AttributeVisitor& value) override {
-        pybind11::gil_scoped_acquire gil;  // Acquire the GIL while in this scope.
-        // Try to look up the overridden method on the Python side.
-        pybind11::function overrided_py_method = pybind11::get_override(this, "visit_attributes");
-        if (overrided_py_method) {                                       // method is found
-            return static_cast<py::bool_>(overrided_py_method(&value));  // Call the Python function.
-        }
-        return false;
+std::shared_ptr<ov::Node> PyOp::clone_with_new_inputs(const ov::OutputVector& new_args) const {
+    py::gil_scoped_acquire gil;  // Acquire the GIL while in this scope.
+    // Try to look up the overridden method on the Python side.
+    py::function overrided_py_method = pybind11::get_override(this, "clone_with_new_inputs");
+    if (overrided_py_method) {                        // method is found
+        auto result = overrided_py_method(new_args);  // Call the Python function.
+        return result.cast<std::shared_ptr<ov::Node>>();
     }
+    // Default implementation for clone_with_new_inputs
+    auto py_handle_type = py_handle.get_type();
+    auto new_py_object = py_handle_type(new_args);
+    return new_py_object.cast<std::shared_ptr<ov::Node>>();
+}
 
-    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 ov::op::Op::type_info_t& PyOp::get_type_info() const {
+    return *m_type_info;
+}
 
-    const type_info_t& get_type_info() const override {
-        PYBIND11_OVERRIDE(const ov::Node::type_info_t&, ov::op::Op, get_type_info);
-    }
+bool PyOp::evaluate(ov::TensorVector& output_values, const ov::TensorVector& input_values) const {
+    PYBIND11_OVERRIDE(bool, ov::op::Op, evaluate, output_values, input_values);
+}
 
-    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 PyOp::has_evaluate() const {
+    PYBIND11_OVERRIDE(bool, ov::op::Op, has_evaluate);
+}
 
-    bool has_evaluate() const override {
-        PYBIND11_OVERRIDE(bool, ov::op::Op, has_evaluate);
+void PyOp::update_type_info() {
+    py::gil_scoped_acquire gil;  // Acquire the GIL while in this scope.
+
+    // Try to look up the overridden method on the Python side.
+    py::function overriden_py_method = pybind11::get_override(this, "get_type_info");
+    if (overriden_py_method) {
+        const auto type_info_from_py = overriden_py_method();
+        if (!py::isinstance<ov::DiscreteTypeInfo>(type_info_from_py)) {
+            OPENVINO_THROW("operation type_info must be an instance of DiscreteTypeInfo, but ",
+                           py::str(py::type::of(type_info_from_py)),
+                           " is passed.");
+        }
+        m_type_info = type_info_from_py.cast<std::shared_ptr<ov::DiscreteTypeInfo>>();
     }
-
-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> op(m, "Op");
 
     op.def(py::init([](const py::object& py_obj) {
         return PyOp(py_obj);
     }));
+
+    op.def("_update_type_info", [](PyOp& self) {
+        self.update_type_info();
+    });
 }

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

@@ -6,6 +6,43 @@
 
 #include <pybind11/pybind11.h>
 
+#include "openvino/op/op.hpp"
+#include "openvino/core/node.hpp"
+#include "pyopenvino/graph/discrete_type_info.hpp"
+
 namespace py = pybind11;
 
+/// Trampoline class to support inheritance from TorchDecoder in Python
+class PyOp : public ov::op::Op {
+public:
+    using ov::op::Op::Op;
+
+    PyOp() = default;
+
+    // Keeps a reference to the Python object to manage its lifetime
+    PyOp(const py::object& py_obj) : py_handle(py_obj) {
+        // Set default value for DiscreteTypeInfo
+        const auto py_class_name = py_handle.get_type().attr("__name__").cast<std::string>();
+        m_type_info = std::make_shared<DiscreteTypeInfoWrapper>(py_class_name, "extension");
+    }
+
+    void update_type_info();
+
+    void validate_and_infer_types() override;
+
+    bool visit_attributes(ov::AttributeVisitor& value) override;
+
+    std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
+
+    const type_info_t& get_type_info() const override;
+
+    bool evaluate(ov::TensorVector& output_values, const ov::TensorVector& input_values) const override;
+
+    bool has_evaluate() const override;
+
+private:
+    py::object py_handle;  // Holds the Python object to manage its lifetime
+    std::shared_ptr<ov::DiscreteTypeInfo> m_type_info;
+};
+
 void regclass_graph_Op(py::module m);

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

@@ -0,0 +1,51 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/core/op_extension.hpp"
+
+#include <pybind11/pybind11.h>
+
+#include <pyopenvino/graph/op_extension.hpp>
+
+#include "pyopenvino/core/common.hpp"
+#include "pyopenvino/core/extension.hpp"
+#include "pyopenvino/graph/discrete_type_info.hpp"
+#include "pyopenvino/graph/node_output.hpp"
+#include "pyopenvino/graph/op.hpp"
+
+namespace py = pybind11;
+
+const ov::DiscreteTypeInfo& PyOpExtension::get_type_info() const {
+    return *m_type_info;
+}
+
+ov::OutputVector PyOpExtension::create(const ov::OutputVector& inputs, ov::AttributeVisitor& visitor) const {
+    py::gil_scoped_acquire acquire;
+
+    const auto node = py_handle_dtype();
+
+    node.attr("set_arguments")(py::cast(inputs));
+    if (node.attr("visit_attributes")(&visitor)) {
+        node.attr("constructor_validate_and_infer_types")();
+    }
+
+    return py::cast<ov::OutputVector>(node.attr("outputs")());
+}
+
+std::vector<ov::Extension::Ptr> PyOpExtension::get_attached_extensions() const {
+    return {};
+}
+
+void regclass_graph_OpExtension(py::module m) {
+    py::class_<PyOpExtension, std::shared_ptr<PyOpExtension>, ov::Extension> op_extension(m, "OpExtension");
+    op_extension.doc() = "openvino.OpExtension provides the base interface for OpenVINO extensions.";
+
+    op_extension.def("__repr__", [](const PyOpExtension& self) {
+        return Common::get_simple_repr(self);
+    });
+
+    op_extension.def(py::init([](py::object dtype) {
+        return PyOpExtension(dtype);
+    }));
+}