wehub-resource-sync
90 lines
3.8 KiB
Python
90 lines
3.8 KiB
Python
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Test configs for einsum."""
|
|
import tensorflow as tf
|
|
from tensorflow.lite.testing.zip_test_utils import create_tensor_data
|
|
from tensorflow.lite.testing.zip_test_utils import make_zip_of_tests
|
|
from tensorflow.lite.testing.zip_test_utils import register_make_test_function
|
|
from tensorflow.python.framework import test_util
|
|
|
|
|
|
@register_make_test_function("make_einsum_tests")
|
|
@test_util.enable_control_flow_v2
|
|
def make_einsum_tests(options):
|
|
"""Make a set of tests to do basic einsum ops."""
|
|
|
|
test_parameters = [
|
|
{
|
|
"dtype": [tf.float32],
|
|
"shapes": [
|
|
((None, None, 8, 64), (4, None, 8, 64), "BQNH,BTNH->BQNT"),
|
|
((1, None, 8, None), (1, None, 8, 64), "BQNT,BTNH->BQNH"),
|
|
((None, None, 8, 64), (8, 8, 64), "ABNH,NDH->ABD"),
|
|
((None, None, 128), (128, 8, 64), "ABD,DNH->ABNH"),
|
|
((3, 4, 5), (3, 5, 6), "ijk,ikm->ijm"),
|
|
((3, 4, 5), (5, 6), "ijk,km->ijm"),
|
|
((2, 5, 7), (5, 2), "LBH,BL->BH"),
|
|
((2, 5, 7), (5, 3, 2), "LBH,BKL->BKH"),
|
|
((2, 5, 7, 3), (2, 4, 7, 3), "BFNH,BTNH->BNFT"),
|
|
((2, 5, 7, 3), (7, 3, 4), "BFND,NDH->BFH"),
|
|
((3, 4, 5), (5, 6, 2), "BFD,DNH->BFNH"),
|
|
((7, 11, 13), (7, 11, 13, 5), "BIN,BINJ->BIJ"),
|
|
((7, 11, 19), (7, 11, 13, 19), "BIJ,BINJ->BIN"),
|
|
((5, 13, 3, 11), (5, 11, 13, 8), "ACBE,AECD->ABCD"),
|
|
((5, 11, 7, 3), (5, 8, 7, 3), "AECD,ABCD->ACBE"),
|
|
((5, 4, 3), (3, 2, 1), "...ij,j...->i..."),
|
|
((5, 4, 3), (3, 2, 1), "...ij,j...->...i"),
|
|
((1, 11, 19), (7, 11, 13, 19), "...IJ,...INJ->...IN"),
|
|
((1, 11, 19), (7, 11, 13, 19), "...IJ,...INJ->IN..."),
|
|
((4, 3, 2, 5), (3, 6, 1), "ij...,jk...->ik..."),
|
|
((4, 3, 2, 5), (3, 6, 1), "ij...,jk...->...ik"),
|
|
],
|
|
},
|
|
]
|
|
|
|
def set_dynamic_shape(shape):
|
|
"""Convert dynamic shapes to static shapes."""
|
|
return [4 if x is None else x for x in shape]
|
|
|
|
def build_graph(parameters):
|
|
"""Build a simple graph with einsum Op."""
|
|
input0_shape = parameters["shapes"][0]
|
|
input1_shape = parameters["shapes"][1]
|
|
equation = parameters["shapes"][2]
|
|
|
|
input0_tensor = tf.compat.v1.placeholder(
|
|
dtype=parameters["dtype"], shape=input0_shape)
|
|
input1_tensor = tf.compat.v1.placeholder(
|
|
dtype=parameters["dtype"], shape=input1_shape)
|
|
out = tf.einsum(equation, input0_tensor, input1_tensor)
|
|
return [input0_tensor, input1_tensor], [out]
|
|
|
|
def build_inputs(parameters, sess, inputs, outputs):
|
|
"""Feed inputs, assign variables, and freeze graph."""
|
|
input0_shape = set_dynamic_shape(parameters["shapes"][0])
|
|
input1_shape = set_dynamic_shape(parameters["shapes"][1])
|
|
input0_value = create_tensor_data(parameters["dtype"], input0_shape)
|
|
input1_value = create_tensor_data(parameters["dtype"], input1_shape)
|
|
output_values = sess.run(
|
|
outputs, feed_dict=dict(zip(inputs, [input0_value, input1_value])))
|
|
return [input0_value, input1_value], output_values
|
|
|
|
make_zip_of_tests(
|
|
options,
|
|
test_parameters,
|
|
build_graph,
|
|
build_inputs,
|
|
use_frozen_graph=True)
|