In squeeze in python/tvm/relax/frontend/nn/op.py:
(1) The axis argument accepts int or list of int.
(2) The nominal default value of axis is -1 (this is just nominal since TVMScript is not Python code, so the default value is not really stored in module), which means removing all axis of dimensions 1, instead of removing the last dim.
In squeeze in python/tvm/relax/op/manipulate.py:
(1) If the axis is int, then it is converted to a list of one element.
In InferLayoutSqueeze in src/relax/op/tensor/manipulate.cc:
(1) The axis attribute is read as-is, and for each value of axis the corresponding element axis_str is set. This means all elements of axis can not be negative, otherwise this is an out-of-bound write.
In Squeeze in python/tvm/relax/frontend/onnx/onnx_frontend.py:
(1) The axis are stored as-is from onnx (see https://github.com/apache/tvm/blob/5398c27e76c7c858d7fae655496e68b297fc99fc/python/tvm/relax/frontend/onnx/onnx_frontend.py#L2424C9-L2424C36). This means if any value of axis is negative, then the result Relax op is invalid.
What is expected: axis must be normalized before op is created.
A simple reproduce:
import tvm, onnx
from tvm.relax.frontend.onnx import from_onnx
model = onnx.load_model_from_string(b'\x08\n\x12\x18onnx.utils.extract_model:\x8c\x02\n5\n\x0btranspose_1\n\x06val_23\x12\x07squeeze\x1a\x0cnode_squeeze"\x07Squeeze\x12\x1bExtracted from {main_graph}*\x16\x08\x01\x10\x07B\x06val_23J\x08\xfe\xff\xff\xff\xff\xff\xff\xffZ*\n\x0btranspose_1\x12\x1b\n\x19\x08\x01\x12\x15\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x02\x08\x01\n\x03\x08\x80\x08b"\n\x07squeeze\x12\x17\n\x15\x08\x01\x12\x11\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x03\x08\x80\x08j*\n\x0btranspose_1\x12\x1b\n\x19\x08\x01\x12\x15\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x02\x08\x01\n\x03\x08\x80\x08j"\n\x07squeeze\x12\x17\n\x15\x08\x01\x12\x11\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x03\x08\x80\x08B\x04\n\x00\x10\x14')
tvm.relax.transform.ConvertLayout({})(from_onnx(model))
In tvm 0.26, it will throw std::bad_alloc.
In
squeezein python/tvm/relax/frontend/nn/op.py:(1) The
axisargument accepts int or list of int.(2) The nominal default value of
axisis -1 (this is just nominal since TVMScript is not Python code, so the default value is not really stored in module), which means removing all axis of dimensions 1, instead of removing the last dim.In
squeezein python/tvm/relax/op/manipulate.py:(1) If the
axisis int, then it is converted to a list of one element.In
InferLayoutSqueezein src/relax/op/tensor/manipulate.cc:(1) The
axisattribute is read as-is, and for each value ofaxisthe corresponding elementaxis_stris set. This means all elements ofaxiscan not be negative, otherwise this is an out-of-bound write.In
Squeezein python/tvm/relax/frontend/onnx/onnx_frontend.py:(1) The
axisare stored as-is from onnx (see https://github.com/apache/tvm/blob/5398c27e76c7c858d7fae655496e68b297fc99fc/python/tvm/relax/frontend/onnx/onnx_frontend.py#L2424C9-L2424C36). This means if any value ofaxisis negative, then the result Relax op is invalid.What is expected:
axismust be normalized before op is created.A simple reproduce:
In tvm 0.26, it will throw
std::bad_alloc.