77 lines
3.3 KiB
Python
77 lines
3.3 KiB
Python
# 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 TestGather(CommonTFLayerTest):
|
|
def _prepare_input(self, inputs_info):
|
|
assert 'params' in inputs_info
|
|
assert 'indices' in inputs_info
|
|
params_shape = inputs_info['params']
|
|
indices_shape = inputs_info['indices']
|
|
inputs_data = {}
|
|
inputs_data['params'] = np.random.randint(-50, 50, params_shape).astype(self.params_type)
|
|
inputs_data['indices'] = np.random.randint(0, self.max_index, indices_shape).astype(self.indices_type)
|
|
return inputs_data
|
|
|
|
def create_gather_net(self, params_shape, params_type, indices_shape, indices_type, axis_value, batch_dims,
|
|
operation_type):
|
|
self.params_type = params_type
|
|
self.indices_type = indices_type
|
|
if batch_dims is None:
|
|
batch_dims = 0
|
|
if axis_value is None:
|
|
axis_value = 0
|
|
axis_norm = axis_value
|
|
if axis_norm < 0:
|
|
axis_norm += len(params_shape)
|
|
assert 0 <= axis_norm < len(params_shape), "Incorrect `axis` value for the test case"
|
|
self.max_index = params_shape[axis_norm]
|
|
|
|
tf.compat.v1.reset_default_graph()
|
|
with tf.compat.v1.Session() as sess:
|
|
params = tf.compat.v1.placeholder(params_type, params_shape, 'params')
|
|
indices = tf.compat.v1.placeholder(indices_type, indices_shape, 'indices')
|
|
if operation_type == "Gather":
|
|
tf.raw_ops.Gather(params=params, indices=indices)
|
|
elif operation_type == "GatherV2":
|
|
axis = tf.constant(axis_value, dtype=tf.int32)
|
|
tf.raw_ops.GatherV2(params=params, indices=indices, axis=axis, batch_dims=batch_dims)
|
|
else:
|
|
assert False, "Incorrect operation type is tested"
|
|
|
|
tf.compat.v1.global_variables_initializer()
|
|
tf_net = sess.graph_def
|
|
|
|
ref_net = None
|
|
|
|
return tf_net, ref_net
|
|
|
|
test_data_precommit = [
|
|
dict(params_shape=[4, 6], params_type=np.float32, indices_shape=[], indices_type=np.int32,
|
|
axis_value=None, batch_dims=None,
|
|
operation_type="Gather"),
|
|
dict(params_shape=[3, 4, 6], params_type=np.float32, indices_shape=[3, 4], indices_type=np.int32,
|
|
axis_value=None, batch_dims=None,
|
|
operation_type="Gather"),
|
|
dict(params_shape=[5, 4, 3], params_type=np.int32, indices_shape=[5, 2, 1], indices_type=np.int64,
|
|
axis_value=2, batch_dims=1,
|
|
operation_type="GatherV2"),
|
|
dict(params_shape=[3, 2, 6, 4], params_type=np.float32, indices_shape=[3, 2, 1, 3], indices_type=np.int32,
|
|
axis_value=-1, batch_dims=-2,
|
|
operation_type="GatherV2"),
|
|
]
|
|
|
|
@pytest.mark.parametrize("params", test_data_precommit)
|
|
@pytest.mark.precommit_tf_fe
|
|
@pytest.mark.nightly
|
|
def test_gather(self, params, ie_device, precision, ir_version, temp_dir, use_new_frontend,
|
|
use_old_api):
|
|
self._test(*self.create_gather_net(**params),
|
|
ie_device, precision, ir_version, temp_dir=temp_dir,
|
|
use_new_frontend=use_new_frontend, use_old_api=use_old_api)
|