tlopex commented on code in PR #20188:
URL: https://github.com/apache/tvm/pull/20188#discussion_r3874846173
##########
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:
Thanks for handling symbolic dimensions. Runtime axes still bypass this
check though, so invalid non-unit axes remain accepted. Please add a runtime
assertion or reject non-constant axes during import, with a regression test.
--
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]
