From 89fe26e3db1f5689e16bd911904334515c56daa0 Mon Sep 17 00:00:00 2001
From: Vitaliy Urusovskij <vitaliy.urusovskij@intel.com>
Date: Tue, 1 Feb 2022 18:46:09 +0300
Subject: [PATCH] Copy RandomUniform `m_state` during clone_with_new_inputs()
 (#10031)

* Copy RandomUniform m_state during clone_with_new_inputs()

* Add `get_state()` for RandomUniform op

* Add copy.random_uniform test
---
 .../include/openvino/op/random_uniform.hpp    |  5 +++
 src/core/src/op/random_uniform.cpp            | 10 +++---
 src/core/tests/copy.cpp                       | 34 +++++++++++++++++++
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/src/core/include/openvino/op/random_uniform.hpp b/src/core/include/openvino/op/random_uniform.hpp
index 1f505282144..533bb230ce2 100644
--- a/src/core/include/openvino/op/random_uniform.hpp
+++ b/src/core/include/openvino/op/random_uniform.hpp
@@ -68,6 +68,11 @@ public:
         m_op_seed = seed2;
     }
 
+    /// \return The state value.
+    std::pair<uint64_t, uint64_t> get_state() const {
+        return m_state;
+    }
+
     OPENVINO_SUPPRESS_DEPRECATED_START
     bool evaluate(const HostTensorVector& outputs, const HostTensorVector& inputs) const override;
     OPENVINO_SUPPRESS_DEPRECATED_END
diff --git a/src/core/src/op/random_uniform.cpp b/src/core/src/op/random_uniform.cpp
index e8d5d8124ca..9f38daed887 100644
--- a/src/core/src/op/random_uniform.cpp
+++ b/src/core/src/op/random_uniform.cpp
@@ -119,12 +119,10 @@ bool op::v8::RandomUniform::visit_attributes(AttributeVisitor& visitor) {
 shared_ptr<Node> op::v8::RandomUniform::clone_with_new_inputs(const OutputVector& new_args) const {
     NGRAPH_OP_SCOPE(v8_RandomUniform_clone_with_new_inputs);
     check_new_args_count(this, new_args);
-    return make_shared<v8::RandomUniform>(new_args[0],
-                                          new_args[1],
-                                          new_args[2],
-                                          m_output_type,
-                                          m_global_seed,
-                                          m_op_seed);
+    auto ru_copy =
+        make_shared<v8::RandomUniform>(new_args[0], new_args[1], new_args[2], m_output_type, m_global_seed, m_op_seed);
+    ru_copy->m_state = this->m_state;
+    return ru_copy;
 }
 
 bool op::v8::RandomUniform::evaluate(const HostTensorVector& outputs, const HostTensorVector& inputs) const {
diff --git a/src/core/tests/copy.cpp b/src/core/tests/copy.cpp
index c16b63d3895..165203dad9e 100644
--- a/src/core/tests/copy.cpp
+++ b/src/core/tests/copy.cpp
@@ -5,9 +5,11 @@
 #include <memory>
 #include <string>
 
+#include "engines_util/execute_tools.hpp"
 #include "gtest/gtest.h"
 #include "ngraph/ngraph.hpp"
 #include "ngraph/opsets/opset5.hpp"
+#include "openvino/opsets/opset8.hpp"
 #include "util/ndarray.hpp"
 #include "util/test_tools.hpp"
 
@@ -402,3 +404,35 @@ TEST(copy, loop) {
     EXPECT_EQ(loop_copy->get_output_shape(1), out1_shape);
     EXPECT_EQ(loop_copy->get_output_shape(2), out2_shape);
 }
+
+TEST(copy, random_uniform) {
+    const auto min_val_param = make_shared<op::Parameter>(element::f32, Shape{1});
+    const auto max_val_param = make_shared<op::Parameter>(element::f32, Shape{1});
+    auto out_shape = make_shared<op::Constant>(element::i64, Shape{3}, std::vector<int64_t>{1, 2, 3});
+    auto ru =
+        std::make_shared<ov::opset8::RandomUniform>(out_shape, min_val_param, max_val_param, element::f32, 150, 10);
+
+    // Call `evaluate` to update m_state
+    ru->evaluate({make_host_tensor<element::i64>(out_shape->get_shape(), out_shape->get_vector<int64_t>()),
+                  make_host_tensor<element::f32>(min_val_param->get_shape(), {0}),
+                  make_host_tensor<element::f32>(max_val_param->get_shape(), {1})},
+                 {make_host_tensor<element::i64>(out_shape->get_shape(), out_shape->get_vector<int64_t>()),
+                  make_host_tensor<element::f32>(min_val_param->get_shape(), {0}),
+                  make_host_tensor<element::f32>(max_val_param->get_shape(), {1})});
+
+    auto out_shape_c = make_shared<op::Constant>(element::i64, Shape{4}, std::vector<int64_t>{4, 3, 2, 1});
+    const auto min_val_param_c = make_shared<op::Parameter>(element::f32, Shape{1});
+    const auto max_val_param_c = make_shared<op::Parameter>(element::f32, Shape{1});
+    OutputVector new_args{out_shape_c, min_val_param_c, max_val_param_c};
+    auto new_ru = ru->clone_with_new_inputs(new_args);
+    auto node_cast = ov::as_type_ptr<ov::opset8::RandomUniform>(new_ru);
+    ASSERT_NE(node_cast, nullptr);
+
+    ASSERT_TRUE(nullptr != new_ru);
+    ASSERT_TRUE(new_args == new_ru->input_values());
+    ASSERT_TRUE(ru->get_out_type() == node_cast->get_out_type());
+    ASSERT_TRUE(out_shape_c->get_shape_val() == node_cast->get_shape());
+    ASSERT_TRUE(ru->get_global_seed() == node_cast->get_global_seed());
+    ASSERT_TRUE(ru->get_op_seed() == node_cast->get_op_seed());
+    ASSERT_TRUE(ru->get_state() == node_cast->get_state());
+}