[TF FE] Support mutable tensors produced by variables and resources

src/frontends/tensorflow/docs/supported_ops.md

@@ -68,10 +68,10 @@ A "supported operation" is one that TensorFlow Frontend can convert to the OpenV
 | AssertNextDataset                                       | NO                            |                               |
 | AssertPrevDataset                                       | NO                            |                               |
 | Assign                                                  | YES                           |                               |
-| AssignAdd                                               | NO                            |                               |
-| AssignAddVariableOp                                     | NO                            |                               |
-| AssignSub                                               | NO                            |                               |
-| AssignSubVariableOp                                     | NO                            |                               |
+| AssignAdd                                               | YES                           |                               |
+| AssignAddVariableOp                                     | YES                           |                               |
+| AssignSub                                               | YES                           |                               |
+| AssignSubVariableOp                                     | YES                           |                               |
 | AssignVariableOp                                        | YES                           |                               |
 | AssignVariableXlaConcatND                               | NO                            |                               |
 | Atan                                                    | YES                           |                               |

src/frontends/tensorflow/include/openvino/frontend/tensorflow/node_context.hpp

@@ -8,6 +8,7 @@
 #include "exception.hpp"
 #include "openvino/core/any.hpp"
 #include "openvino/frontend/node_context.hpp"
+#include "variable.hpp"
 
 namespace ov {
 namespace frontend {
@@ -21,19 +22,50 @@ class NodeContext : public ov::frontend::NodeContext {
     using Ptr = std::shared_ptr<NodeContext>;
     NodeContext(const std::shared_ptr<DecoderBase>& decoder,
                 const OutputVector& inputs,
+                const VariableMap::Ptr& ov_variable_state_map = nullptr,
                 TranslateSession* translate_session = nullptr)
         : ov::frontend::NodeContext(decoder->get_op_type()),
           m_decoder(decoder),
           m_translate_session(translate_session),
-          m_inputs(inputs) {}
+          m_inputs(inputs),
+          m_ov_variable_state_map(ov_variable_state_map) {}
 
     /// Detects if there is at least one input attached with a given name
     bool has_input(const size_t& port_index) const {
         return port_index < m_inputs.size();
     }
 
+    /// Retrieve the input by reference. It is needed for operations working with reference inputs
+    Output<Node> get_input_by_reference(int port_index) const {
+        auto input = m_inputs.at(port_index);
+        auto ref_node = ov::as_type_ptr<Variable>(input.get_node_shared_ptr());
+        if (ref_node) {
+            auto ref_node_name = ref_node->get_name();
+            Variable::Ptr found_var = nullptr;
+            // get the latest state of Variable with the required name
+            if (m_ov_variable_state_map->get_variable_state(get_name(), ref_node_name, found_var)) {
+                // no need to extract variable value and return variable output
+                return found_var->output(0);
+            }
+        }
+        return input;
+    }
+
+    /// Retrieve the input by value. If variable comes to the input, it tries to extracts its value at this moment
     Output<Node> get_input(int port_index) const override {
-        return m_inputs.at(port_index);
+        auto input = m_inputs.at(port_index);
+        auto ref_node = ov::as_type_ptr<Variable>(input.get_node_shared_ptr());
+        if (ref_node) {
+            auto ref_node_name = ref_node->get_name();
+            Variable::Ptr found_var = nullptr;
+            // get the latest state of Variable with the required name
+            if (m_ov_variable_state_map->get_variable_state(get_name(), ref_node_name, found_var)) {
+                // resolve Variable node by extracting its value from variable state map
+                return found_var->get_value();
+            }
+        }
+
+        return input;
     }
 
     size_t get_input_size() const override {
@@ -60,12 +92,19 @@ class NodeContext : public ov::frontend::NodeContext {
         return m_translate_session;
     }
 
+    VariableMap::Ptr get_variable_state_map() const {
+        return m_ov_variable_state_map;
+    }
+
 private:
     ov::Any apply_additional_conversion_rules(const ov::Any& data, const std::type_info& type_info) const override;
 
     std::shared_ptr<DecoderBase> m_decoder;
     TranslateSession* m_translate_session;
     const OutputVector& m_inputs;
+
+    // it contains variables states at each node of the graph
+    VariableMap::Ptr m_ov_variable_state_map;
 };
 
 using CreatorFunctionIndexed = std::function<ov::OutputVector(const ov::frontend::tensorflow::NodeContext&)>;

src/frontends/tensorflow/include/openvino/frontend/tensorflow/variable.hpp

@@ -0,0 +1,165 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/util/framework_node.hpp"
+
+namespace ov {
+namespace frontend {
+namespace tensorflow {
+
+class Variable : public ov::op::util::FrameworkNode {
+public:
+    using Ptr = std::shared_ptr<Variable>;
+    OPENVINO_OP("TFVariable", "ov::frontend::tensorflow", ::ov::op::util::FrameworkNode);
+
+    Variable(const std::string& name,
+             const ov::Shape& shape,
+             const ov::element::Type& type,
+             const std::shared_ptr<DecoderBase>& decoder)
+        : ov::op::util::FrameworkNode(ov::OutputVector{}, 1),
+          m_name(name),
+          m_shape(shape),
+          m_type(type),
+          m_decoder(decoder),
+          m_is_initialized(false),
+          m_init_counter(0) {
+        validate_and_infer_types();
+    }
+
+    Variable(const std::string& name,
+             const ov::Shape& shape,
+             const ov::element::Type& type,
+             const ov::Output<ov::Node>& value,
+             const std::shared_ptr<DecoderBase>& decoder)
+        : Variable(name, shape, type, decoder) {
+        m_value = value;
+        // reset names of tensor corresponding to variable value
+        // that is because variable can have multiple values during inference
+        m_value.set_names({});
+        m_is_initialized = true;
+        ++m_init_counter;
+    }
+
+    Variable(const Variable& other, const ov::Output<ov::Node>& value) : Variable(other) {
+        m_value = value;
+        // reset names of tensor corresponding to variable value
+        // that is because variable can have multiple values during inference
+        m_value.set_names({});
+        m_is_initialized = true;
+        ++m_init_counter;
+    }
+
+    void validate_and_infer_types() override {
+        set_output_type(0, m_type, m_shape);
+    }
+
+    bool is_initialized() const {
+        return m_is_initialized;
+    }
+
+    ov::Output<ov::Node> get_value() const {
+        FRONT_END_GENERAL_CHECK(
+            m_is_initialized,
+            "[TensorFlow Frontend] internal error: get_value() is called for uninitialized variable");
+        return m_value;
+    }
+
+    std::string get_name() const {
+        return m_name;
+    }
+
+    uint64_t get_init_counter() const {
+        return m_init_counter;
+    }
+
+private:
+    std::string m_name;
+    ov::Shape m_shape;
+    ov::element::Type m_type;
+    std::shared_ptr<DecoderBase> m_decoder;
+    bool m_is_initialized;
+    ov::Output<ov::Node> m_value;
+    // this member is used to select the latest state of Variable
+    uint64_t m_init_counter;
+};
+
+// a container of Variables state for each operation node in a graph
+class VariableMap {
+public:
+    using Ptr = std::shared_ptr<VariableMap>;
+    bool get_variable_state(const std::string& node_name,
+                            const std::string& variable_name,
+                            Variable::Ptr& found_variable) const {
+        if (m_variables_state.count(node_name) > 0) {
+            for (const auto& variable : m_variables_state.at(node_name)) {
+                if (variable && variable->get_name() == variable_name && variable->is_initialized()) {
+                    found_variable = variable;
+                    return true;
+                }
+            }
+        } else {
+            return false;
+        }
+        return false;
+    }
+
+    void initialize_variable_state_map_for_node(const std::vector<std::string>& control_dependencies,
+                                                const std::vector<std::string>& data_dependencies,
+                                                const std::string& node_name) {
+        m_variables_state[node_name] = {};
+        for (const auto& dependency : control_dependencies) {
+            for (const auto& dependency_variable : m_variables_state[dependency]) {
+                update_variable_state_map_for_node(node_name, dependency_variable);
+            }
+        }
+
+        for (const auto& dependency : data_dependencies) {
+            for (const auto& dependency_variable : m_variables_state[dependency]) {
+                update_variable_state_map_for_node(node_name, dependency_variable);
+            }
+        }
+    }
+
+    void update_variable_state_map_for_node(const std::string& node_name, const Variable::Ptr& update_variable) {
+        FRONT_END_GENERAL_CHECK(
+            update_variable && update_variable->is_initialized(),
+            "[TensorFlow Frontend] internal error: variable maps must be updated with initialized variable");
+        auto variable_name = update_variable->get_name();
+
+        size_t remove_ind = 0;
+        bool remove_old_variable = false;
+        bool found_better = false;
+        // remove old variable state if exists
+        for (size_t ind = 0; ind < m_variables_state[node_name].size(); ++ind) {
+            auto checked_variable = m_variables_state[node_name][ind];
+            if (checked_variable->get_name() == variable_name && checked_variable->is_initialized() &&
+                checked_variable->get_init_counter() < update_variable->get_init_counter()) {
+                remove_ind = ind;
+                remove_old_variable = true;
+                break;
+            } else if (checked_variable->get_name() == variable_name && checked_variable->is_initialized() &&
+                       checked_variable->get_init_counter() >= update_variable->get_init_counter()) {
+                found_better = true;
+            }
+        }
+
+        if (remove_old_variable) {
+            // update the variable map with new variable
+            m_variables_state[node_name].erase(m_variables_state[node_name].begin() + remove_ind);
+        }
+
+        if (!found_better) {
+            m_variables_state[node_name].push_back(update_variable);
+        }
+    }
+
+private:
+    std::unordered_map<std::string, std::vector<Variable::Ptr>> m_variables_state;
+};
+
+}  // namespace tensorflow
+}  // namespace frontend
+}  // namespace ov

src/frontends/tensorflow/src/op/assign.cpp

@@ -0,0 +1,44 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "common_op_table.hpp"
+#include "openvino/frontend/tensorflow/node_context.hpp"
+#include "openvino/frontend/tensorflow/variable.hpp"
+
+using namespace std;
+using namespace ov;
+using namespace ov::op;
+using namespace ov::frontend::tensorflow;
+
+namespace ov {
+namespace frontend {
+namespace tensorflow {
+namespace op {
+OutputVector translate_assign_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"Assign"});
+    auto ref = as_type_ptr<Variable>(node.get_input_by_reference(0).get_node_shared_ptr());
+    TENSORFLOW_OP_VALIDATION(
+        node,
+        ref,
+        "[TensorFlow Frontend] internal error: Assign operation expects variable by the first input");
+
+    // create new variable operation node that will represent the same variable but it will be initialized with value
+    auto value = node.get_input(1);
+    auto new_ref = make_shared<Variable>(*ref, value);
+
+    // since this operation produces new state of the variable
+    // it needs to update variables map
+    auto variables_state_map = node.get_variable_state_map();
+    TENSORFLOW_OP_VALIDATION(node,
+                             variables_state_map,
+                             "[TensorFlow Frontend] internal error: variable state map is nullptr");
+    variables_state_map->update_variable_state_map_for_node(node.get_name(), new_ref);
+
+    return {new_ref};
+}
+
+}  // namespace op
+}  // namespace tensorflow
+}  // namespace frontend
+}  // namespace ov

