1#
2# Copyright (C) 2021 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16def test(name, input0, input1, adj0, adj1, output, input0_data, input1_data,
17    output_data):
18  model = Model().Operation("BATCH_MATMUL", input0, input1, adj0, adj1).To(
19      output)
20  quant8_signed = DataTypeConverter().Identify({
21      input0: ("TENSOR_QUANT8_ASYMM_SIGNED", 0.25, 0),
22      input1: ("TENSOR_QUANT8_ASYMM_SIGNED", 0.50, -64),
23      output: ("TENSOR_QUANT8_ASYMM_SIGNED", 1.00, -128),
24  })
25  Example({
26      input0: input0_data,
27      input1: input1_data,
28      output: output_data,
29  }, model=model,
30      name=name).AddVariations("float16", "int32", quant8_signed)
31
32
33test(
34    name="Simple",
35    input0=Input("op1", "TENSOR_FLOAT32", "{1, 2, 3}"),
36    input1=Input("op2", "TENSOR_FLOAT32", "{1, 3, 4}"),
37    adj0=BoolScalar("adj0", False),
38    adj1=BoolScalar("adj1", False),
39    output=Output("op3", "TENSOR_FLOAT32", "{1, 2, 4}"),
40    input0_data=[1, 2, 3, 4, 5, 6],
41    input1_data=[7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18],
42    output_data=[74., 80., 86., 92., 173., 188., 203., 218.],
43)
44
45test(
46    name="RHSAdjoint",
47    input0=Input("op1", "TENSOR_FLOAT32", "{1, 2, 3}"),
48    input1=Input("op2", "TENSOR_FLOAT32", "{1, 4, 3}"),
49    adj0=BoolScalar("adj0", False),
50    adj1=BoolScalar("adj1", True),
51    output=Output("op3", "TENSOR_FLOAT32", "{1, 2, 4}"),
52    input0_data=[1, 2, 3, 4, 5, 6],
53    input1_data=[7, 11, 15, 8, 12, 16, 9, 13, 17, 10, 14, 18],
54    output_data=[74., 80., 86., 92., 173., 188., 203., 218.],
55)
56
57test(
58    name="LHSAdjoint",
59    input0=Input("op1", "TENSOR_FLOAT32", "{1, 3, 2}"),
60    input1=Input("op2", "TENSOR_FLOAT32", "{1, 3, 4}"),
61    adj0=BoolScalar("adj0", True),
62    adj1=BoolScalar("adj1", False),
63    output=Output("op3", "TENSOR_FLOAT32", "{1, 2, 4}"),
64    input0_data=[1, 4, 2, 5, 3, 6],
65    input1_data=[7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18],
66    output_data=[74., 80., 86., 92., 173., 188., 203., 218.],
67)
68
69test(
70    name="TwoBatchSize",
71    input0=Input("op1", "TENSOR_FLOAT32", "{2, 2, 3}"),
72    input1=Input("op2", "TENSOR_FLOAT32", "{2, 3, 4}"),
73    adj0=BoolScalar("adj0", False),
74    adj1=BoolScalar("adj1", False),
75    output=Output("op3", "TENSOR_FLOAT32", "{2, 2, 4}"),
76    input0_data=[1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6],
77    input1_data=[7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
78                 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18],
79    output_data=[74., 80., 86., 92., 173., 188., 203., 218.,
80                 74., 80., 86., 92., 173., 188., 203., 218.],
81)
82