[GitHub] [tvm] mbrookhart commented on a change in pull request #8251: [Frontend, Tensorflow] Support for broadcasting in batch_matmul when shapes differ

[GitBox <gi...@apache.org>]

mbrookhart commented on a change in pull request #8251:
URL: https://github.com/apache/tvm/pull/8251#discussion_r651276848



##########
File path: python/tvm/relay/frontend/tensorflow_ops.py
##########
@@ -1157,11 +1154,18 @@ def _impl(inputs, attr, params, mod):
                 new_shape_y = _op.concatenate(_op.Tuple(new_shape_y), axis=0)
 
             input_x = _op.reshape(input_x, newshape=new_shape_x)
-            input_y = _op.reshape(input_y, newshape=new_shape_y)
+
+            if np.prod(orig_shape_y) < np.prod(new_shape_y):
+                input_y = _op.broadcast_to(input_y, new_shape_y)

Review comment:
       Broadcasting here could cause memory sizes to explode, would it be better to use implicit broadcasting in the batch_matmul kernel if possible? We do that in the onnx importer, but converting back to ND can be a pain.




-- 
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org