You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "xiaosibai (via GitHub)" <gi...@apache.org> on 2023/04/10 06:21:51 UTC
[GitHub] [tvm] xiaosibai commented on a diff in pull request #14423: [TFLite] Support for BATCH_MATMUL tflite operator
xiaosibai commented on code in PR #14423:
URL: https://github.com/apache/tvm/pull/14423#discussion_r1161467217
##########
python/tvm/relay/frontend/tflite.py:
##########
@@ -2959,6 +2977,135 @@ def convert_batch_to_space_nd(self, op):
return out
+ def convert_batch_matmul(self, op):
+ """batch_matmul implementation."""
+ try:
+ from tflite.BatchMatMulOptions import BatchMatMulOptions
+ except ImportError:
+ raise ImportError("The tflite package must be installed")
+
+ input_tensors = self.get_input_tensors(op)
+
+ assert len(input_tensors) == 2, "two input tensor arguments expected"
+
+ batch_matmul_options = BatchMatMulOptions()
+ op_options = op.BuiltinOptions()
+ batch_matmul_options.Init(op_options.Bytes, op_options.Pos)
+
+ input_a = self.get_expr(input_tensors[0].tensor_idx)
+ input_b = self.get_expr(input_tensors[1].tensor_idx)
+
+ shape_a = shape_of(input_a)
+ shape_b = shape_of(input_b)
+ rank_a = _infer_shape(shape_a)[0]
+ rank_b = _infer_shape(shape_b)[0]
+
+ if rank_a > 2 or rank_b > 2:
+ # Determine the output batch dimension
+ new_a_shape = shape_a
+ new_b_shape = shape_b
+ if rank_a > rank_b:
+ rank_diff = rank_a - rank_b
+ new_b_shape = _op.concatenate(
+ [
+ _expr.const([1] * rank_diff, dtype=_infer_type(b_shape).checked_type.dtype),
Review Comment:
the "b_shape" here is a mistake? should be "shape_b"?
##########
python/tvm/relay/frontend/tflite.py:
##########
@@ -2959,6 +2977,135 @@ def convert_batch_to_space_nd(self, op):
return out
+ def convert_batch_matmul(self, op):
+ """batch_matmul implementation."""
+ try:
+ from tflite.BatchMatMulOptions import BatchMatMulOptions
+ except ImportError:
+ raise ImportError("The tflite package must be installed")
+
+ input_tensors = self.get_input_tensors(op)
+
+ assert len(input_tensors) == 2, "two input tensor arguments expected"
+
+ batch_matmul_options = BatchMatMulOptions()
+ op_options = op.BuiltinOptions()
+ batch_matmul_options.Init(op_options.Bytes, op_options.Pos)
+
+ input_a = self.get_expr(input_tensors[0].tensor_idx)
+ input_b = self.get_expr(input_tensors[1].tensor_idx)
+
+ shape_a = shape_of(input_a)
+ shape_b = shape_of(input_b)
+ rank_a = _infer_shape(shape_a)[0]
+ rank_b = _infer_shape(shape_b)[0]
+
+ if rank_a > 2 or rank_b > 2:
+ # Determine the output batch dimension
+ new_a_shape = shape_a
+ new_b_shape = shape_b
+ if rank_a > rank_b:
+ rank_diff = rank_a - rank_b
+ new_b_shape = _op.concatenate(
+ [
+ _expr.const([1] * rank_diff, dtype=_infer_type(b_shape).checked_type.dtype),
+ shape_b,
+ ],
+ 0,
+ )
+ elif rank_a < rank_b:
+ rank_diff = rank_b - rank_a
+ new_a_shape = _op.concatenate(
+ [
+ _expr.const([1] * rank_diff, dtype=_infer_type(a_shape).checked_type.dtype),
Review Comment:
the "a_shape" here is a mistake? should be "shape_a"?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org