src/frontends/tensorflow/src/op/assign_add.cpp

@@ -0,0 +1,47 @@
+// Copyright (C) 2018-2024 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "common_op_table.hpp"
+#include "openvino/frontend/tensorflow/node_context.hpp"
+#include "openvino/frontend/tensorflow/variable.hpp"
+#include "openvino/op/add.hpp"
+
+using namespace std;
+using namespace ov;
+using namespace ov::op;
+using namespace ov::frontend::tensorflow;
+
+namespace ov {
+namespace frontend {
+namespace tensorflow {
+namespace op {
+
+OutputVector translate_assign_add_op(const NodeContext& node) {
+    default_op_checks(node, 2, {"AssignAdd"});
+    auto ref = as_type_ptr<Variable>(node.get_input_by_reference(0).get_node_shared_ptr());
+    TENSORFLOW_OP_VALIDATION(
+        node,
+        ref,
+        "[TensorFlow Frontend] internal error: AssignAdd operation expects variable by the first input");
+
+    // create new variable operation node that will represent the same variable but it will be initialized with value
+    auto value_to_add = node.get_input(1);
+    auto current_value = ref->get_value();
+    auto new_value = make_shared<v1::Add>(current_value, value_to_add);
+    auto new_ref = make_shared<Variable>(*ref, new_value);
+
+    // since this operation produces new state of the variable
+    // it needs to update variables map
+    auto variables_state_map = node.get_variable_state_map();
+    TENSORFLOW_OP_VALIDATION(node,
+                             variables_state_map,
+                             "[TensorFlow Frontend] internal error: variable state map is nullptr");
+    variables_state_map->update_variable_state_map_for_node(node.get_name(), new_ref);
+    return {new_ref};
+}
+
+}  // namespace op
+}  // namespace tensorflow
+}  // namespace frontend
+}  // namespace ov