OmarAzizi commented on code in PR #20188:
URL: https://github.com/apache/tvm/pull/20188#discussion_r3870876674
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -2173,6 +2173,23 @@ def _impl_v13(cls, bb, inputs, attr, params):
axis = tuple([int(x) for x in axis.data.numpy()])
return cls._squeeze(bb, data, axis)
+ @classmethod
+ def _check_squeeze_axes_are_unit_dims(cls, data, axes):
+ """Raise if any axis to be squeezed has a statically known extent
other than 1."""
+ rank = _get_known_tensor_rank(data)
+ if rank is None:
+ return
+ ty = data.ty
+ if not (isinstance(ty, relax.TensorType) and isinstance(ty.shape,
relax.ShapeExpr)):
+ return
+ for axis in _normalize_constant_axes(list(axes), rank, "Squeeze"):
+ extent = ty.shape.values[axis]
+ if isinstance(extent, tirx.IntImm) and int(extent.value) != 1:
Review Comment:
You're right, thank you. I went with the second approach because it's
similar to how other parts in the frontend handle these cases.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
