diff --git a/src/frontends/tensorflow_common/src/op/pack.cpp b/src/frontends/tensorflow_common/src/op/pack.cpp index d4502ad75c4..51773bf4725 100644 --- a/src/frontends/tensorflow_common/src/op/pack.cpp +++ b/src/frontends/tensorflow_common/src/op/pack.cpp @@ -14,18 +14,21 @@ namespace tensorflow { namespace op { OutputVector translate_pack_op(const NodeContext& node) { - auto axis = node.get_attribute("axis"); + default_op_checks(node, 1, {"Pack", "PACK"}); + auto num_size = static_cast(node.get_input_size()); + + auto axis = node.get_attribute("axis", 0); auto axis_const = make_shared(element::i64, Shape{}, axis); OutputVector concat_inputs; - for (size_t i = 0; i < node.get_input_size(); ++i) { - auto in = node.get_input(static_cast(i)); + for (int ind = 0; ind < num_size; ++ind) { + auto in = node.get_input(ind); concat_inputs.push_back(make_shared(in, axis_const)); } - auto res = make_shared(concat_inputs, axis); - set_node_name(node.get_name(), res); - return res->outputs(); + auto pack = make_shared(concat_inputs, axis); + set_node_name(node.get_name(), pack); + return {pack}; } } // namespace op } // namespace tensorflow diff --git a/tests/layer_tests/tensorflow_tests/test_tf_Pack.py b/tests/layer_tests/tensorflow_tests/test_tf_Pack.py new file mode 100644 index 00000000000..cfee54616fc --- /dev/null +++ b/tests/layer_tests/tensorflow_tests/test_tf_Pack.py @@ -0,0 +1,53 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +import tensorflow as tf +from common.tf_layer_test_class import CommonTFLayerTest + + +class TestPack(CommonTFLayerTest): + def _prepare_input(self, inputs_info): + inputs_data = {} + for input_name, input_shape in inputs_info.items(): + inputs_data[input_name] = np.random.randint(-5, 5, input_shape).astype(self.input_type) + return inputs_data + + def create_pack_net(self, input_shape, input_num, axis, input_type): + self.input_type = input_type + tf.compat.v1.reset_default_graph() + # Create the graph and model + with tf.compat.v1.Session() as sess: + inputs = [] + type_map = { + np.float32: tf.float32, + np.int32: tf.int32, + } + assert input_type in type_map, "Test error: need to update type_map" + tf_type = type_map[input_type] + for ind in range(input_num): + inputs.append(tf.compat.v1.placeholder(tf_type, input_shape, 'input' + str(ind))) + if axis is not None: + tf.raw_ops.Pack(values=inputs, axis=axis) + else: + tf.raw_ops.Pack(values=inputs) + tf.compat.v1.global_variables_initializer() + + tf_net = sess.graph_def + + return tf_net, None + + test_data_basic = [ + dict(input_shape=[2, 4], input_num=2, axis=None, input_type=np.float32), + dict(input_shape=[3, 1, 2], input_num=3, axis=1, input_type=np.int32), + ] + + @pytest.mark.parametrize("params", test_data_basic) + @pytest.mark.precommit_tf_fe + @pytest.mark.nightly + def test_pack_basic(self, params, ie_device, precision, ir_version, temp_dir, + use_new_frontend, use_old_api): + self._test(*self.create_pack_net(**params), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_new_frontend=use_new_frontend, use_old_api=use_old_api)