62 lines
2.8 KiB
Python
62 lines
2.8 KiB
Python
# Copyright (C) 2018-2022 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pytest
|
|
|
|
from common.tf_layer_test_class import CommonTFLayerTest
|
|
|
|
|
|
class TestGather(CommonTFLayerTest):
|
|
|
|
def create_indices_constant(self):
|
|
pass
|
|
|
|
def create_gather_net(self, data_shape, indices, axis, batch_dims, use_new_frontend, **kwargs):
|
|
import tensorflow as tf
|
|
|
|
tf.compat.v1.reset_default_graph()
|
|
|
|
with tf.compat.v1.Session() as sess:
|
|
data = tf.compat.v1.placeholder(tf.float32, data_shape, 'data')
|
|
indices = tf.constant(indices, dtype=tf.int32)
|
|
gather = tf.gather(data, indices, axis=axis, batch_dims=batch_dims, name='gather_output')
|
|
|
|
tf.compat.v1.global_variables_initializer()
|
|
tf_net = sess.graph_def
|
|
|
|
ref_net = None
|
|
|
|
return tf_net, ref_net
|
|
|
|
test_data_precommit = [
|
|
dict(data_shape=[6, 8, 10, 12], indices=[[0, 2, 4], [5, 7, 9]], axis=2, batch_dims=0),
|
|
dict(data_shape=[4, 6, 8, 10, 12], indices=[2, 5], axis=1, batch_dims=0),
|
|
dict(data_shape=[4, 6, 8, 10, 12], indices=[2, 5], axis=-1, batch_dims=0)
|
|
]
|
|
|
|
@pytest.mark.parametrize("params", test_data_precommit)
|
|
@pytest.mark.precommit
|
|
def test_gather(self, params, ie_device, precision, ir_version, temp_dir, use_new_frontend):
|
|
self._test(*self.create_gather_net(**params, ir_version=ir_version, use_new_frontend=use_new_frontend),
|
|
ie_device, precision, ir_version, temp_dir=temp_dir, use_new_frontend=use_new_frontend)
|
|
|
|
test_data_nightly = [
|
|
dict(data_shape=[2, 3], axis=1, indices=[0, 2], batch_dims=0),
|
|
dict(data_shape=[10, 12], axis=0, indices=[3, 6], batch_dims=0),
|
|
dict(data_shape=[10, 12], axis=1, indices=[[0, 1, 3, 4, 5], [6, 7, 9, 10, 11]], batch_dims=0),
|
|
dict(data_shape=[8, 10, 12], axis=0, indices=[3, 6], batch_dims=0),
|
|
dict(data_shape=[8, 10, 12], axis=-1, indices=[5, 8], batch_dims=0),
|
|
dict(data_shape=[6, 8, 10, 12], axis=0, indices=[2, 5], batch_dims=0),
|
|
dict(data_shape=[6, 8, 10, 12], axis=-1, indices=[5, 8], batch_dims=0),
|
|
dict(data_shape=[6, 8, 10, 12], axis=2, indices=[[0, 2, 4], [5, 7, 9]], batch_dims=0),
|
|
dict(data_shape=[2, 14, 10, 12], axis=1, indices=[[0, 1, 3, 4, 5], [6, 7, 9, 10, 11]], batch_dims=1),
|
|
dict(data_shape=[4, 6, 8, 10, 12], axis=0, indices=[1, 3], batch_dims=0),
|
|
dict(data_shape=[4, 6, 8, 10, 12], axis=-1, indices=[5, 8], batch_dims=0),
|
|
]
|
|
|
|
@pytest.mark.parametrize("params", test_data_nightly)
|
|
@pytest.mark.nightly
|
|
def test_gather_nightly(self, params, ie_device, precision, ir_version, temp_dir, use_new_frontend):
|
|
self._test(*self.create_gather_net(**params, use_new_frontend=use_new_frontend),
|
|
ie_device, precision, ir_version, temp_dir=temp_dir, use_new_frontend=use_new_frontend)
|