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 arg_min_max.""" 16from __future__ import absolute_import 17from __future__ import division 18from __future__ import print_function 19 20import random 21 22import tensorflow.compat.v1 as tf 23from tensorflow.lite.testing.zip_test_utils import create_tensor_data 24from tensorflow.lite.testing.zip_test_utils import make_zip_of_tests 25from tensorflow.lite.testing.zip_test_utils import register_make_test_function 26 27 28@register_make_test_function() 29def make_arg_min_max_tests(options): 30 """Make a set of tests to do arg_max.""" 31 32 test_parameters = [{ 33 "input_dtype": [tf.float32, tf.int32], 34 "input_shape": [[], [1, 1, 1, 3], [2, 3, 4, 5], [2, 3, 3], [5, 5], [10]], 35 "output_type": [tf.int32, tf.int64], 36 "is_arg_max": [True], 37 "dynamic_range_quantize": [False, True], 38 }] 39 40 def build_graph(parameters): 41 """Build the topk op testing graph.""" 42 input_value = tf.compat.v1.placeholder( 43 dtype=parameters["input_dtype"], 44 name="input", 45 shape=parameters["input_shape"]) 46 axis = random.randint(0, max(len(parameters["input_shape"]) - 1, 0)) 47 if parameters["is_arg_max"]: 48 out = tf.math.argmax( 49 input_value, axis, output_type=parameters["output_type"]) 50 else: 51 out = tf.math.argmin( 52 input_value, axis, output_type=parameters["output_type"]) 53 return [input_value], [out] 54 55 def build_inputs(parameters, sess, inputs, outputs): 56 input_value = create_tensor_data(parameters["input_dtype"], 57 parameters["input_shape"]) 58 return [input_value], sess.run( 59 outputs, feed_dict=dict(zip(inputs, [input_value]))) 60 61 make_zip_of_tests( 62 options, 63 test_parameters, 64 build_graph, 65 build_inputs, 66 expected_tf_failures=8) 67