From 56ada41bd05f47ceeddd86f1e99f61858e26a2ba Mon Sep 17 00:00:00 2001
From: Mikhail Treskin <mikhail.treskin@intel.com>
Date: Mon, 7 Jun 2021 18:19:52 +0300
Subject: [PATCH] Add get_constant_from_source capability to pyngraph (#6018)

---
 ngraph/python/src/pyngraph/util.cpp           | 19 +++++++++++++
 ngraph/python/tests/test_ngraph/test_utils.py | 28 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 ngraph/python/tests/test_ngraph/test_utils.py

diff --git a/ngraph/python/src/pyngraph/util.cpp b/ngraph/python/src/pyngraph/util.cpp
index a2eac625bab..5178e84fe90 100644
--- a/ngraph/python/src/pyngraph/util.cpp
+++ b/ngraph/python/src/pyngraph/util.cpp
@@ -5,6 +5,7 @@
 #include <pybind11/numpy.h>
 
 #include "pyngraph/util.hpp"
+#include "ngraph/validation_util.hpp"
 
 namespace py = pybind11;
 
@@ -18,4 +19,22 @@ void regmodule_pyngraph_util(py::module m)
 {
     py::module mod = m.def_submodule("util", "ngraph.impl.util");
     mod.def("numpy_to_c", &numpy_to_c);
+    mod.def("get_constant_from_source",
+            &ngraph::get_constant_from_source,
+            py::arg("output"),
+            R"(
+                    Runs an estimation of source tensor.
+
+                    Parameters
+                    ----------
+                    output : Output
+                        output node
+
+                    Returns
+                    ----------
+                    get_constant_from_source : Constant or None
+                        If it succeeded to calculate both bounds and
+                        they are the same returns Constant operation
+                        from the resulting bound, otherwise Null.
+                )");
 }
diff --git a/ngraph/python/tests/test_ngraph/test_utils.py b/ngraph/python/tests/test_ngraph/test_utils.py
new file mode 100644
index 00000000000..49b90017305
--- /dev/null
+++ b/ngraph/python/tests/test_ngraph/test_utils.py
@@ -0,0 +1,28 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import ngraph as ng
+from ngraph.impl import Shape
+
+
+def test_get_constant_from_source_success():
+    dtype = np.int
+    input1 = ng.parameter(Shape([5, 5]), dtype=dtype, name="input_1")
+    input2 = ng.parameter(Shape([25]), dtype=dtype, name="input_2")
+    shape_of = ng.shape_of(input2, name="shape_of")
+    reshape = ng.reshape(input1, shape_of, special_zero=True)
+    folded_const = ng.impl.util.get_constant_from_source(reshape.input(1).get_source_output())
+
+    assert folded_const is not None
+    assert folded_const.get_vector() == [25]
+
+
+def test_get_constant_from_source_failed():
+    dtype = np.int
+    input1 = ng.parameter(Shape([5, 5]), dtype=dtype, name="input_1")
+    input2 = ng.parameter(Shape([1]), dtype=dtype, name="input_2")
+    reshape = ng.reshape(input1, input2, special_zero=True)
+    folded_const = ng.impl.util.get_constant_from_source(reshape.input(1).get_source_output())
+
+    assert folded_const is None