[PYTHON] Expose layout helpers (#8507)

* create layout_helpers

* @pszmel

* fix typo

* Create tests

* style fix

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

@@ -51,3 +51,4 @@ from openvino.pyopenvino import Output
 from openvino.pyopenvino import Layout
 from openvino.pyopenvino import ConstOutput
 from openvino.pyopenvino import util
+from openvino.pyopenvino import layout_helpers

runtime/bindings/python/src/pyopenvino/graph/layout_helpers.cpp

@@ -0,0 +1,26 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "pyopenvino/graph/layout_helpers.hpp"
+
+#include <pybind11/pybind11.h>
+
+#include "openvino/core/layout.hpp"
+
+namespace py = pybind11;
+
+void regmodule_graph_layout_helpers(py::module m) {
+    py::module mod = m.def_submodule("layout_helpers", "openvino.impl.layout_helpers");
+
+    mod.def("has_batch", &ov::layout::has_batch, py::arg("layout"));
+    mod.def("batch_idx", &ov::layout::batch_idx, py::arg("layout"));
+    mod.def("has_channels", &ov::layout::has_channels, py::arg("layout"));
+    mod.def("channels_idx", &ov::layout::channels_idx, py::arg("layout"));
+    mod.def("has_depth", &ov::layout::has_depth, py::arg("layout"));
+    mod.def("depth_idx", &ov::layout::depth_idx, py::arg("layout"));
+    mod.def("has_height", &ov::layout::has_height, py::arg("layout"));
+    mod.def("height_idx", &ov::layout::height_idx, py::arg("layout"));
+    mod.def("has_width", &ov::layout::has_width, py::arg("layout"));
+    mod.def("width_idx", &ov::layout::width_idx, py::arg("layout"));
+}

runtime/bindings/python/src/pyopenvino/graph/layout_helpers.hpp

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

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

@@ -31,6 +31,7 @@
 #include "pyopenvino/core/version.hpp"
 #include "pyopenvino/graph/dimension.hpp"
 #include "pyopenvino/graph/layout.hpp"
+#include "pyopenvino/graph/layout_helpers.hpp"
 #include "pyopenvino/graph/ops/constant.hpp"
 #include "pyopenvino/graph/ops/parameter.hpp"
 #include "pyopenvino/graph/ops/result.hpp"
@@ -76,6 +77,7 @@ PYBIND11_MODULE(pyopenvino, m) {
     regclass_graph_op_Constant(m_op);
     regclass_graph_op_Parameter(m_op);
     regclass_graph_op_Result(m_op);
+
 #if defined(NGRAPH_ONNX_FRONTEND_ENABLE)
     regmodule_graph_onnx_import(m);
 #endif
@@ -83,6 +85,7 @@ PYBIND11_MODULE(pyopenvino, m) {
     regclass_graph_Function(m);
     regmodule_graph_passes(m);
     regmodule_graph_util(m);
+    regmodule_graph_layout_helpers(m);
     regclass_graph_Variant(m);
     regclass_graph_VariantWrapper<std::string>(m, std::string("String"));
     regclass_graph_VariantWrapper<int64_t>(m, std::string("Int"));

runtime/bindings/python/tests/test_ngraph/test_basic.py

@@ -12,7 +12,7 @@ import openvino as ov
 from openvino.pyopenvino import VariantInt, VariantString
 
 from openvino.exceptions import UserInputError
-from openvino.impl import Function, PartialShape, Shape, Type
+from openvino.impl import Function, PartialShape, Shape, Type, layout_helpers
 from openvino.impl.op import Parameter
 from tests.runtime import get_runtime
 from tests.test_ngraph.util import run_op_node
@@ -538,3 +538,56 @@ def test_layout():
     layout = ov.Layout("N...C")
     assert layout == "N...C"
     assert layout != "NC?"
+
+
+def test_layout_helpers():
+    layout = ov.Layout("NCHWD")
+    assert(layout_helpers.has_batch(layout))
+    assert(layout_helpers.has_channels(layout))
+    assert(layout_helpers.has_depth(layout))
+    assert(layout_helpers.has_height(layout))
+    assert(layout_helpers.has_width(layout))
+
+    assert layout_helpers.batch_idx(layout) == 0
+    assert layout_helpers.channels_idx(layout) == 1
+    assert layout_helpers.height_idx(layout) == 2
+    assert layout_helpers.width_idx(layout) == 3
+    assert layout_helpers.depth_idx(layout) == 4
+
+    layout = ov.Layout("N...C")
+    assert(layout_helpers.has_batch(layout))
+    assert(layout_helpers.has_channels(layout))
+    assert not(layout_helpers.has_depth(layout))
+    assert not(layout_helpers.has_height(layout))
+    assert not (layout_helpers.has_width(layout))
+
+    assert layout_helpers.batch_idx(layout) == 0
+    assert layout_helpers.channels_idx(layout) == -1
+
+    with pytest.raises(RuntimeError):
+        layout_helpers.height_idx(layout)
+
+    with pytest.raises(RuntimeError):
+        layout_helpers.width_idx(layout)
+
+    with pytest.raises(RuntimeError):
+        layout_helpers.depth_idx(layout)
+
+    layout = ov.Layout("NC?")
+    assert(layout_helpers.has_batch(layout))
+    assert(layout_helpers.has_channels(layout))
+    assert not(layout_helpers.has_depth(layout))
+    assert not(layout_helpers.has_height(layout))
+    assert not (layout_helpers.has_width(layout))
+
+    assert layout_helpers.batch_idx(layout) == 0
+    assert layout_helpers.channels_idx(layout) == 1
+
+    with pytest.raises(RuntimeError):
+        layout_helpers.height_idx(layout)
+
+    with pytest.raises(RuntimeError):
+        layout_helpers.width_idx(layout)
+
+    with pytest.raises(RuntimeError):
+        layout_helpers.depth_idx(layout)