Add Support for MatMulInteger (#2072)

* Add Support for MatMulInteger MatMulInteger was supported in ONNX opset v10 (not checked in proposed change, the error can be addressed on save), this specific type combination is support in TensorFlow, but the node type not identified and handled properly here. Handles #2071 Signed-off-by: Gregory Morse <gregory.morse@live.com> * Update math.py Signed-off-by: Gregory Morse <gregory.morse@live.com> * Update support_status.md Signed-off-by: Gregory Morse <gregory.morse@live.com> * Update test_backend.py Signed-off-by: Gregory Morse <gregory.morse@live.com> Signed-off-by: Gregory Morse <gregory.morse@live.com> Co-authored-by: Jay Zhang <36183870+fatcat-z@users.noreply.github.com>

Author: GregoryMorse

support_status.md

@@ -27,6 +27,7 @@
 | AvgPool3D | 1 ~ 17 |
 | BatchMatMul | 1 ~ 17 |
 | BatchMatMulV2 | 1 ~ 17 |
+| BatchMatMulV3 | 1 ~ 17 |
 | BatchToSpaceND | 1 ~ 17 |
 | BiasAdd | 1 ~ 17 |
 | BiasAddV1 | 1 ~ 17 |

tests/test_backend.py

@@ -1075,6 +1075,15 @@ def func(x, y):
  return tf.identity(x_, name=_TFOUTPUT)
  self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: x_val}, rtol=1e-5)
 
+ @check_tf_min_version("2.6")
+ def test_matmulinteger(self):
+ x_val = np.array([1, 2, -3, -4], dtype=np.int8).reshape((2, 2))
+ y_val = np.array([1, 2, -3, -4], dtype=np.int8).reshape((2, 2))
+ def func(x, y):
+ x_ = tf.matmul(x, y, output_type=tf.int32)
+ return tf.identity(x_, name=_TFOUTPUT)
+ self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
  @check_onnxruntime_incompatibility("Sub")
  def test_sub(self):
  x_val = np.array([1.0, 2.0, -3.0, -4.0], dtype=np.float32).reshape((2, 2))

tf2onnx/onnx_opset/math.py

@@ -363,13 +363,13 @@ def version_1(cls, ctx, node, **kwargs):
  name=op_name, shapes=shapes, dtypes=dtypes)
 
 
-@tf_op(["MatMul", "BatchMatMul", "BatchMatMulV2"])
+@tf_op(["MatMul", "BatchMatMul", "BatchMatMulV2", "BatchMatMulV3"])
 class MatMul:
  @classmethod
  def version_1(cls, ctx, node, **kwargs):
  # tensorflow allows transpose and conjugated. If found, insert the required transpose.
  # We could use Gemm as well but tensorflow does not pass bias in matmul.
- node.type = "MatMul"
+ if node.type != "MatMulInteger": node.type = "MatMul"
 
  attrs = ["transpose_a", "transpose_b", "adjoint_a", "adjoint_b", "adj_x", "adj_y"]
  attrs_val = [node.get_attr(attr) for attr in attrs]
@@ -408,7 +408,19 @@ def version_1(cls, ctx, node, **kwargs):
  val = node.get_attr(i)
  if val is not None and val.i != 0:
  raise ValueError(node.type + " attribute " + i + " is not supported")
-
+ @classmethod
+ def version_10(cls, ctx, node, **kwargs):
+ if (ctx.get_dtype(node.input[0]) in [onnx_pb.TensorProto.INT8, onnx_pb.TensorProto.UINT8] and
+ ctx.get_dtype(node.input[1]) in [onnx_pb.TensorProto.INT8, onnx_pb.TensorProto.UINT8] and
+ ctx.get_dtype(node.output[0]) == onnx_pb.TensorProto.INT32):
+ node.type = "MatMulInteger"
+ zpdata_a = np.zeros(1, dtype=utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[0])))
+ zero_point_node_a = ctx.make_const(utils.make_name("zero_point_a"), zpdata_a)
+ zpdata_b = np.zeros(1, dtype=utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[1])))
+ zero_point_node_b = ctx.make_const(utils.make_name("zero_point_b"), zpdata_b)
+ ctx.replace_inputs(node, [node.input[0], node.input[1],
+ zero_point_node_a.output[0], zero_point_node_b.output[0]])
+ cls.version_1(ctx, node, **kwargs)
 
 @tf_op("Erf")
 class Erf: