• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Test configs for gather_nd."""
16import tensorflow.compat.v1 as tf
17from tensorflow.lite.testing.zip_test_utils import create_tensor_data
18from tensorflow.lite.testing.zip_test_utils import make_zip_of_tests
19from tensorflow.lite.testing.zip_test_utils import register_make_test_function
20
21
22@register_make_test_function()
23def make_gather_nd_tests(options):
24  """Make a set of tests to do gather_nd."""
25
26  test_parameters = [
27      {
28          "params_dtype": [tf.float32, tf.int32, tf.int64, tf.string],
29          "params_shape": [[5, 1]],
30          "indices_dtype": [tf.int32, tf.int64],
31          "indices_shape": [[1, 1]],
32      },
33      {
34          "params_dtype": [tf.float32, tf.int32, tf.int64, tf.string],
35          "params_shape": [[5, 5]],
36          "indices_dtype": [tf.int32, tf.int64],
37          "indices_shape": [[2, 1], [2, 2]],
38      },
39      {
40          "params_dtype": [tf.float32, tf.int32, tf.int64, tf.string],
41          "params_shape": [[5, 5, 10]],
42          "indices_dtype": [tf.int32, tf.int64],
43          "indices_shape": [[3, 1], [2, 2], [2, 3], [2, 1, 3]],
44      },
45      {
46          "params_dtype": [tf.float32, tf.string],
47          "params_shape": [[1, 0]],
48          "indices_dtype": [tf.int64],
49          "indices_shape": [[0, 2]],
50      },
51  ]
52
53  def build_graph(parameters):
54    """Build the gather_nd op testing graph."""
55    params = tf.compat.v1.placeholder(
56        dtype=parameters["params_dtype"],
57        name="params",
58        shape=parameters["params_shape"])
59    indices = tf.compat.v1.placeholder(
60        dtype=parameters["indices_dtype"],
61        name="indices",
62        shape=parameters["indices_shape"])
63    out = tf.gather_nd(params, indices)
64    return [params, indices], [out]
65
66  def build_inputs(parameters, sess, inputs, outputs):
67    params = create_tensor_data(parameters["params_dtype"],
68                                parameters["params_shape"])
69    indices = create_tensor_data(parameters["indices_dtype"],
70                                 parameters["indices_shape"], 0,
71                                 parameters["params_shape"][0] - 1)
72    return [params, indices], sess.run(
73        outputs, feed_dict=dict(zip(inputs, [params, indices])))
74
75  make_zip_of_tests(options, test_parameters, build_graph, build_inputs)
76